Repository navigation
Fix silent Metal grid truncation in softmax, logsumexp, rms_norm and layer_norm past 2^27 rows - #4717
ArupKantiDas wants to merge 2 commits into
Conversation
zcbenz
added
the
await response
|
Yes, it is, but only after one small extension: optional y/z grid sizes, default to 1. The reduction's offset, which is set from the callback, stays as a kernel argument because it needs the absolute column for its bounds check. The dispatch_col_reduce goes away, and its 4 callers in reduce.cpp use the shared helper. Should I add it to this PR, or as a follow-up if you'd rather keep this one small. |
softmax, logsumexp, rms_norm and layer_norm (and the layer_norm VJP and the looped rms_norm VJP) launch one threadgroup per row in a 1D grid. Once rows * threadgroup size passes 2^32, Metal truncates the grid width, so most rows are never computed and silently come back as zeros. Launch the rows in slices that each fit in one grid, and bind the input and output buffers at each slice's first row. Launches that already fit are unchanged and still use a single dispatch. Fixes ml-explore#4624
dispatch_col_reduce from ml-explore#4437 does the same slicing as the new helper. Give the helper optional grid_y and grid_z sizes, use it for the four column reduction dispatches, and remove dispatch_col_reduce. The reductions still set each slice's start as the column_offset kernel argument (from the bind callback), because the kernel needs the absolute column for its bounds check.
|
Please add it to this PR, thanks for looking into it! |
8c08681 to
6ca0320
Compare
|
Thanks for the confirmation. It's done, I've added the change to this PR as a separate commit and also rebased onto the latest Reduction benchmarks (float16, ms per call, M5 Pro):
|
zcbenz
removed
the
await response
Issue
For running operations like softmax, logsumexp, rms_norm and layer_norm on the Mac GPU, each row used to get 1 threadgroup launched in one long grid, set up one after the other.
The grid passed 2^32 threads at somewhere around 2^27 rows (of 128 numbers or fewer).
Metal cuts the grid off silently, because of which most of the rows were never even computed and came back as 0.
The layer_norm gradient and the long-row rms_norm gradient have the exact same problem.
Fix
The new helper in
utils.hnameddispatch_threadgroups_in_slicessplits the oversized launches into slices that actually fit, and each slice then binds the input and output at its first row, so the GPU kernels stay unchanged.The inputs that already fit still use exactly one launch, which is the same idea as the reductions fix done in #4437.
The row counts changed from
inttoint64_t, and the rms_norm gradient for short rows is also untouched, since it already uses at most about 512 teams.Test and Reproduction
First I reproduced the issue on my M5 Pro, 0.32.4, at the following boundaries: [2^27 - 1] rows correct, [2^27] rows all wrong, [2^27 + 1000] rows with only first 1000 correct.
Then 134M and 402M random rows compared against MLX's CPU result, of which there were 0 mismatches after the fix.
Added 3 new tests (softmax/logsumexp, rms_norm/layer_norm, layer_norm gradient). They fail on 0.32.4 and pass with the fix applied.
The gradient test is skipped in CI because it uses about 6.5 GB of memory.
The full suite ran with 996 of them, out of which 56 were skipped.
Fixes #4624
Benchmarks
Milliseconds per call, float16 inputs, both main and the fix built with Xcode, measured on my M5 Pro.
The cross_entropy and the attention fallback use the same launch pattern and are not fixed in this PR. Can be taken up in a later PR.