Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention (huggingface.co)

Researchers analyzed why training transformers with flash attention in low-precision formats (e.g., FP16/FP8) sometimes produces catastrophic loss explosions and offer both a mechanistic diagnosis and a simple fix. They show the failure isn’t random: attention layers can collapse to similar, effectively low-rank query/key representations, and low-precision arithmetic introduces biased rounding errors. Those two effects interact—low-rank activations amplify bias in intermediate sums, which then skews gradients and weight updates, creating a runaway cycle that derails training. The paper validates this causal chain experimentally and proposes a minimal change to the flash attention implementation that reduces the rounding bias and stabilizes training. That fix is lightweight but important: it makes low-precision flash-attention training reliable without reverting to higher precision or expensive workarounds. For the AI/ML community, this clarifies a longstanding instability, informs safer mixed-precision and quantized training recipes, and highlights that numerical bias and representation rank must be considered when designing attention kernels and hardware-aware optimizations. The result should accelerate safe adoption of low-precision routines for large-scale transformer training, improving compute and memory efficiency while avoiding hidden numerical pitfalls.