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

cs.LGSep 28, 2026

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

Authors: Junlin Chen, Daize Dong, Huanwei Di, Haolong Jia, Jiawei Wu, Haotian Xie, Mingkai Zheng, Yang Li, +4 more

Organizations: Rutgers University · Carnegie Mellon University · Oracle · New York University · MBZUAI

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.

Figures & tables

	Input: BF16 $Q,K,V,O,dO\in\mathbb{R}^{N\times d}$ ; FP32 $\mathrm{LSE}\in\mathbb{R}^{N}$ ; scale $\alpha$ ; key tiles $K_{j},V_{j}$ , query tiles $Q_{i},O_{i},dO_{i}$ .
	Output: BF16 $dQ,dK,dV$ . Accumulators and row statistics are FP32.
1	$D\leftarrow\mathrm{rowsum}(dO\odot O)$ ; $dQ^{\rm acc}\leftarrow 0$ ; $B^{\rm acc}\leftarrow 0$ , $\rho\leftarrow 0$ , $m\leftarrow 0$
2	for each key tile $j$ in parallel do
3	Load $K_{j},V_{j}$ ; $dK^{\rm acc}_{j}\leftarrow 0$ , $dV_{j}\leftarrow 0$
4	for each query tile $i$ that attends to tile $j$ do
5	Load $Q_{i},dO_{i},\mathrm{LSE}_{i},D_{i}$ ; $S\leftarrow\alpha\,Q_{i}K_{j}^{\top}$ ; $P\leftarrow\exp(S-\mathrm{LSE}_{i})$ (masked)

Algorithm 1 FA3 backward pass (black) with the GProj additions ( orange )

Method	$dQ$ (%)	$dK$ (%)	$dV$ (%)
FA3	773	85.6	0.369
QY-shift (source-faithful) ‡	1307	264	80.5
FA3-SBS	219	13.3	0.334
GProj-Q	0.342	13.3	0.334
GProj	0.342	0.371	0.334
FP32 attention (eager) †	0.384	0.371	0.165

Table 1: Full-tensor relative

L_{2}

error (%) of BF16 gradients against an FP64 reference, median over eight captured attention inputs of the 450M model. FA3-SBS and the GProj rows share identical forward outputs. †,‡ See Appendix E .

Method	$dQ$	$dK$	$dV$
FA3	0.2601 / 4794.33	0.2537 / 4.1051	0.2239 / 0.7657
FA3-SBS	0.2601 / 4794.53	0.2535 / 4.1050	0.2239 / 0.7657
GProj	0.2238 / 2.4938	0.2219 / 2.7565	0.2239 / 0.7657

Table 2: Synthetic suite: relative

L_{2}

error percentages, median / maximum over the 58 test cases, excluding zero targets separately for each gradient (

dQ/dK/dV

: 53/51/57 nonzero targets); the two analytic calibrations are reported in Appendix F .

Run (tokens)	Layers with logit scale $>10^{3}$	Logit scale L5 / L11	Layer 11 rms $Q$ / $K$	Layer 11 rows $p_{\max}>0.999$	Layer 11 $dQ$ error FA3 / cuDNN
FA3 (23.1B)	5, 11	5,436 / 5,519	71 / 78	86%	4.5 / 2.2
FA3 (33.6B)	5, 11	12,843 / 51,057	213 / 240	99.4%	253 / 5.6
FA3-SBS (33.6B)	4, 5, 11	39,205 / 49,158	216 / 227	99.6%	85 / 2.5
Key smoothing (33.6B)	5, 11	10,905 / 17,272	124 / 139	74%	30 / 1.1
GProj (33.6B)	none (max 26)	8 / 7	2 / 3	0%	0.12 / 0.12
FP32 attention (33.6B)	none (max 27)	7 / 7	2 / 3	0%	0.13 / 0.13

Table 3: Regime statistics at matched checkpoints, measured with exact FP32 attention on two fixed documents. Errors are relative

L_{2}

errors of

dQ

against an FP64 reference on the same BF16 operands, given as ratios rather than percentages (median of the two documents’ means; values above 1 mean the error exceeds the gradient). With accurate gradients (GProj, FP32) the layers stay small and FA3’s error equals cuDNN’s BF16 floor.

(1) GProj comparison
Method	ms/update	vs. FA3	peak GiB
FA3	227.357	+0.00%	9.191
FP32 attention (eager)	719.420	+216.43%	9.182
GProj	238.120	+4.73%	9.191
(2) optimized FP32 comparison
FA3	231.031	+0.00%	9.191

Table 4: Complete training-step time and peak memory; overhead is relative to the FA3 row of the same panel. Measurement protocol in Appendix E .

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Forward	Backward	Saved output
FA3	BF16 FA3	Native mixed precision	BF16
FA3-SBS	Subtract-before-scale FA3	Native mixed precision	BF16
Key smoothing	FA3 on mean-subtracted keys	Native mixed precision	BF16
QY-shift	Author-source dynamic shift	Author-source BF16 arithmetic	BF16
GProj	Subtract-before-scale FA3	Dense query/key projection	BF16
FP32 attention	FP32 attention	FP32 attention	See below

Appendix

Table 5: Primary method definitions. All methods expose BF16 output and gradient tensors to the model. GProj retains BF16 matrix products with FP32 accumulation and corrections. Eager FP32 recomputes its output; fused FP32 saves an internal FP32 output.

Capture	GProj $dQ$	$dK$	$dV$	FP32 $dQ$	$dK$	$dV$
0	0.3223	0.2942	0.1699	0.4203	0.3290	0.1592
1	0.2504	0.2397	0.1924	0.2245	0.2216	0.1648
2	0.2669	0.3183	0.1734	0.2799	0.3411	0.1603
3	0.2668	0.2855	0.1966	0.4146	0.3760	0.1639
4	1.2604	1.8577	1.0888	1.0951	2.7637	0.1771
5	0.4158	0.5350	0.4715	0.3552	0.3657	0.1648

Appendix

Table 6: Per-capture relative

L_{2}

errors (%) of GProj and fused FP32 attention against FP64 references; Table 7 identifies the captures.

Capture	Checkpoint	Layer	Document
0	5500	11	P
1	5500	5	P
2	5500	11	D
3	5500	5	D
4	8000	11	P
5	8000	5	P

Appendix

Table 7: Identity key for Table 6 . Layers are zero-based; P is a PG19 book and D a ProofPile document.

Tokens (B)	FA3	FA3-SBS	GProj	Key smoothing	FP32 attention
0.8	4.624 / 0.827	4.635 / 0.790	4.623 / 0.918	4.618 / 0.804	4.625 / 0.852
12.6	1.870 / 0.406	1.871 / 0.360	1.869 / 0.312	1.870 / 0.373	1.869 / 0.372
23.1	1.783 / 0.358	1.779 / 0.375	1.773 / 0.403	1.776 / 0.329	1.772 / 0.359
29.4	1.909 / 3.871	1.783 / 0.325	1.773 / 0.385	1.790 / 0.355	1.773 / 0.288
33.6	2.122 / 375.508	1.780 / 0.364	1.769 / 0.396	1.937 / 7.381	1.769 / 0.354
37.7	2.107 / 446.931	1.753 / 0.279	1.740 / 0.324	1.932 / 4.789	1.740 / 0.375

Appendix

Table 8: Matched from-scratch training (seed 42, same data order): all-rank mean training loss and mean pre-clipping gradient norm per 10-step logging interval. Cells show loss / gnorm. All five arms complete 50.0B tokens. The Qiu–Yao arm, stopped at 13.2B tokens, is reported in Appendix G (Figure 4 ). “Final 0.84B” gives mean loss / median gnorm over the 20 records in the final 0.84B tokens; “Norm

>

10 at” gives the tokens at which the gradient norm first exceeds 10.

Setting	Configuration
Parameters	449,465,344; referred to as 450M in the main text
Transformer	24 layers, hidden width 1024
Attention	16 query heads, four key/value heads, head dimension 64
Vocabulary	64,256; untied input and output embeddings
Sequence length	4096 tokens
Attention backend	Native BF16 FlashAttention-3

Appendix

Table 9: Configuration of the 450M model and of every training run.

Ckpt.	Layer	$dQ$ rel. err.	Resid. frac.	Err. cos.	Saved- $O$ red.	rel. err.	cos.
						After saved $O$
5500	5	6.517	0.0055	1.0000	2.98 $\times$	2.165	0.419
5500	11	6.120	0.0060	1.0000	2.45 $\times$	2.496	0.366
8000	5	17.889	0.0067	1.0000	15.73 $\times$	1.055	0.686
8000	11	411.104	0.0197	0.9999	72.02 $\times$	5.988	0.156

Appendix

Table 10: Confirmation on 127 documents per checkpoint/layer, with 128 strided query rows and all 4096 keys. Entries are per-document medians; relative errors are ratios, not percentages. Resid. frac. and Err. cos. compare Equation ( 19 ) with observed native error. Saved-

O

red. is the median paired reduction in absolute

dQ

error after replacing only saved

O

by an FP64-recomputed output rounded to BF16; the last two columns give the remaining relative

dQ

error and the cosine between the computed and exact

dQ

Backward intervention at layers 5 and 11	6500	8000
Native	36.653	5353.856
Replace saved $O$ , first row only	36.402	5353.857
Exact first-row VJP only	36.402	5353.857
Replace saved $O$ , all rows	2.450	19.746
FP32 VJP, all rows	2.258	18.185

Appendix

Table 11: Median full-model parameter-gradient norm over eight documents. All 24 attention outputs and the scalar loss are bitwise identical across modes at each document/checkpoint. “All rows” still changes only two attention layers. Saved-

O

replacement uses FP64 recomputation followed by BF16 rounding, as in Table 10 .

Saved $O$	Saved LSE	$dQ$ error reduction	$dK$ error reduction
FA3-SBS	FA3	$10.822$ – $42.971$	$11.968$ – $58.601$
FA3	FA3-SBS	$0.99676$ – $0.99806$	$0.99636$ – $0.99780$
FA3-SBS	FA3-SBS	$10.815$ – $43.005$	$11.954$ – $58.562$

Appendix

Table 12: Ranges of paired absolute-error reduction on the four checkpoint-8000 full-row captures. The improvement from the source intervention is recovered by its saved output with the FA3 LSE.

Backend	$dQ$ error	$dK$ error	$dQ$ cosine	Late L11 $dQ$ error	One-key violations
FA3	7.729	0.8555	0.131	162.3	6/96
FA3-SBS	2.187	0.1327	0.414	4.12	0/96
FA2, 2.8.3.post1	7.729	0.8555	0.131	162.3	6/96
SDPA cuDNN	2.206	0.1287	0.412	4.16	0/48
SDPA flash	2.217	0.1289	0.410	3.90	0/96
SDPA efficient, KV expanded	2.147	0.1271	0.422	4.04	0/96

Appendix

Table 13: Fixed-operand backend sweep on the from-scratch FA3 run’s captures. Errors are relative

L_{2}

ratios, not percentages; the first three numeric columns are medians over eight full-row captures. “Late L11” is the checkpoint-8000 PG19 layer-11 capture. One-key violations count prescribed cases with

o_{0}\neq v_{0}

(the same cases have nonzero

dq_{0}

); cuDNN cannot run the 48 length-one cases.

Seed	LR	Rel. err. 0	Rel. err. 250	Rel. err. 500	Abs. err. 500	Cosine 500
2026	1	0.00177832	0.00187079	0.00261074	0.00498338	0.999996597
2026	4	0.00177832	0.00340782	0.0243191	0.0514253	0.999706727
2027	1	0.00193911	0.00194124	0.00264846	0.00479022	0.999996495
2027	4	0.00193911	0.00345874	0.0156458	0.0314603	0.999878256
2028	1	0.00186521	0.00200166	0.0026835	0.00528956	0.999996407
2028	4	0.00186521	0.00360284	0.015028	0.0348366	0.999897825

Appendix

Table 14: Full-model gradient analyses of three independent 101,986,816-parameter plain transformer initializations. All attention forward outputs and loss are bitwise matched; only all-attention backward is replaced by FP32 reference arithmetic. Relative error uses the reference gradient norm.

Seed	Peak step	Native norm	Reference norm	Rel. error	Cosine
2026	5054	4913.875	4988.690	0.018294	0.999944

Appendix

Table 15: The LR4 training-batch peak of seed 2026. The analysis replaces every attention backward on the identical native forward and accumulates all eight sequence gradients in FP32 with training’s scored-token weights.

Explore similar work

One QK Channel, Many Sources: Tracing Low-Precision Attention Collapse
Aug 3, 2026Shuxiao Xie, Shuyang Xie, Yuan Cao +3Matched Fp16 Intermediate Transformer Attention
Beijing Tongming Lake Information Technology Application Innovation Center (TLAIC), China · Fudan University Institute of Systems for Advanced Computing, China · Harbin Institute of Technology, China +1
FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention
Sep 21, 2026Anthony Givans, Michael Crawshaw, Mingrui LiuBlock Sparse Flash Attention Gumbel-Softmax Relaxation
FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training
Jun 22, 2026Dikshant Kukreja, Kritarth Prasad, Avinash Anand +6Matched Fp16 Intermediate Optimizer Design
Puch AI · Singapore Institute of Technology · Nanyang Technological University +2

Aug 3, 2026, cs.LGJ/K move · Enter open · S save

One QK Channel, Many Sources: Tracing Low-Precision Attention Collapse

Shuxiao Xie, Shuyang Xie, Yuan Cao, Dezhi Ran+2

Beijing Tongming Lake Information Technology Application Innovation Center (TLAIC), China · Fudan University Institute of Systems for Advanced Computing, China · Harbin Institute of Technology, China+1

Matched Fp16 Intermediate Transformer Attention Channel Dissociation

Open paper AskPDF arXiv

A bfloat16 transformer can train normally, then collapse abruptly. Prior work links collapse to structured attention errors and shows QK normalization disrupts their compounding. Distinct low-precision errors trigger the same collapse, leaving unclear whether each needs a fix at its source or one shared route can be blocked instead. We isolated the fault behind a reproduced GPT-2-class collapse to the streaming-softmax accumulator, where an fp32 streaming core repairs it, and turned it into an assay for moving a controlled error across sources. Using it, we found that errors placed outside attention still drove the same QK spectral runaway, and that correcting only QK kept training stable while the fault stayed active. This is a source-channel dissociation: fault source is not failure channel. It held across tested architectures and scales, and reproduced on a second GPU architecture. As a causal probe, projecting each update off the current QK weights' leading three singular directions held the query projection's largest singular value to 11.1, whereas removing equal energy elsewhere left it at 237: the QK channel causally drives the early runaway. What lets the injected error in is temporal sign-coherence, its per-head sign persisting across steps, not aggregate deviation; once inside, the runaway shows as attention-logit saturation. QK-Guard, a dormant controller, tests this by switching on parameter-free QK normalization at the first monitored threshold crossing. On the runs designated for this test, the QK-local action prevented the failure of each matched or same-configuration unguarded run; on plain GPT-2, all 12 final train and validation losses were within 0.03 nat of same-configuration always-on QK-norm, and both methods ran 60k steps without collapse. Intervention at the QK locus therefore suffices in place of a fix at each source.

Abstract

Explore similar work

One QK Channel, Many Sources: Tracing Low-Precision Attention Collapse

FlashBoB: I/O-Efficient Exact Backward-over-Backward for Softmax Attention

FORGE: Fused On-Register Gradient Elimination for Memory-Efficient LLM Training