AI/ ai research · transformers · low-precision training · flash attention

Why Flash Attention Breaks During Low-Precision Training

A paper traces Flash Attention training collapses to a compounding rounding-error loop - and ships a minimal code fix.

Researchers have identified the root cause of a Flash Attention training failure that has been quietly derailing AI experiments for years.

Low-precision formats - FP8, BF16 - cut the compute cost of training transformers significantly, which is why nearly every serious training run now uses them. But they carry a known risk: loss explosions, where training error spikes suddenly and the run is effectively over. A new paper provides the first mechanistic explanation for why this specific failure happens when Flash Attention is in the stack. Two phenomena are to blame: the attention mechanism begins converging toward similar low-rank representations, while biased rounding errors in low-precision arithmetic compound on each other in a feedback loop that corrupts weight updates until training collapses entirely.

That matters because Flash Attention and low-precision formats are now default infrastructure for large-model training, not optional extras. Teams have historically responded to these instabilities with trial-and-error - raising precision mid-run, tuning learning rates, or restarting with a different seed and hoping. A mechanistic explanation turns a black-box failure into a diagnosable bug and points directly at where to intervene.

The researchers back the analysis with a minimal modification to Flash Attention that breaks the rounding-error cycle and stabilizes training, with code published publicly. Whether the patch finds its way into the widely-used upstream library is a separate question the paper, reasonably, leaves open.

TR

The Revision

Written by an AI system from the public sources credited above. How we write →