PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Haiquan Qiu, Quanming Yao

arXiv:2510.04212v4cs.LGcs.AI

TL;DR

The paper examines a persistent GPT-2 failure in BF16 flash-attention training and traces it to similar low-rank structures and biased rounding errors. A targeted modification mitigates the rounding bias and stabilizes training, while the authors note that broader generalization remains to be investigated.

  • Problem

    GPT-2 training with BF16 flash attention can undergo sudden loss explosion after several thousand steps, while standard attention or FP32 workarounds sacrifice efficiency.

    Source: §3.1 · p. 3

  • Method

    The paper analyzes high- and low-precision flash-attention gradients to trace how similar low-rank structures and biased rounding errors affect weight updates, then tests a targeted mitigation.

    Source: §1 · p. 2, §3.3.1 · p. 5

  • Results

    The authors identify interacting low-rank representations and biased rounding errors as the failure mechanism, and report that a minimal flash-attention modification stabilizes training.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The analytical workflow is offered as a blueprint for diagnosing numerical instabilities in other architectures, scales, and low-precision formats.

    Source: §5 · p. 10

    The analysis focuses on one GPT-2 failure case; generalization to other architectures, larger scales, or FP8 needs investigation, and the mitigation may not address other instability sources.

    Source: §5 · p. 10

Abstract

from arXiv · show

The pursuit of computational efficiency has driven the adoption of low-precision formats for training transformer models. However, this progress is often hindered by notorious training instabilities. This paper provides the first mechanistic explanation for a long-standing and unresolved failure case where training with flash attention in low-precision settings leads to catastrophic loss explosion. Our in-depth analysis reveals that the failure is not a random artifact but caused by two intertwined phenomena: the emergence of similar low-rank representations within the attention mechanism and the compounding effect of biased rounding errors inherent in low-precision arithmetic. We demonstrate how these factors create a vicious cycle of error accumulation that corrupts weight updates, ultimately derailing the training dynamics. To validate our findings, we introduce a minimal modification to the flash attention that mitigates the bias in rounding errors. This simple change stabilizes the training process, confirming our analysis and offering a practical solution to this persistent problem. Code is available at https://github.com/ucker/why-low-precision-training-fails.

1 INTRODUCTION

The introduction frames low precision as a route to faster, lower-memory transformer training, then explains a persistent flash-attention failure through interacting representation and rounding effects. A targeted change that mitigates rounding bias stabilizes training.

1 INTRODUCTION

  • Low-precision formats reduce memory use and can speed training, but attention mechanisms are especially sensitive to numerical precision.The introduction identifies understanding failure mechanisms as a barrier to further precision reduction.Source: §1 · p. 1
  • Similar low-rank representations across tokens and training steps combine with biased rounding errors to accumulate biased weight updates and derail training.The resulting weight and activation spectral norms increase abnormally.Source: §1 · pp. 1–2
  • A minimal flash-attention modification mitigates rounding bias, allows low-rank weight updates to cancel, and stabilizes training.The authors present this experiment as validation of their analysis.Source: §1 · p. 2

2 PRELIMINARY

The preliminary section explains low-precision training and BF16 rounding, then describes how flash attention’s tiled computation reduces memory use while retaining forward and backward attention operations.

2.1 LOW-PRECISION TRAINING

  • Low-precision formats reduce memory use and accelerate computation; mixed-precision training commonly uses 16-bit computations with an FP32 master weight copy.The section describes BF16 as a preferred choice for large language model training because it retains FP32’s dynamic range.Source: §2.1 · p. 2
  • BF16 has 7 significand bits, so rounding can introduce errors that become biased across operations on data with a specific distribution.The usual round-to-nearest, ties-to-even method is designed to be unbiased for random data.Source: §2.1 · p. 2

2.2 FLASH ATTENTION

  • Flash attention tiles Q, K, and V and processes blocks iteratively, reducing sequence-length memory complexity from O(N^2) to O(N).The algorithm moves blocks between high-bandwidth memory and on-chip SRAM to minimize memory transfers.Source: §2.2 · p. 2
  • The forward pass uses online softmax with query and key/value blocks, accumulating unnormalized outputs before normalizing the final output block.It also computes output O and log-sum-exp statistics L.Source: §2.2 · pp. 2–3
  • The backward pass computes δ = rowsum(dO ◦ O), recomputes attention scores, and accumulates dQ, dK, and dV block-wise.The tiled backward pass maintains I/O efficiency.Source: §2.2 · p. 3

3 ROOT CAUSES OF INSTABILITY IN FLASH ATTENTION

The section traces BF16 flash-attention loss explosion to errors in the output computation and explains how biased errors combine with similar low-rank representations to corrupt weight updates.

3.1 THE FAILURE CASE OF LOW-PRECISION FLASH ATTENTION

  • BF16 flash attention produces a sudden loss explosion, whereas the stable configuration converges.The failure is reported after several thousand training steps.Source: Figure 2 · p. 3, §3.1 · p. 3

3.2 ISOLATING THE SOURCE OF FAILURE WITHIN FLASH ATTENTION

  • Disabling tiling does not prevent failure; experiments localize its origin to flash attention in layer 2 and identify BF16 errors in O as the likely source.Using flash attention only in layer 2 reproduces failure, while replacing it there restores stability; high-precision O also stabilizes training.Source: §3.2 · pp. 4–5

3.3 THE ROOT CAUSES OF TRAINING FAILURE

  • The failure combines similar low-rank representations with positively biased BF16 rounding errors, which make error accumulate instead of cancel.The analysis traces the positive coefficients to biased rounding during the ¯PV computation.Source: §3.3 · p. 5

3.3.1 CAUSE 1. SIMILAR LOW-RANK MATRICES BIAS WEIGHT UPDATES

  • Similar rank-1 components across tokens and steps, weighted by δ errors, produce gradient differences that bias weight updates.The positive cumulative sum of (δlp −δhp)[T] compounds error along the shared low-rank direction, increasing weight spectral norms and activations before loss explosion.Source: §3.3.1 · pp. 5–6
  • The output error is linked to aligned negative dO and O errors, while unit attention probabilities multiplying predominantly negative V values create conditions for biased rounding.When ¯P[T,t] = 1, adding negative BF16 values can overflow the significand; rounding up then introduces a larger negative error than the positive error from rounding down.Source: §3.3.1 · pp. 7–9
  • Computing ¯O = ¯PV in FP32 is sufficient to stabilize training, linking this product’s rounding bias to the failure mechanism.The paper reports that BF16 rounding events in this accumulation can create systematic negative errors in O.Source: §3.3.1 · pp. 7, 9

4 EXPERIMENT: MITIGATE BIAS IN ROUNDING ERROR

The modification targets repeated row maxima that make attention probabilities round to 1, adjusting softmax normalization to avoid biased rounding and restore training stability.

4 EXPERIMENT: MITIGATE BIAS IN ROUNDING ERROR

  • The targeted softmax modification prevents biased rounding from repeated maxima and restores training stability.It detects multiple identical maxima in a row of S and adjusts normalization so all elements of P̄ remain below 1.Source: §4 · p. 9
  • When a maximum is repeated, the method adjusts the normalization factor dynamically so the maximum exponent is strictly negative.For a positive repeated maximum it uses m = βrm with β > 1; for a negative repeated maximum it sets m = 0.Source: §4 · p. 9
  • The normalization change is equivalent to standard attention in exact arithmetic and is integrated without altering the backward pass.The paper reports β values from 2 to 8; smaller values risk rounding back to 1, while larger values risk underflow.Source: §4 · pp. 9–10

5 CONCLUSION

The paper explains loss explosion in low-precision flash attention through similar low-rank representations and biased BF16 rounding, and presents a targeted mitigation that restores stability.

5 CONCLUSION

  • The paper presents a mechanistic explanation of low-precision flash-attention loss explosion, tracing it to similar low-rank representations interacting with biased BF16 rounding errors.A targeted modification restores stability and validates the analysis.Source: §5 · p. 10
  • Attention sinks attract high scores and are more likely to produce probabilities of 1, linking sink behavior to biased rounding in the P̄V computation.The discussion also attributes abnormal weight spectral-norm growth to accumulated low-rank error matrices in gradients.Source: §5 · p. 10
  • The analysis is limited to a specific GPT-2 failure case; transfer to other architectures, larger scales, or formats such as FP8 remains unverified.The proposed mitigation targets the identified rounding error and may not address other numerical-instability sources.Source: §5 · p. 10
  • The authors propose applying their error-localization and root-cause workflow to other architectures and precision formats, and developing automated detection and mitigation tools.They identify FP8 training, larger models, and different architectures as possible extensions.Source: §5 · p. 10

A.1 MIXED-PRECISION BF16 TRAINING.

Mixed-precision BF16 is widely used to balance training cost and numerical range, though instability remains possible in BF16 pretraining.

A.1 MIXED-PRECISION BF16 TRAINING.

  • BF16 uses an 8-bit exponent and 7-bit mantissa, retaining FP32’s dynamic range while halving storage cost.Earlier FP16 training used an FP32 master weight copy and fixed loss scaling, but its narrower exponent range can cause gradient underflow (Micikevicius et al., 2017).Source: §A.1 · p. 14
  • BF16 mixed precision is used in large-scale language-model training, including GPT-3, PaLM, Gopher, Chinchilla, and LLaMA.Training frameworks such as Megatron and DeepSpeed integrate BF16 with ZeRO (Rajbhandari et al., 2020).Source: §A.1 · p. 14
  • About 10% of GPT-2 pretraining runs diverged under pure BF16, compared with 0% under TF32, according to Lee et al. (2024).The passage notes that BF16 improves stability but complementary stabilization may still be needed at scale.Source: §A.1 · p. 14

A.2 STABILIZING LOW-PRECISION TRAINING

Prior stabilization work modifies precision scaling, optimizers, activations, and architecture; Flash Attention instead reduces attention’s memory complexity through tiling and fused computation.

A.2 STABILIZING LOW-PRECISION TRAINING

A.3 FLASH ATTENTION

  • Flash Attention tiles query, key, and value blocks and uses fused computation with online softmax, reducing memory complexity from O(N^2) to O(N).The algorithm processes blocks while maintaining running maxima and exponential sums, without materializing the full attention matrix.Source: §A.3 · p. 15
  • Flash Attention 2 improves work partitioning and reduces non-matrix-multiplication overhead; its pseudocode absorbs the attention scaling factor α = 1/√d into Q.The passage describes the algorithm as an I/O-aware design for long-context models.Source: §A.3 · p. 15

B BF16 ADDITION

BF16 retains FP32’s dynamic range with fewer fraction bits, and its addition process introduces rounding error. For data with a specific distribution, accumulated rounding errors can become biased and contribute to low-precision training failure.

B BF16 ADDITION

  • BF16 uses 1 sign bit, 8 exponent bits, and 7 fraction bits, matching FP32’s dynamic range with lower precision.Its 16-bit representation is used in deep learning for its balance of efficiency and numerical range.Source: B BF16 ADDITION · p. 15
  • BF16 addition aligns exponents, adds significands, normalizes the result, and rounds it to fit the 7-bit fraction.Exponent alignment can discard bits shifted beyond the available precision.Source: B BF16 ADDITION · p. 15
  • The described rounding rule rounds up or down according to whether the truncated part is above or below half an LSB, with ties going to an even LSB.Source: B BF16 ADDITION · p. 16
  • For data with a specific distribution, repeated rounding errors can accumulate in one direction despite round-to-nearest, ties-to-even.The paper identifies this accumulation of biased error as a critical factor in low-precision training failure.Source: B BF16 ADDITION · p. 17

C DESIGN CONSIDERATIONS FOR MITIGATING BIASED ROUNDING ERROR IN FLASH ATTENTION

The proposed flash-attention modification targets rounding bias while avoiding changes in cases that could introduce numerical instability. The paper also reports similar representation patterns in Llama-3.1-8B, without claiming that this model will necessarily fail in BF16 training.

C DESIGN CONSIDERATIONS FOR MITIGATING BIASED ROUNDING ERROR IN FLASH ATTENTION

D SIMILAR PATTERNS DISCOVERED IN LLAMA-3.1-8B

D.1 STRUCTURALLY SIMILAR (PK)[T]⊤X[T]

  • The paper visualizes representations and similar columns for attention head 13 in layer 1 of Llama-3.1-8B using a sample sentence.Source: §D.1 · p. 17

D.2 MULTIPLE MAXIMUM IN ¯P

Llama-3.1-8B exhibits multiple maxima in attention scores, and the frequency of this phenomenon rises before the reported loss explosion. The paper describes that timing as a leading indicator of impending training failure.

D.2 MULTIPLE MAXIMUM IN ¯P

  • Multiple maximum values occur in Llama-3.1-8B attention scores.Source: §D.2 · p. 17
  • The first token’s high attention scores also exhibit attention-sink behavior.Source: §D.2 · p. 17

E.1 RELATION BETWEEN MULTIPLE MAXIMA AND LOSS

  • The frequency of multiple maxima begins to increase before loss explodes, around step 7000, and is described as a leading indicator of impending failure.The reported relationship is a strong correlation.Source: §E.1 · p. 18

E.2 DERIVING EQUATION 1

This section derives the difference between high- and low-precision query gradients. It simplifies the expression through cancellation and rewrites row-wise scaling as diagonal-matrix multiplication.

E.2 DERIVING EQUATION 1

  • The derivation targets dQhp − dQlp, substitutes definitions of dQ and dS, and simplifies the expression by canceling the dP terms.Source: §E.2 · pp. 18–19
  • For v = δlp − δhp, row-wise scaling P ◦ v is rewritten as diag(v)P, since each row i of P is scaled by vi.Source: §E.2 · pp. 19–20
  • The stabilized forward-pass algorithm requires Q, K, V, block sizes, β > 1, and ϵ > 0, and returns output O and log-sum-exp L.Source: §E.2 · p. 20

F SUGGESTIONS FOR HANDLING FAILURE IN LOW-PRECISION TRAINING

The paper proposes a workflow for diagnosing low-precision numerical instability, from isolating the faulty operation to tracing systematic arithmetic bias. This workflow can help researchers develop principled solutions and may generalize to formats such as FP8.

F SUGGESTIONS FOR HANDLING FAILURE IN LOW-PRECISION TRAINING

Related papers

Loading 2510.04212v4…