The authors pretrained a 450 million-parameter transformer for 50 billion tokens, comparing BF16 FlashAttention-3 with FP32 attention and modified backward passes.
Restoring softmax gradient symmetry prevents late BF16 attention instability
A row-wise projection corrected large gradient errors in FlashAttention-3 and matched FP32 attention during transformer pretraining.
Top University
Junlin Chen · Daize Dong · Huanwei Di · Haolong Jia · Jiawei Wu · Haotian Xie · +6 more
Rutgers University · Carnegie Mellon University · Oracle · New York University · MBZUAI
Research Digest··3 min read
Chen and colleagues trace a late-training failure in BF16 FlashAttention-3 to two numerical errors, including a previously overlooked violation of a softmax conservation law.
Why this paper
From Carnegie Mellon University and 4 others
In one line
FlashAttention-3 in BF16 causes gradient blow-up due to broken softmax score gradient sum conservation, fixable with gauge projection.
What we could check
- ·No code link found
- ·No weights link found
- ·No dataset link found
- ·No compute details found
- ✓Limitations stated by the authors (2 noted)
- ·No benchmark numbers found
Observed from the paper text and links we have. Absence here means we did not find it, not that it does not exist.
§