Conversation
On Metal, CrossEntropy::use_fallback returned true and both eval_gpu overloads threw NYI, so mx.fast.cross_entropy always ran the unfused logsumexp - take_along_axis graph. Add the forward and VJP kernels and enable them, following the CUDA implementation from ml-explore#3947. Forward: one threadgroup per row, single-pass online logsumexp with float32 accumulation for every input dtype. The loss is formed as (max - x_t) + log(normalizer) so the two close values are subtracted first. The host shrinks the threadgroup to ceil(V / N_READS) rounded to a SIMD multiple, so short rows take one iteration of the same looped kernel; the cross-SIMD reduction only reads the slots that were written. VJP: exp((x - x_t) - loss) is softmax(x), so the backward pass needs no reduction and the one-hot target is never materialized. When the logits buffer can be donated the gradient is written in place, with a device memory barrier between the reads of x_t and the first write. Negative targets wrap, matching the take_along_axis fallback this replaces. The JIT library name is passed explicitly because deriving it from the kernel name would drop the "cross_" prefix. The float16/bfloat16 tolerances in test_cross_entropy tighten to 1e-3 on the GPU. The fallback fails this because its logsumexp runs in the input dtype; the fused kernels are within ~4e-6 for all three dtypes. The CPU path keeps the old tolerances.
nn.losses.cross_entropy only took the mx.fast.cross_entropy path when CUDA was available. The Metal kernel exists now, so gate on the default device being the GPU alone. The half precision example in the docstring now shows what the fast path returns: the loss is cast back to the logits dtype, so it is bfloat16, not float32.
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
1 participant
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.
Metal kernels for
mx.fast.cross_entropy, same idea as the CUDA ones from #3947. On Metal this currently falls back tologsumexp - take_along_axis.The fallback does the
lse - x_tsubtraction in the logits dtype and only then casts to float32, so with bf16 logits you lose bits right there. The kernel does everything in fp32. Max abs error vs float64 numpy on random(4, 7, 8192)logits:I tightened
test_cross_entropyto 1e-3 for fp16/bf16 on GPU, it fails on main with Metal. I couldn't run it on CUDA, no hardware here, but that kernel accumulates in fp32 as well.For the backward I use
exp((x - x_t) - loss), which is the softmax, so there is no second reduction and I never build the one-hot. That is where the time and memory go. Times in ms on an M3 Pro, mean of 20, scripts at the bottom:Forward is bandwidth bound either way and the differences there are noise. fwd+bwd is about 3x at the big vocab sizes and 2x at V=1024, peak memory ~1.5x lower.
I went with one threadgroup per row,
ceil(V / N_READS)threads rounded up to a simd multiple and capped at 1024. So only the looped variant, I didn't write the block kernel that softmax has for small V. If you want it for the small V case I can add it.Negative targets wrap, matching
take_along_axis. The CUDA kernel doesn't check that and reads out of bounds.One thing to flag: with a
-inftarget logit the gradient row is now NaN where the unfused path gave finite values, because the VJP reconstructs the softmax from the loss. As far as I can tell from reading the CUDA kernel it does the same, so I left it.The second commit drops the
mx.cuda.is_available()check innn.losses.cross_entropyso Metal also goes throughfast.cross_entropy. Also fixed the docstring example there, it showed float32 output for a bfloat16 input.python/tests/run.pypasses, also underMLX_METAL_JIT=ON.