Fix jaccl ring all_gather: direction 1 slice was not mirrored - #4443
Fix jaccl ring all_gather: direction 1 slice was not mirrored#4443Drifter4242 wants to merge 2 commits into
Conversation
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.
|
Thank you, will look into this as soon as possible. |
|
Do you mind covering it with tests for all_gather? |
nastya236
added
the
await response
…cally 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.
|
I went ahead and wrote a test for it (actually I got Fable to write the test and run it, but I did review the test and results).
I have 2 M3 Ultras with thunderbolt that I can do the test on. I think the standard mlx CI doesn't do the test on that hardware which is why it got missed. I hope that's enough. I admit that this isn't my area of expertise. |
zcbenz
added
await verification
|
i'm seeing something very similar on my setup: 4 Mac Studio, running mlx 0.32.1 with jaccl-ring over thunderbolt for distributed pretraining. i reproduced incomplete/incorrect results from native all_gather with both int32 and float32 arrays, including explicitly evaluated inputs. meanwhile, all_sum passed checks for now, my workaround is just for small control messages: each rank puts its message into its own section of a zero-filled array, then we use all_sum so every rank gets all four messages. keeping checksums and sequence checks in place. gradient reductions still use native all_sum. haven't tested this pr's patch yet, so i can't confirm it fixes my 4-Mac studio setup, but the symptoms and the two issues described here line up closely with what I see. |
5 participants
I confess that Fable found and fixed this bug, and I struggle to understand exactly what the problem is. But since the change is fairly small, I thought it reasonable to create a PR for review.
My understanding is that the all_gather was refactored recently and has a bug.
Both sides were sending the lower half of the slice. Instead it should use the same code as the all_reduce_wire and send upper half in the second direction.
There is also a second bug fix in this PR, where there is a one-line change to a size parameter for ring_pass.