Fine at 32 Bits, Broken at 4: Where the Softmax Bug Shows Up
Last week we re-tested a fix for a rounding bug in the loss function almost every model trains with. The bug was real; the accuracy gain was small. The obvious next question is whether it affects your training. We mapped where the same arithmetic turns up. In standard 32-bit training it almost never matters. In the 8- and 4-bit formats that newer chips use to make training cheaper, it hits most of every batch, and the place it hurts most is attention.
01 — one rule
Fewer bits, earlier cutoff
A quick recap. Training nudges a model towards the right answer, and the nudge is worked out as "1 minus the probability the model gave that answer". Once that probability is so close to 1 that the number format stores it as exactly 1, the nudge becomes exactly zero. The model stops learning from that example, even though there was still a little left to learn.
How close is "close enough" depends only on how many digits the format keeps. That gives one rule for every format. We had already measured the cutoff in 16-, 32- and 64-bit, and each landed within a few percent of the rule, so we trust it for the 8- and 4-bit formats we haven't run.
Read the confidence labels as "the right answer gets no signal once the model is this sure". In 32-bit that only covers answers the model is already all but certain of, which have almost nothing left to learn anyway. In 4-bit it covers anything the model gets right with 80% confidence: most of a normal batch.
02 — at 32 and 16 bits
It matters only where confidence keeps climbing
Everyday fine-tuning and pretraining are fine. In our earlier tests, PyTorch on an NVIDIA GPU worked out its loss in 32-bit even when the model ran in 16-bit, so the earlier 16-bit cutoff never appeared. We judged our own fine-tuning runs unaffected on that basis.
Three kinds of training keep pushing confidence up without limit, and those do reach the cutoff:
Training many passes until the data is memorised. Measured: the fix gained between 1.5 and 7 points of accuracy depending on the model and format, and cost points at one setting. A student model that starts as an exact copy of its teacher. Measured: the standard loss pushed it away from the teacher in every format, even 64-bit. Image–text matching with a very sharp setting, as in CLIP-style training, where large margins are routine. Untested; we expect the effect to be small.
03 — at 8 and 4 bits
Attention is where it hurts most
Attention, the part of a language model that decides which earlier words to look at, uses the same probability step in every layer. When it learns, the nudges across each row of attention should add up to exactly zero. The fast attention code most training uses (FlashAttention) relies on that, but it works out the nudges from probabilities that have already been rounded down to the low-precision format, so the row no longer adds up.
Worked out from the same rounding rules: in 16-bit the leftover is about one part in 500 of the signal, which training shrugs off. In 4-bit it is about a quarter, which is not noise any more. This matches a problem reported in FlashAttention the week after the original paper, and explains why fixing it is described as a step towards practical 4-bit training.
The cutoff shows up here too. Many models have rows where almost all attention goes to one spot, often the first word. The most extreme of those rows sit past even the 32-bit cutoff, so they get zero or noise instead of their true, tiny signal. Harmless in 16- and 32-bit; not in 4-bit.
We haven't run this yet. The test fits our existing setup: round attention to 16, 8 and simulated 4 bits, and compare the current calculation with two corrected ones.
04 — the map
Who should care, and how sure we are
A few entries need a sentence each.
Reinforcement learning depends on one conversion. If the model's log-probabilities are kept in 16-bit, any token it is more than 99.6% sure of gets no signal. That is fine when the token was right and a problem when it was wrong: a confidently wrong answer can't be pushed down. Converting the log-probabilities to 32-bit removes it.
Sigmoid losses split in two. PyTorch's BCEWithLogitsLoss has the same bug with two classes. logsigmoid, which DPO and most reward models are built on, calculates it in a way that never rounds to zero.
LoRA adds nothing of its own. The adapters don't contain this kind of calculation; they inherit whatever the loss above them does. Speculative decoding doesn't train, so there is no learning signal to lose.
05 — practical lessons
What to check in your own setup
Find out what precision your loss really runs in. On PyTorch with an NVIDIA GPU, the standard cross-entropy is 32-bit inside even for 16-bit models. Fused or custom loss code, compiled models and JAX need checking one by one. We haven't surveyed them.
Moving to 8- or 4-bit training? Assume the bug is active. That includes the probabilities inside attention, not just the final loss. Keep these steps in higher precision, or use code that applies the fix, and check that it does.
Keep a 32-bit copy of the weights. In our tests, pure 16-bit weights stopped training altogether, a separate rounding problem that standard mixed-precision setups already avoid.
Don't start a student exactly at its teacher, or use the fixed loss if you do. We measured the drift in every format.
Doing reinforcement learning? Convert log-probabilities to 32-bit. It's one line, and it means a confidently wrong token can still be corrected.
At 32 bits this bug is a footnote. At 4 bits it is in every layer, and attention is the place to fix first.