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

Source: arxiv.org/abs/2609.34272

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…

JSON

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}'