Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention

TL;DR

A mechanistic explanation for a two-year-old, well-reported failure case: training GPT-2 with Flash Attention in BF16 produces a catastrophic loss explosion after several thousand steps. The paper traces this to two intertwined phenomena — the emergence of similar low-rank representations in the attention backward pass and a systematic biased rounding error in BF16 addition during the P V matrix product when softmax probabilities saturate to exactly 1. A minimal softmax modification that prevents P from ever reaching 1 stabilizes training, validating the analysis.

Key claims

The failure is localized to a single layer (layer 2) and a small set of attention heads in that layer, identified via spectral-norm spikes in W_Q; computing those heads’ output O = P V in FP32 is sufficient to restore stability [§3.2, Fig. 3].
The numerical error enters the gradient via the backward-pass term D = rowsum(dO ⊙ O); replacing the BF16-computed O with an FP32 recomputation, or computing D via the alternative D = rowsum(dP ⊙ P) formulation, both stabilize training [§3.2].
Tiling is not the cause — disabling tiling (block size = sequence length) still produces the loss explosion, ruling out the block-wise online softmax as the source [§3.2].
The gradient error decomposes as a weighted sum of rank-1 outer products Δd_i^Q ⊗ x_i, and the rows Δd_i^Q and x_i are highly similar across training steps and token positions, so the error matrix is effectively a single low-rank direction [§3.3.1, Fig. 4, Eqn. 3].
The scalar coefficients of this low-rank direction are systematically positive — cumulative sums tracked from step 6580 to 6680 stay monotonically positive — so the low-rank error accumulates into a biased weight update rather than canceling out, driving ‖W_Q‖_2 and activations to grow abnormally [§3.3.1, Fig. 5(a), Fig. 9].
The positive bias is traced to BF16 addition rounding in the unnormalized output Õ = P̃ V: when softmax produces an exact 1 (because the pre-softmax score equals the row maximum), the product term is just a BF16 V element, and adding two same-sign BF16 numbers whose sum overflows the 7-bit significand triggers asymmetric round-up/round-down behavior that produces a negative-biased rounding error on the negative-dominated value features [§3.3.2, Fig. 6, Eqn. 4].
Attention sinks are mechanistically linked to the instability: sinks attract high attention scores and so are more likely to produce attention probabilities of exactly 1, which is the precise trigger for the biased BF16 rounding error in P V [§5 Discussion].
A targeted softmax modification — when a row’s pre-softmax maximum is repeated, subtract an extra ε = 2^{-5} (or floor it to a strictly negative value) so the maximum exponent argument is strictly negative and all P elements stay strictly less than 1 — is mathematically equivalent to standard attention in exact arithmetic (shift-invariance of softmax) and empirically restores stable training on the original failure case [§4, Algorithm 1, Fig. 7].
Results reproduce across NVIDIA A100, RTX 4090, and Huawei Ascend 910B, indicating the bug is in the BF16 arithmetic, not a vendor kernel quirk [§5].

Method

The investigation is diagnostic, not a new algorithm: the authors reproduce the nanoGPT Issue #303 failure deterministically by recording and replaying the exact batch sequence that produced the original loss explosion (GPT-2, 12 layers, 768 dim, 1024 context, 4× A100, BF16 mixed-precision, AdamW), then narrow the failure step by step. They first localize it to one layer (spectral-norm spike in layer 2), then to specific attention heads (heads 1, 7, 8, 9, 11, 12 in that layer), then to the backward-pass D term, then to O from the forward pass, then to the unnormalized product Õ = P̃ V, and finally to BF16 rounding at token positions where P[i,j] is exactly 1 (i.e. j = argmax).

The analytical decomposition Δ∇W_Q = Σ_i (1/√d) Δd_i^Q ⊗ x_i (Eqn. 2) makes the rank-1 structure explicit, and tracking the cumulative sum of Δd_i^Q[k] across steps shows the positive bias. The BF16 addition trace in §3.3.2 walks through one specific overflow → right-shift → round-to-nearest event where the sticky bit (activated by accumulated small FP32 residuals) forces a round-up on negative operands, producing the asymmetric negative error.

The fix (Algorithm 1, with one extra magenta line in the standard FA tiling algorithm) dynamically detects rows whose maximum is repeated and subtracts ε = 2^{-5} from the safe-softmax shift, guaranteeing e^{s_{ij} - m_i} < 1 for all entries. ε is chosen empirically to be small enough not to underflow and large enough that the result doesn’t round back to 1.

Results

The original failure (nanoGPT Issue #303-style loss spike after several thousand steps) is reproduced deterministically and then eliminated by the proposed softmax modification (Fig. 7); training is otherwise unchanged.
FP32 ablations isolating individual computations confirm the localization: FP32 O in the forward pass, FP32 D recomputation in the backward, or FP32 O = P V for only the 6 outlier heads in layer 2 all restore stability [§3.2].
The biased rounding analysis is supported numerically by a step-by-step BF16 addition trace (-0.85546875 + small residual + -0.0859375 rounds to -0.94140625 instead of the FP32-exact -0.94189453125, i.e. a negative error of -0.000487) that propagates into the visibly negative-stepping cumulative error in Fig. 6(c).
No degradation reported from the softmax fix at the analyzed scale; it integrates into standard FA tiling without altering the backward pass.

Why it’s interesting

This paper resolves a specific open question in the Training stability at scale cluster: prior work (A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training) noted that attention sinks correlate with low-precision instability and proposed Gated Attention as a fix, but didn’t pin down the arithmetic mechanism. This paper supplies that — sinks produce attention probabilities of exactly 1, which is the precise input pattern under which BF16 addition in P V develops a systematic negative bias on negative-dominated value features. It gives the StreamingLLM and the Discovery of Attention Sinks line a direct numerical interpretation: sinks aren’t just an attention-pattern curiosity, they’re the trigger for a specific floating-point pathology. It also complements the QK-Clip: Taking Muon Further on the Scaleup Journey / MuonClip stability fixes by explaining why spectral-norm clipping on Q/K helps — clipping suppresses the saturation that would otherwise produce the P = 1 events. The fix itself is the smallest possible diff (one line in safe-softmax), which is a clean alternative to the architectural rewrites (Gated Attention, QK-Norm) for teams that can’t change their architecture mid-training.