Bfloat16 errors cause large transformer training instabilities fixed by gauge projection
Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training
Machine LearningDistributed, Parallel, and Cluster Computing
Summary
Training large AI models using a faster, lower-precision number format called BF16 sometimes causes sudden huge errors in learning, even without crashing. The authors found that small mistakes in handling certain math steps in attention calculations lead to these errors late in training. They fixed the problem by correcting these small errors with a method called gauge projection, which brings the results back in line with full precision without slowing down training too much. This fix helps training large transformers remain stable and accurate.
What this means in practice
- •For machine learning engineers: Prevent training instabilities in large transformer models using low-precision BF16 attention kernels by applying gauge projection corrections.
- •For deep learning framework developers: Improve or patch fused attention implementations like FlashAttention to handle BF16 rounding errors and avoid gradient blow-up during large-scale pretraining.
Authors
Junlin Chen, Daize Dong, Huanwei Di, Haolong Jia, Jiawei Wu, Haotian Xie, Mingkai Zheng, Yang Li, Leshang Chen, Huishu Wang, Eric P. Xing, Hongyi Wang
Abstract
BF16 is now standard in large-scale pretraining, including in fused attention kernels such as FlashAttention, and these kernels are widely trusted. When we used FlashAttention-3 to pretrain a 450M-parameter transformer on 50B tokens, however, we ran into a problem: training was healthy for 25B tokens, then the gradient norm grew a thousandfold and the loss ended 0.2 nats above FP32 attention, without a single NaN. Recomputing the attention backward of just two layers in FP32 removes almost all of the excess gradient. Part of the cause is known: a fused multiply-add in the forward softmax, so far treated as an extreme-input NaN case and never fixed in FlashAttention-3. Repairing it stops the blow-up, but the query gradient is still wrong by more than its own size, and training still drives attention logits to thousands of times their size under accurate gradients. The remaining error comes from a broken conservation law. The softmax score gradient sums to zero along every row, which makes the query gradient blind to where the keys sit as a group; rounding it to BF16 leaves a small nonzero sum that leaks the mean key into the gradient, and the leak grows exactly as late training makes keys large and attention sharp. We introduce GProj (gauge projection), which restores the zero sum after the cast with two rank-one corrections per row. It cuts the remaining median query/key gradient errors from 219%/13% to 0.34%/0.37%, on par with FP32 attention, for 4.7% more time per training step. In matched from-scratch runs it trains to the same loss as FP32 attention, while FlashAttention-3 and key smoothing both destabilize.