From 4e00919e5ce09add1a2343b92e7b18aaa52c9d93 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Thu, 26 Feb 2026 09:26:20 +0100 Subject: [PATCH] [CUDA][NCCL] group split (#3172) --- mlx/distributed/nccl/nccl.cpp | 64 ++++++++++++++++++------- python/tests/nccl_test_distributed.py | 67 +++++++++++++++++++++++++++ 2 files changed, 114 insertions(+), 17 deletions(-) diff --git a/mlx/distributed/nccl/nccl.cpp b/mlx/distributed/nccl/nccl.cpp index d8244bf9..c7bce4e0 100644 --- a/mlx/distributed/nccl/nccl.cpp +++ b/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 + create(int numRanks, int rank, ncclUniqueId commId) { + ncclComm_t raw; + CHECK_NCCL(ncclCommInitRank(&raw, numRanks, commId, rank)); + return std::make_shared(raw, rank, numRanks); + } + + static std::shared_ptr 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(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 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 split(int color, int key = -1) override { - throw std::runtime_error("[nccl] Group split not supported."); + std::shared_ptr 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( + 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(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 comm_; bool initialized_ = false; }; diff --git a/python/tests/nccl_test_distributed.py b/python/tests/nccl_test_distributed.py index e6eed8a9..637d5832 100644 --- a/python/tests/nccl_test_distributed.py +++ b/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()