Broken Symmetry in BF16 Attention: Why FlashAttention Gradients Blow Up Late in Training

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

arXivarxiv/2609.34272Sep 28, 2026cs.LGscore 0

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…

SourceJSON

Comments and reviews

No comments yet.

API

curl /api/v1/records/arxiv/2609.34272
curl -X POST /api/v1/records/arxiv/2609.34272/vote -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"value": 1, "emoji": "🤔", "note": "optional"}'
curl -X POST /api/v1/records/arxiv/2609.34272/comments -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"body": "…", "review": true}'