Skip to content

Add sum_scatter to the ring backend - #4236

Open
erwinzhang7 wants to merge 1 commit into
ml-explore:mainfrom
erwinzhang7:ring-sum-scatter
Open

Add sum_scatter to the ring backend#4236
erwinzhang7 wants to merge 1 commit into
ml-explore:mainfrom
erwinzhang7:ring-sum-scatter

Conversation

@erwinzhang7

Copy link
Copy Markdown
Contributor

RingGroup::sum_scatter threw, so mx.distributed.sum_scatter was unavailable on ring and
nn.fully_shard could not run there at all: it reduce-scatters gradients in the backward pass.
Since mlx.launch uses ring locally, FSDP was in practice CUDA only. NCCL and JACCL already
implement the op, so this makes ring the third of the four backends.

No new communication machinery needed. A ring all reduce is already a scatter reduce
followed by a gather, so a reduce scatter is the first of those two passes. all_reduce_impl
takes a phases argument to stop after the first, and a first_segment to start a segment
ahead, which is what leaves each rank holding its own segment rather than one dictated by the
direction of travel.

The pipeline keeps a send and a recv in flight, so the loop was consuming one less entry than
it issued. With two phases the trailing entry belongs to the gather, which writes into the data
directly and needs no reduction, so nothing was wrong. With one phase every entry has to be
reduced and the last one was not. The loop now runs the extra step to drain it, which also
removes the valid() checks that were there to avoid consuming a future twice.

Splitting the work over the sockets had to change direction. The all reduce gives each wire a
chunk of the array, which for a reduce scatter would leave every rank owning a piece of each
chunk instead of one segment of the whole. So each wire takes a slice of every segment
instead, and the segment a rank returns is complete once the wires finish. That is the
part_offset and part_size pair; the all reduce passes zeros and behaves exactly as before.

The one cost is a scratch buffer the size of the input. The reduction accumulates in place and
the input is const, and unlike the all reduce there is no output big enough to borrow, since it
is 1 / size of the input.

It also needs a one line fix in mlx/backend/cpu/distributed.cpp. ReduceScatter::eval_cpu
asserted inputs.size() == 0 and then read inputs[0], copied from Recv::eval_cpu above it
where no inputs is right. Nothing reached it before: ring threw, mpi throws, nccl evaluates on
the gpu, and jaccl needs hardware the runners do not have. It aborts a debug build as soon as
ring can serve the op.

Verified

macOS 26.6, M5 Max. One address per rank already gives a socket in each direction, so the split
runs there too; more addresses take it further. Measured on a 1M element float32, where each
wire ends up with two packets:

addresses per rank sockets scatters
1 1 left, 1 right 2
2 2 left, 2 right 4

Small inputs stay on a single scatter.

4 ranks, 2 addresses each 14/14 every rank
2 ranks, 4 addresses each 14/14 every rank
1 address, 2 / 4 / 8 ranks 14/14 every rank

Run again against a debug build, which is what turned up the assertion and what the linux legs
of CI use.

test_sum_scatter in ring_test_distributed.py follows the nccl one and checks the result
against the matching slice of all_sum, over float32, float16 and bfloat16 and four shapes. The
all reduce tests in the same suite cover the code it shares.

Separately, nn.fully_shard now trains on ring: forward, backward, optimizer step and a second
forward with the loss moving.

@zcbenz zcbenz added the await verification This pull request is non-trivial and requires a human expert to verify its correctness. label Aug 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await verification This pull request is non-trivial and requires a human expert to verify its correctness.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants