One barrier short

The kernels are a sensible optimization. A KL divergence between two 151,936-way distributions, once per token, is a lot of memory traffic if you materialize the log-softmaxes, so since v0.26.0 in July 2025 mlx-⁠lm has computed it in one fused pass on the GPU, forward and backward, for KL, and from v0.26.3 for Jensen-Shannon.1 In August a user running DWQ on an M5 Max reported that the backward pass sometimes returned -3.39e38, about the most negative value a 32-bit float can hold, and that the first training step corrupted the weights of models above roughly a billion parameters. Routing the loss through plain MLX ops made the problem go away.

The report came with a careful hypothesis about a padding sentinel and the input domain of fast::exp, offered as a starting point and labelled as one. Two days later a contributor, AxelNoun, posted a different one: the kernels reuse their shared scratch buffer for two exchanges with no barrier in between, which fits a failure that comes and goes with batch size and gradient checkpointing far better than a deterministic math-domain bug would. That diagnosis was right. What follows is what it took to prove it on a machine, and why a one-line fix still took until the end of October to land.

Sixty-four floats

Each call launches one threadgroup of 1,024 threads per row of logits. A thread strides through the row four values at a time, keeping a running maximum and a running sum of exponentials relative to that maximum, the usual online log-sum-exp. At the end the 1,024 partial results have to be combined. Inside a SIMD group of 32 lanes that is free: simd_max and simd_sum are register shuffles.2 Across the 32 groups it needs memory, and the kernels declare exactly enough for one exchange, threadgroup float shared[32 * 2]: two floats per group, one for the student distribution q and one for the teacher p.

Here is the exchange as it shipped, from the KL forward kernel. The other three kernels have the same lines.

// Share the maxs across the threadgroup
float prev_max_q = max_q;
float prev_max_p = max_p;
max_q = simd_max(max_q);
max_p = simd_max(max_p);
// step 1: lane 0 of group g writes into slots 2g and 2g+1
if (simd_lane_id == 0) {
    shared[simd_group_id * 2 + 0] = max_q;
    shared[simd_group_id * 2 + 1] = max_p;
}
// step 2
threadgroup_barrier(mem_flags::mem_threadgroup);
// step 3: lane L reads the maximums of group L
max_q = shared[simd_lane_id * 2 + 0];
max_p = shared[simd_lane_id * 2 + 1];
max_q = simd_max(max_q);
max_p = simd_max(max_p);
// nothing here
// Share the sum_exp across the threadgroup
sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
sum_exp_q = simd_sum(sum_exp_q);
sum_exp_p = simd_sum(sum_exp_p);
// step 4: the same slots
if (simd_lane_id == 0) {
    shared[simd_group_id * 2 + 0] = sum_exp_q;
    shared[simd_group_id * 2 + 1] = sum_exp_p;
}
// step 5
threadgroup_barrier(mem_flags::mem_threadgroup);
sum_exp_q = shared[simd_lane_id * 2 + 0];
sum_exp_p = shared[simd_lane_id * 2 + 1];
sum_exp_q = simd_sum(sum_exp_q);
sum_exp_p = simd_sum(sum_exp_p);

Nothing orders step 3 in one group against step 4 in another. That is the whole bug.

A barrier is a rendezvous, not a lock

It is easy to read the barrier in step 2 as protecting the array. It does not. A threadgroup barrier means every thread has to arrive before any thread leaves. Once they have all left, each group runs at its own pace again, and nothing orders them until the next barrier. So a fast group can finish step 3, rescale, do its simd_sum, and have lane 0 write its partial sum into slot 2g while a slower group is still on its way to reading slot 2g as a maximum.

group 7, fast                       group 20, slow
write max_q, max_p into [14], [15]  write max_q, max_p into [40], [41]
======== barrier, step 2: all 32 groups arrive, all 32 leave ========
read [0..63] as maximums
simd_max: row maximum in hand
rescale, simd_sum
write sum_q, sum_p into [14], [15]  <- slot 14 now holds a partial sum
                                    read [0..63] as maximums
                                    lane 7 gets group 7's SUM, not its max
                                    simd_max folds it into the row maximum
                                    this group's normalizer is now wrong
======== barrier, step 5 ========

The slow group now has the wrong maximum for the row, so the normalizer it computes from it is wrong, and so is that row’s loss or gradient. Whether the row actually comes out wrong depends on which slot was overwritten and what it held. A partial sum that happens to sit below the true maximum changes nothing, which is why some overwritten rows are still correct and why the bug is intermittent rather than constant. In the backward kernel the wrong normalizer multiplies into a full vocabulary-wide gradient row, which is how AxelNoun connected it to the saturated -3.39e38 gradients in the report.

The Metal Shading Language specification is plain about this. Section 6.10.1 of version 4.1 gives the barrier one ordering guarantee, that all threads reach it before any continue; what memory it fences is set by the flag, and its Table 6.13 describes a barrier with mem_none as an execution barrier only, with no memory fence. An unsynchronized read and write of the same location by different threads is a data race by definition. MLX core’s own log-sum-exp kernel, logsumexp.h, does the same two exchanges with two separate arrays, local_max and local_normalizer, and never reuses a slot.

Catching it in the act

A diagnosis that fits is not the same as a diagnosis that is true, so the first thing I wanted was to watch the overwrite happen. The obvious instrumentation, storing every value a thread reads to a debug buffer as it reads it, mostly made the bug disappear: 0 and 18 events across two runs that should have shown hundreds, and no wrong rows. Extra memory traffic changing whether a bug appears is itself typical of a timing race, but it is not proof.3

The version that worked keeps everything in registers. Each thread holds the two values it read in step 3, lane 0 of each group holds what its group wrote in steps 1 and 4, and nothing is stored until the kernel is otherwise finished. Afterwards every read is compared with what the slot’s owner wrote as its maximum.

4,982 reads returned something other than the owner’s maximum, and every one of them was, bit for bit, the partial sum that owner had written in step 4. Every wrong row had an overwrite in it. No row went wrong without one. With one barrier added after step 3, the count is zero everywhere.

How wrong, how often

Then the public function, called the way a user would call it, on identical inputs 30 times per setting:

Wrong means a loss off by more than 1e-4 relative or a gradient entry off by more than 1e-4 absolute, 1e-2 in bf16. The bf16 rows are flagged on their gradients: a bf16 loss error near 0.4% is rounding, and the same figure shows up in runs with no wrong rows. Two runs are shown for each bf16 shape, and the worst errors are the worst across both.

Two things in that table explain the fifteen months. First, the tests. tests/test_losses.py checked the kernels on inputs of shape 2 × 4 × 4000, which is 8 rows of 4,000: zero wrong in 30 calls, and 1 in 30 at 16 rows of 32,000. The window only opens wide at vocabulary-sized rows, where each thread’s share of the work is long enough for the groups to drift apart. Second, the dtype. DWQ runs the loss in bf16 by default, and on this M1 Pro bf16 was hit in 2 of 230 calls at 1,024 rows and in 2 of 130 at 4,096 rows, close to DWQ’s default batch of four 1,025-token sequences.4 fp32 is hit most of the time. A bug that passes its tests and corrupts one gradient in fifty looks, from the outside, like noise.

Real logits behave the same way. With Qwen3-0.6B in bf16 as the teacher, its own 4-bit, group-size-64 quantization as the student, 4 × 1,024 tokens of the repository’s Markdown as input, and the loss called exactly as DWQ calls it, kl_div_loss(0.5 * student, 0.5 * teacher): in fp32, on 1,024 of those tokens, every one of 30 calls had wrong rows, 340 in all, and the batch loss DWQ would have trained on was off by more than 0.1% in 18 of 20 calls, by 4.0% at worst. In bf16, on all 4,096, 1 of 20 gradients was off by 5.4%. DWQ makes one such call per step, so a run of a few hundred steps absorbs a handful of corrupted gradients on this chip, and many more on an M5 Max.

Cause, not correlation

The last thing a race diagnosis needs is to rule out everything else that looked suspicious in those kernels. There were two candidates: the slot reuse, and a barrier in the KL forward kernel’s final reduction that used mem_none where a write-then-read through shared memory calls for mem_threadgroup. Same inputs, same 30 calls, one change at a time. The p-values are two-sided Fisher exact tests against upstream.

The row that settles it is separate arrays. Give the sums their own 64 floats and add no synchronization whatsoever, and every error disappears. The reuse is the cause, not a side effect of extra waiting. The mem_none flag sits outside what the spec guarantees, and fixing it alone changed nothing measurable. The Jensen-Shannon kernels behave identically: 23 of 30 upstream, 0 of 30 with the barrier.

The fix

One barrier, after the maximums have been read and before the sums are written, in each of the four kernels:

         max_q = simd_max(max_q);
         max_p = simd_max(max_p);

+        // All threads must read the maxs before the sums overwrite them
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+
         // Share the sum_exp across the threadgroup
         sum_exp_q *= metal::fast::exp(prev_max_q - max_q);

The pull request, #2010, also switches the KL forward kernel’s final-reduction barrier from mem_none to mem_threadgroup, because that step does write shared memory and then read it, and the matching step in the JS kernel already used the stronger flag. That change is spec-driven; no measurement shows it. The diff is 44 lines added and 1 removed across two files, and most of the added lines are the test: fp32 logits of shape 128 × 151,936, loss and gradient checked against the CPU path at a relative tolerance of 1e-4, then ten more calls that have to be bit-identical to the first. Against main, both new tests failed in all five runs I tried. With the barrier, they pass. Speed is unchanged within noise: across 16 configurations, the fixed kernels averaged 0.1% faster than main, inside a run-to-run spread of up to 24% between two copies of main.

michalk8 approved it on October 27 and merged it on October 29, closing #1700. As of this writing the fix is on main and not yet in a tagged release; v0.32.2, from October 16, still has the race. One thing #2010 does not touch: at a logit scale of 20 the Jensen-Shannon loss still overflows to non-finite values on main, and so does the plain MLX reference at that scale. That is a separate, deterministic bug, one that AxelNoun’s original PR had also fixed with a stable softplus. It has not been fixed on main.

Run it yourself

The whole reproduction is twenty lines and a few seconds on any Apple GPU. It calls the public function on identical inputs 20 times and checks each result against the plain-MLX reference. A correct kernel gives one answer, the right one, every time.

import mlx.core as mx
from mlx_lm.tuner.losses import kl_div_loss

mx.random.seed(0)
q = mx.random.normal((512, 151936)) * 4  # student, Qwen-sized vocabulary
p = mx.random.normal((512, 151936)) * 4  # teacher

lq = q - mx.logsumexp(q, axis=-1, keepdims=True)
lp = p - mx.logsumexp(p, axis=-1, keepdims=True)
ref = mx.sum(mx.exp(lp) * (lp - lq), axis=-1)
ref_grad = mx.exp(lq) - mx.exp(lp)

answers, bad = set(), 0
for i in range(20):
    kl = kl_div_loss(q, p)
    grad = mx.grad(lambda x: kl_div_loss(x, p).sum())(q)
    err = mx.max(mx.abs(kl - ref) / mx.abs(ref)).item()
    gerr = mx.max(mx.abs(grad - ref_grad)).item()
    answers.add(bytes(memoryview(kl)))
    bad += err > 1e-4 or gerr > 1e-4
    print(f"call {i:2d}: loss rel error {err:.1e}, grad error {gerr:.1e}")
print(f"{bad}/20 calls wrong; {len(answers)} different losses for one input")

On October 31, against the last pre-fix main, this gave 15 wrong calls out of 20 and 13 different losses for one input, with a worst loss error of 19% and a worst gradient error of 0.28. Against main with #2010, 0 of 20 and one answer, with errors at the 1e-6 floor.

Fifteen months, in dates

I want to be careful about the lesson, because the easy one is unfair. The maintainers were carrying a queue they said outright they could not review, and from inside that queue a five-commit PR fixing three things at once, for a bug the shipped tests could not see, is hard to tell from the pile. The diagnosis never changed between August and October. The packaging did. What landed was one fix and one flag change in 44 lines, with its own test that failed on main in every run and credit in the description to the person who found the cause. The earlier PR had the fix and the tests too, plus two more fixes and 186 lines, and eight days after it was reopened the maintainers were closing pull requests to stay afloat. If you are sitting on a correct fix that was closed for capacity, split it and resend the smallest piece.

If you run DWQ

Install mlx-⁠lm from main, or wait for the first release after v0.32.2 and check that it includes #2010. Until then the fp32 path is the dangerous one, and the default bf16 path is hit occasionally on an M1 Pro and often on an M5 Max. The reporter’s workaround is unchanged: make mlx_lm.tuner.losses.can_run_metal return False before the first loss call, which routes the loss through plain MLX ops at a cost in memory and speed, the two things the fused kernels were written to save. And the cheapest self-check for any fused reduction kernel is the one in the script above. Call it twice on the same large input and compare the bytes. If they differ, you do not need a theory yet. You need a barrier.