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 4 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. |
|
part 2 / also noticed the low priority label. would it be worth reconsidering that given this can return incorrect data rather than just fail with an error? getting this verified would help people use the ring w/o custom workarounds and potentially avoid incorrect results going unnoticed. thanks. |
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.