From 9b6b10c766d776174d43da2cf956e4ac83a7f5a1 Mon Sep 17 00:00:00 2001 From: Dave Hind Date: Sun, 30 Aug 2026 10:06:44 -0500 Subject: [PATCH 1/4] Fix jaccl ring all_gather: direction 1 slice was not mirrored The 3900 refactor gives all_gather_wire identical wire_offset/wire_end for both ring directions, so both cover the same lower-half slice of each peer region and the upper half is never transferred (stale/zero bytes in the output). all_reduce_wire computes the directional term lr * n_wires * size_per_wire correctly, which is why all_sum works; mirror that logic here. Repro: MLX_JACCL_RING=1, 2 ranks: all_gather corrupts at every size while all_sum passes. --- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 32 +++++++++++++++------ 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index 15b81a2177..dab0dab318 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -198,7 +198,7 @@ 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,22 +215,36 @@ 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, ((rank_ + 1) % size_) * n_bytes}; + // The size argument is the valid extent of the offset coordinate space, + // which ring_pass clamps copies against (size - offset). For all_gather + // the send/recv offsets are absolute positions in the full output buffer + // and every rank's region is complete, so the extent is the whole output + // (n_bytes * size_) - passing a single region's n_bytes made the clamp + // underflow to 0 for any rank whose peer region starts above offset 0 + // (every rank but the one owning region 0 received only zeros). ring_pass<2, char>( lw, - n_bytes, + n_bytes * size_, n_bytes, n_bytes * size_, n_bytes_per_wire, From e3b493304c252a945549a7454b936a8ff24224ab Mon Sep 17 00:00:00 2001 From: Dave Hind Date: Tue, 1 Sep 2026 20:03:15 -0500 Subject: [PATCH 2/4] Added all_gather test to show the bug in the all_gather code. Specifically this test uses distinct data at larger sizes. Tested on 2 M3 Ultras with thunderbolt. Fable wrote the code and performed the tests. I reviewed the results and wrote the commit message. --- python/tests/mlx_distributed_tests.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) 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() From ddc69d5704795b5010cbe2d6a1729973cebefcc3 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Fri, 11 Sep 2026 11:50:22 +0200 Subject: [PATCH 3/4] drop claude comment --- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 7 ------- 1 file changed, 7 deletions(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index dab0dab318..c74a54da54 100644 --- a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h +++ b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h @@ -235,13 +235,6 @@ class RingImpl { ((rank_ + size_ - 1) % size_) * n_bytes, ((rank_ + 1) % size_) * n_bytes}; - // The size argument is the valid extent of the offset coordinate space, - // which ring_pass clamps copies against (size - offset). For all_gather - // the send/recv offsets are absolute positions in the full output buffer - // and every rank's region is complete, so the extent is the whole output - // (n_bytes * size_) - passing a single region's n_bytes made the clamp - // underflow to 0 for any rank whose peer region starts above offset 0 - // (every rank but the one owning region 0 received only zeros). ring_pass<2, char>( lw, n_bytes * size_, From ece789538dbc6b1b2df7e19aa92ed75b5bc066d3 Mon Sep 17 00:00:00 2001 From: Anastasiia Filippova Date: Wed, 16 Sep 2026 01:12:10 +0200 Subject: [PATCH 4/4] pre-commit --- mlx/distributed/jaccl/lib/jaccl/ring_impl.h | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/mlx/distributed/jaccl/lib/jaccl/ring_impl.h b/mlx/distributed/jaccl/lib/jaccl/ring_impl.h index c74a54da54..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), n_wires, lw); + out_ptr, + n_bytes, + static_cast(n_bytes_per_wire), + n_wires, + lw); }); }