Fused Metal KL-divergence kernel overflows (−3.39e38) in backward pass during DWQ for models >~1B

[yamyr]

Description

While running DWQ (mlx_lm.quant.dwq) on models larger than ~1B parameters on an M-series Mac (M5 Max, 137 GB), the fused Metal KL-divergence kernel's backward pass overflows to -3.39e38 (≈ -FLT_MAX) on certain training batches. The initial validation loss computes fine; the first training step corrupts the parameters, and training diverges immediately.

Environment

• mlx-lm 0.31.3, macOS (Apple Silicon M5 Max, 137 GB unified memory)
• Observed during DWQ distillation (8-bit teacher → 4-bit student), e.g. Qwen2.5-1.5B-Instruct; smaller (<~1B) models unaffected on the same setup

Where

mlx_lm/tuner/losses.py:

• _make_kl_backward_kernel() — the fused backward kernel
• _kl_div_loss.vjp — dispatches it as grid=(1024, cotangent.size, 1), threadgroup=(1024, 1, 1), templated on ("V", logits_q.shape[-1])
• can_run_metal() — returns mx.default_device() == mx.gpu and mx.metal.is_available(); kl_div_loss() falls back to the nn.losses.kl_div_loss reference path when it is false

Reproduction signals

• Triggered by the default fused path; also triggered independently by --grad-checkpoint and by --batch-size > 1.
• Forcing the numerically stable reference KL path avoids it entirely:

import mlx_lm.tuner.losses as _losses
_losses.can_run_metal = lambda: False   # use reference KL path (stable)
from mlx_lm.quant.dwq import main
main()

With that one-line monkeypatch, DWQ completes normally and the resulting weights pass our perplexity gate (teacher 8.34 < DWQ 9.45 < plain-affine 9.86 on our pilot).

Possible mechanism (hypothesis — not isolated)

Offered only as a starting point; we have not confirmed this and would not want it taken as a diagnosis.

Both KL kernels pad the tail block with a -1e30 sentinel:

vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;

and later rescale partial sums with metal::fast::exp(prev_max_q - max_q). For a thread whose whole slice falls in the padding region, that argument is ≈ -1e30. Metal's metal::fast:: math functions are only defined over a limited input domain, so if fast::exp does not return exactly 0 there, sum_exp_q — and therefore lse_q = max_q + fast::log(sum_exp_q) — is garbage, and the per-element write

out[offset + j] = static_cast<T>(c * (fast::exp(vals_q[j] - lse_q) - fast::exp(vals_p[j] - lse_p)));

saturates. That would also explain why the forward pass looks fine: it reduces to one scalar per row, whereas the backward writes a full V-wide gradient, so a corrupted lse is far more consequential.

The same padding-and-rescale pattern appears in the JS kernels in this file, if the theory holds up.

Expected

The fused Metal KL kernel's backward should match the reference implementation numerically (or clamp/guard against the overflow), for models above ~1B params.

Notes

Happy to provide the exact DWQ invocation and gate script — we've documented the workaround in our internal recipe and can share timings/configs. We can also test candidate patches on the M5 Max setup where this reproduces.

[AxelNoun]

Following up on this: the evidence points to a synchronization bug in the
fused kernels rather than the -1e30 sentinel itself. The inter simdgroup
reduction reuses the same threadgroup buffer for the max exchange and the
sum exchange, with no barrier between the read back of the maxes and the
write of the partial sums, so a fast simdgroup can overwrite a slot before
a slow one has read it. That collapses the row logsumexp and saturates the
gradient, which matches the intermittent, load dependent behavior (batch
size, grad checkpointing) much better than a fast math domain issue would,
since that one would be deterministic. Another data point against the
sentinel as the trigger: the forward pass performs the same out of range
fast::exp on every row and computes correctly. The sentinel call sites are
hardened anyway in a separate commit.

I do not own Apple hardware, so measurements come from GitHub macOS arm64
runners (virtualized Apple M series with Metal): running the backward
twice on identical inputs gives 50/50 nondeterministic on the worst KL
configuration before the fix (float32, shape (512, 151936)) and 0/50
after, with the pipelined stress case going from 3/400 to 0/400. While
validating that, the diagnostic also exposed a second deterministic bug
in the JS kernels: naive log(1+exp(log-ratio)) overflows at logit scale
20; that is fixed with a stable softplus in its own commit. Fix and full
logs in #1714.

You offered to test candidate patches on the M5 Max where this
reproduces. The branch installs with

pip install git+https://github.com/AxelNoun/mlx-lm.git@fix/kl-kernel-threadgroup-race

and the determinism discriminator lives in debug_1700/check_kl.py on the
ci and ci-fix branches of the same fork: run it once on each branch to
compare before and after on your machine. Your DWQ recipe, unchanged and
without the can_run_metal monkeypatch, is the definitive check: no
-3.39e38 values, finite loss at every step, and final perplexity in line
with the 9.45 you measured on the reference path.