A Silent BF16 Bug in FlashAttention-3 Can Derail Entire Pretraining Runs
A new fix called GProj restores a broken conservation law in attention gradients, stopping late-training blowups that FlashAttention-3 silently causes in BF16.
- FlashAttention-3 in BF16 silently corrupts gradients late in pretraining, causing 1000x gradient norm spikes with no NaNs.
- Root cause: BF16 rounding breaks the zero-sum conservation law in softmax score gradients, leaking the mean key into query gradients.
- The leak scales with key magnitude and attention sharpness, so it only manifests after tens of billions of tokens.
- A known forward FMA bug in FA3, previously dismissed as a NaN edge case, also contributes to the blowup.
- GProj restores the zero-sum property with two rank-one corrections per row, cutting median query gradient error from 219% to 0.34%.
- Costs 4.7% more time per step and matches FP32 attention loss in from-scratch runs where FA3 destabilizes.
A BF16 gradient error can derail FlashAttention-3 pretraining
A team pretraining a 450-million-parameter transformer on 50 billion tokens saw the gradient norm rise 1,000-fold after 25 billion stable tokens. The final loss was 0.2 nats higher than a matched run using FP32 attention, where nats are natural-log units of cross-entropy. No NaN or infinity appeared. A paper attributes the failure to separate numerical errors in the forward and backward passes of BF16 FlashAttention-3.
BF16 halves storage relative to FP32 and can improve accelerator throughput, while fused attention kernels reduce memory traffic by combining several operations. BF16 retains FP32’s exponent range but stores only seven fraction bits, making rounding errors more likely. Training stacks often accept those errors as small deviations from FP32. This study shows how one deviation can accumulate silently across transformer layers and surface late in a long pretraining run.
Two error paths converge
To isolate the source, the authors recomputed the attention backward pass for two layers in FP32 while leaving the rest of the model unchanged. That intervention removed almost all of the excess gradient and identified two error channels:
- Forward softmax error:
This story is for Pro members
You've reached the end of the free preview. Upgrade to AlphaSignal Pro to read the full article - and everything else behind the paywall.
Comments
