Distributed layers

mlx/distributed/distributed.cpp

@@ -60,6 +60,10 @@ class EmptyGroup : public GroupImpl {
     throw std::runtime_error(
         "Communication not implemented in an empty distributed group.");
   }
+  void barrier() override {
+    throw std::runtime_error(
+        "Barrier not implemented in an empty distributed group.");
+  }
 };
 
 } // namespace detail
@@ -80,6 +84,10 @@ Group Group::split(int color, int key /* = -1 */) const {
   return Group(group_->split(color, key));
 }
 
+void Group::barrier() {
+  return group_->barrier();
+}
+
 Group init(bool strict /* = false */) {
   auto init_group = [strict]() {
     auto default_group = mpi::init(strict);

mlx/distributed/distributed.h

@@ -41,6 +41,8 @@ struct Group {
     return group_;
   }
 
+  void barrier();
+
  private:
   std::shared_ptr<detail::GroupImpl> group_{nullptr};
 };

mlx/distributed/distributed_impl.h

@@ -19,6 +19,7 @@ class GroupImpl {
   virtual void all_gather(const array& input, array& output) = 0;
   virtual void send(const array& input, int dst) = 0;
   virtual void recv(array& out, int src) = 0;
+  virtual void barrier() = 0;
 };
 
 /* Return the communication stream. */

mlx/distributed/mpi/mpi.cpp

@@ -74,6 +74,7 @@ struct MPIWrapper {
     LOAD_SYMBOL(MPI_Allgather, all_gather);
     LOAD_SYMBOL(MPI_Send, send);
     LOAD_SYMBOL(MPI_Recv, recv);
+    LOAD_SYMBOL(MPI_Barrier, barrier);
     LOAD_SYMBOL(MPI_Type_contiguous, mpi_type_contiguous);
     LOAD_SYMBOL(MPI_Type_commit, mpi_type_commit);
     LOAD_SYMBOL(MPI_Op_create, mpi_op_create);
@@ -198,6 +199,7 @@ struct MPIWrapper {
   int (*comm_free)(MPI_Comm*);
   int (*send)(const void*, int, MPI_Datatype, int, int, MPI_Comm);
   int (*recv)(void*, int, MPI_Datatype, int, int, MPI_Comm, MPI_Status*);
+  int (*barrier)(MPI_Comm);
 
   // Objects
   MPI_Comm comm_world_;
@@ -319,6 +321,10 @@ class MPIGroup : public GroupImpl {
         &status);
   }
 
+  void barrier() override {
+    mpi().barrier(comm_);
+  }
+
  private:
   MPI_Comm comm_;
   bool global_;

python/mlx/nn/layers/__init__.py

@@ -60,6 +60,12 @@
     ConvTranspose2d,
     ConvTranspose3d,
 )
+from mlx.nn.layers.distributed import (
+    AllToShardedLinear,
+    QuantizedAllToShardedLinear,
+    QuantizedShardedToAllLinear,
+    ShardedToAllLinear,
+)
 from mlx.nn.layers.dropout import Dropout, Dropout2d, Dropout3d
 from mlx.nn.layers.embedding import Embedding
 from mlx.nn.layers.linear import Bilinear, Identity, Linear