diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 15b81a2177..d6305d9d4b 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -198,7 +198,11 @@ class RingImpl { size_t n_bytes_per_wire = (n_bytes + (2 * n_wires) - 1) / (2 * n_wires); dispatch_wires(n_wires, [&](int lw) { all_gather_wire( - out_ptr, n_bytes, static_cast(n_bytes_per_wire), lw); + out_ptr, + n_bytes, + static_cast(n_bytes_per_wire), + n_wires, + lw); }); } @@ -215,14 +219,21 @@ class RingImpl { char* out_ptr, int64_t n_bytes, int64_t n_bytes_per_wire, + int n_wires, int lw) { - // Both directions send the same contiguous slice of each rank's region. - int64_t slice = static_cast(lw) * n_bytes_per_wire; - int64_t wire_offset[2] = {slice, slice}; - // Clamp to this wire's slice end so the last frame can't spill into the - // next wire's slice. - int64_t wire_end_bytes = std::min(n_bytes, slice + n_bytes_per_wire); - int64_t wire_end[2] = {wire_end_bytes, wire_end_bytes}; + // Each rank's region is divided into 2 directional halves, each split + // into n_wires contiguous slices - mirroring all_reduce_wire. Without + // the lr term both directions cover the same lower-half slice and the + // upper half of every peer's region is never transferred. + int64_t wire_offset[2]; + int64_t wire_end[2]; + for (int lr = 0; lr < 2; lr++) { + wire_offset[lr] = lr * n_wires * n_bytes_per_wire + + static_cast(lw) * n_bytes_per_wire; + int64_t region_end = + std::min(n_bytes, (lr + 1) * n_wires * n_bytes_per_wire); + wire_end[lr] = std::min(region_end, wire_offset[lr] + n_bytes_per_wire); + } int64_t send_offset[2] = {rank_ * n_bytes, rank_ * n_bytes}; int64_t recv_offset[2] = { ((rank_ + size_ - 1) % size_) * n_bytes, @@ -230,7 +241,7 @@ class RingImpl { ring_pass<2, char>( lw, - n_bytes, + n_bytes * size_, n_bytes, n_bytes * size_, n_bytes_per_wire, diff --git a/python/tests/mlx_distributed_tests.py b/python/tests/mlx_distributed_tests.py index cbb9663046..7d6e56aaac 100644 --- a/python/tests/mlx_distributed_tests.py +++ b/python/tests/mlx_distributed_tests.py @@ -325,6 +325,22 @@ def test_all_gather(self): self.assertEqual(y.shape, (world.size() * 2, 2, 4)) self.assertTrue(mx.all(y == 1)) + def test_all_gather_rank_data(self): + # Every rank contributes distinct data so that a missing, stale, or + # misplaced peer region is detected (an all-ones gather cannot tell). + # Sizes include one large enough that backends which slice transfers + # across directions/wires (e.g. the jaccl ring) exercise every slice, + # and odd sizes that exercise the tail clamping. + world = mx.distributed.init() + for size in [7, 1024, 1000003, 4 * 1024 * 1024]: + base = mx.arange(size, dtype=mx.int32) % 100003 + x = base + world.rank() + y = mx.distributed.all_gather(x) + self.assertEqual(y.shape, (world.size() * size,)) + regions = y.reshape(world.size(), size) + for r in range(world.size()): + self.assertTrue(mx.all(regions[r] == base + r).item()) + def test_clip_grad_norm_sharded(self): world = mx.distributed.init() N = world.size()