[CUDA][NCCL] group split (#3172)

mlx/distributed/nccl/nccl.cpp

@@ -267,28 +267,54 @@ inline void bootstrap_unique_id(
 
 } // namespace detail
 
+// helper struct to manage communicator
+struct NCCLComm {
+  ncclComm_t comm;
+  int rank_;
+  int size_;
+
+  NCCLComm(ncclComm_t c, int rank, int size)
+      : comm(c), rank_(rank), size_(size) {}
+
+  static std::shared_ptr<NCCLComm>
+  create(int numRanks, int rank, ncclUniqueId commId) {
+    ncclComm_t raw;
+    CHECK_NCCL(ncclCommInitRank(&raw, numRanks, commId, rank));
+    return std::make_shared<NCCLComm>(raw, rank, numRanks);
+  }
+
+  static std::shared_ptr<NCCLComm> split(NCCLComm* source, int color, int key) {
+    ncclComm_t raw;
+    // default config, blocking comm creation
+    ncclConfig_t config = NCCL_CONFIG_INITIALIZER;
+    CHECK_NCCL(ncclCommSplit(source->comm, color, key, &raw, &config));
+    int new_rank, new_size;
+    CHECK_NCCL(ncclCommUserRank(raw, &new_rank));
+    CHECK_NCCL(ncclCommCount(raw, &new_size));
+    return std::make_shared<NCCLComm>(raw, new_rank, new_size);
+  }
+
+  NCCLComm(const NCCLComm&) = delete;
+  NCCLComm& operator=(const NCCLComm&) = delete;
+};
+
 using GroupImpl = mlx::core::distributed::detail::GroupImpl;
 class NCCLGroup : public GroupImpl {
  public:
   NCCLGroup(int worldRank, int worldSize, const std::string initMethod)
-      : rank_(worldRank),
-        size_(worldSize),
-        comm_(nullptr),
-        initMethod_(initMethod) {
+      : rank_(worldRank), size_(worldSize), initMethod_(initMethod) {
     if (initialized_)
       return;
     int ndev;
     CHECK_CUDA(cudaGetDeviceCount(&ndev));
     CHECK_CUDA(cudaSetDevice(rank_ % ndev));
     detail::bootstrap_unique_id(uniqueId_, rank_, size_, initMethod_);
-    CHECK_NCCL(ncclCommInitRank(&comm_, size_, uniqueId_, rank_));
+    comm_ = NCCLComm::create(size_, rank_, uniqueId_);
     initialized_ = true;
   }
-
-  ~NCCLGroup() {
-    ncclCommDestroy(comm_);
-    initialized_ = false;
-  }
+  // Used by split() to wrap an already-created communicator
+  NCCLGroup(std::shared_ptr<NCCLComm> comm, int rank, int size)
+      : rank_(rank), size_(size), comm_(std::move(comm)) {}
 
   Stream communication_stream(StreamOrDevice s) override {
     return to_stream(s, Device::gpu);
@@ -309,8 +335,11 @@ class NCCLGroup : public GroupImpl {
     });
   }
 
-  virtual std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
-    throw std::runtime_error("[nccl] Group split not supported.");
+  std::shared_ptr<GroupImpl> split(int color, int key = -1) override {
+    key = (key < 0) ? rank() : key;
+    auto new_comm = NCCLComm::split(comm_.get(), color, key);
+    return std::make_shared<NCCLGroup>(
+        new_comm, new_comm->rank_, new_comm->size_);
   }
 
   void all_gather(const array& input, array& output, Stream stream) override {
@@ -322,7 +351,7 @@ class NCCLGroup : public GroupImpl {
           gpu_ptr<T>(output),
           input.size(),
           dt,
-          comm_,
+          comm_->comm,
           encoder.stream()));
     });
   }
@@ -371,7 +400,7 @@ class NCCLGroup : public GroupImpl {
         input.size(),
         dt,
         op,
-        comm_,
+        comm_->comm,
         encoder.stream()));
   }
 
@@ -390,14 +419,15 @@ class NCCLGroup : public GroupImpl {
         output.size(),
         dt,
         op,
-        comm_,
+        comm_->comm,
         encoder.stream()));
   }
 
-  int rank_, size_;
+  int rank_;
+  int size_;
   std::string initMethod_;
   ncclUniqueId uniqueId_;
-  ncclComm_t comm_;
+  std::shared_ptr<NCCLComm> comm_;
   bool initialized_ = false;
 };
 

python/tests/nccl_test_distributed.py

@@ -47,6 +47,73 @@ class TestNCCLDistributed(mlx_distributed_tests.MLXDistributedCommonTestCase):
                 maxrelerror = maxrelerror.max()
                 self.assertLessEqual(maxrelerror, rtol)
 
+    def test_groups(self):
+        world = mx.distributed.init()
+        self.assertEqual(world.size(), 8)
+        self.assertTrue(0 <= world.rank() < 8)
+
+        world2 = mx.distributed.init()
+        self.assertEqual(world.size(), world2.size())
+        self.assertEqual(world.rank(), world2.rank())
+
+        sub = world.split(world.rank() % 2)
+        self.assertEqual(sub.size(), 4)
+        self.assertEqual(sub.rank(), world.rank() // 2)
+
+        sub = world.split(world.rank() // 2)
+        self.assertEqual(sub.size(), 2)
+
+    def test_all_reduce_split(self):
+        world = mx.distributed.init()
+        dtypes = [
+            (mx.float32, 1e-6),
+            (mx.float16, 5e-3),
+            (mx.bfloat16, 1e-1),
+        ]
+        sizes = [
+            (7,),
+            (10,),
+            (1024,),
+            (1024, 1024),
+        ]
+        key = mx.random.key(0)
+        group = world.split(world.rank() % 2)
+
+        for dt, rtol in dtypes:
+            for sh in sizes:
+                x = (
+                    mx.random.uniform(shape=(group.size(),) + sh, key=key) * 10
+                ).astype(dt)
+
+                # All sum
+                y = mx.distributed.all_sum(x[group.rank()], group=group)
+                z = x.sum(0)
+                maxrelerror = (y - z).abs()
+                if rtol > 0:
+                    maxrelerror /= z.abs()
+                maxrelerror = maxrelerror.max()
+                self.assertLessEqual(maxrelerror, rtol)
+
+                # All max
+                y = mx.distributed.all_max(x[group.rank()], group=group)
+                z = x.max(0)
+                self.assertTrue(mx.all(y == z))
+
+                # All min
+                y = mx.distributed.all_min(x[group.rank()], group=group)
+                z = x.min(0)
+                self.assertTrue(mx.all(y == z))
+
+    def test_all_gather_split(self):
+        world = mx.distributed.init()
+        dtypes = [mx.float32, mx.float16, mx.bfloat16]
+        sub = world.split(world.rank() % 2)
+        for dt in dtypes:
+            x = mx.ones((2, 2, 4), dtype=dt)
+            y = mx.distributed.all_gather(x, group=sub)
+            self.assertEqual(y.shape, (sub.size() * 2, 2, 4))
+            self.assertTrue(mx.all(y == 1))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()