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

2026-02-26 09:26:20 +01:00
parent 6ec0192270
commit 4e00919e5c
2 changed files with 114 additions and 17 deletions
@@ -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()