Add sum_scatter to the ring backend - #4236
Open
erwinzhang7 wants to merge 1 commit into
Open
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
RingGroup::sum_scatterthrew, somx.distributed.sum_scatterwas unavailable on ring andnn.fully_shardcould not run there at all: it reduce-scatters gradients in the backward pass.Since
mlx.launchuses ring locally, FSDP was in practice CUDA only. NCCL and JACCL alreadyimplement 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_impltakes a
phasesargument to stop after the first, and afirst_segmentto start a segmentahead, 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_offsetandpart_sizepair; 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 / sizeof the input.It also needs a one line fix in
mlx/backend/cpu/distributed.cpp.ReduceScatter::eval_cpuasserted
inputs.size() == 0and then readinputs[0], copied fromRecv::eval_cpuabove itwhere 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:
Small inputs stay on a single scatter.
Run again against a debug build, which is what turned up the assertion and what the linux legs
of CI use.
test_sum_scatterinring_test_distributed.pyfollows the nccl one and checks the resultagainst the matching slice of
all_sum, over float32, float16 and bfloat16 and four shapes. Theall reduce tests in the same suite cover the code it shares.
Separately,
nn.fully_shardnow trains on ring: forward, backward, optimizer step and a secondforward with the loss moving.