Repository navigation
Fix all_gather benchmark collapsing its input to a scalar - #4017
Conversation
| def fn(x): | ||
| for _ in range(its_per_eval): | ||
| x = mx.distributed.all_gather(x)[0] | ||
| x = mx.distributed.all_gather(x)[: a.shape[0]] |
There was a problem hiding this comment.
Should it just slice all?
| x = mx.distributed.all_gather(x)[: a.shape[0]] | |
| x = mx.distributed.all_gather(x)[:] |
There was a problem hiding this comment.
That works for a single process, but i think it grows once there is more than one rank. all_gather concatenates along the first axis, so the output first dimension is group_size times the input, and [:] keeps all of it. With N ranks the input to the next iteration is 5N, then 5N^2, and the loop runs 100 times, so for N=8 it would blow up almost immediately.
Slicing back to the original first dimension keeps every iteration doing the same amount of work, which is what makes the per iteration number meaningful:
iter0 -> (5, 5)
iter1 -> (5, 5)
iter2 -> (5, 5)
That is why i went with [: a.shape[0]] rather than [0] or [:].
I could only test with a group of one here, so if you know all_gather behaves differently under a real multi rank group, say so and i will push [:] instead. Happy either way.
Proposed changes
synchronize_bench.pycrashes before it can report the all gather timing:The loop indexes with
[0]where it means to slice.all_gathergrows axis 0, but[0]removes that axis instead of trimming it, so the array loses a dimension on every pass and runs out after two:Slicing back to the original first dimension keeps the shape fixed, which is what a benchmark loop wants:
With that change the script completes and prints both numbers:
Note this is not specific to a single process.
all_gatheronly scales axis 0, so[0]drops a dimension no matter how many ranks there are, and the third iteration fails either way. I could only run it with a group of one here, so I have not exercised it undermlx.launchwith several ranks.Checklist
Put an
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes(benchmark script, so no test; verified by running it before and after)
honest note on process: freshman contributor, Claude Code helps me sweep, but i found this by running the benchmark scripts, then printed the shape at each iteration to work out why it died rather than guessing at the fix. flagging clearly that the multi rank path is untested on my machine, in case that changes how you want it written.