為什么BF16的FlashAttention會把訓(xùn)練「炸掉」？清華給出機制解釋

一句話總結(jié)：社區(qū)里困擾了多年的一個 “玄學(xué)” 現(xiàn)象終于被拆解清楚了：在 BF16 等低精度訓(xùn)練里，F(xiàn)lashAttention 不是隨機出 bug，而是會在特定條件下觸發(fā)有方向的數(shù)值偏置，借助注意力中涌現(xiàn)的相似低秩更新方向被持續(xù)放大，最終把權(quán)重譜范數(shù)和激活推到失控，導(dǎo)致 loss 突然爆炸。論文還給出一個幾乎不改模型、只在 safe softmax 里做的極小修改，實測能顯著穩(wěn)定訓(xùn)練。

因果鏈總覽（論文 Figure 1）

• 標題：Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention
• 作者：邱海權(quán)，姚權(quán)銘
• 機構(gòu)：清華大學(xué) 電子工程系
• 投稿：ICLR 2026 Oral
• 關(guān)鍵詞：低精度訓(xùn)練，BF16，F(xiàn)lashAttention，數(shù)值穩(wěn)定性，舍入誤差（rounding error），低秩表示（low-rank）
• 論文鏈接：https://arxiv.org/abs/2510.04212
• 代碼鏈接：https://github.com/ucker/why-low-precision-training-fails

背景：低精度訓(xùn)練越來越 “剛需”，但注意力比你想的更敏感

大模型訓(xùn)練的現(xiàn)實是：顯存和吞吐決定一切。工業(yè)界普遍在混合精度里使用 BF16/FP16，甚至把 FFN 推到 FP8，以換取更高的訓(xùn)練效率。但工程實踐同樣殘酷：越接近 “極限精度”，訓(xùn)練越容易出現(xiàn)難以解釋的不穩(wěn)定。

Flash Attention 是長上下文訓(xùn)練的關(guān)鍵加速組件，幾乎成了標配。問題在于，社區(qū)長期存在一個可復(fù)現(xiàn)卻難以解釋的失敗案例：

• 用 FlashAttention + BF16 訓(xùn)練 GPT-2，一開始正常收斂，但在幾千 step 之后突然 loss 爆炸。
• 你可以通過回退到標準注意力、或把關(guān)鍵計算提高到 FP32 來 “救火”，但代價是吞吐和顯存優(yōu)勢沒了。

這類問題被報告了多年（相關(guān) issue 在多個開源項目里反復(fù)出現(xiàn)），卻一直缺少一條能 “從數(shù)值誤差一路解釋到 loss 爆炸” 的機制鏈。

作者的做法很工程，且足夠 “可復(fù)現(xiàn)”：

機制解釋 1：相似低秩結(jié)構(gòu)，讓誤差變成 “持續(xù)推力” 而不是噪聲

結(jié)果就是：權(quán)重更新被 “帶偏”，譜范數(shù)和激活異常增長，最終把訓(xùn)練推到 loss 爆炸。

低秩結(jié)構(gòu)相似性與偏置累積（論文 Figure 4/5）

機制解釋 2：偏置從哪來？safe softmax + BF16 舍入誤差里藏著一個 “離散觸發(fā)器”

作者把問題追到了 FlashAttention 前向里的未歸一化輸出：

• 檢測一行 S 中最大值是否出現(xiàn)多次
• 一旦出現(xiàn) “重復(fù)最大值”，就動態(tài)調(diào)整safe softmax 的行移位常數(shù) m，讓最大位置的指數(shù)也變成嚴格小于 1

論文給出的實現(xiàn)（概念上）如下：

實驗結(jié)果：穩(wěn)定訓(xùn)練不再 “突然炸”

論文在 BF16 設(shè)置下驗證了上述分析與修復(fù)：

• GPT-2S：使用修改后的 FlashAttention，在 AdamW 與 Muon 兩種優(yōu)化器下，都能穩(wěn)定訓(xùn)練到 600K steps
• GPT-2M：同樣能在 AdamW 下穩(wěn)定訓(xùn)練（論文展示到 100K steps）
• 論文還提到該現(xiàn)象與結(jié)論在多種硬件上保持一致（包括 A100、RTX 4090、Ascend 910B）

驗證集 loss 曲線對比（論文 Figure 7）

更重要的啟示：別把低精度誤差當成 “零均值噪聲”

這篇論文的價值不只在 “修了一個 bug”，更在于給出了一個可遷移的診斷范式：

• 數(shù)值誤差未必是隨機噪聲。在特定分布與離散事件（如重復(fù)最大值、概率精確為 1）下，舍入誤差可能形成系統(tǒng)性偏置。
• 模型結(jié)構(gòu)會放大偏置。注意力里涌現(xiàn)的相似低秩更新方向，讓偏置誤差更容易 “同向疊加”。
• 經(jīng)驗修復(fù)為什么有效也能被解釋：論文討論了 attention sinks 與多最大值的關(guān)系，并給出了一個數(shù)值層面的連接；同時也指出一些穩(wěn)定化技巧（如 QK normalization、Gated Attention）可能通過 “打散結(jié)構(gòu)相似性” 來阻止誤差同向累積。

作者介紹

邱海權(quán)是清華大學(xué)在讀博士研究生，研究方向涵蓋機器學(xué)習(xí)理論、表示學(xué)習(xí)與大模型機制分析。他的研究圍繞模型表達能力、結(jié)構(gòu)歸納偏置以及參數(shù)空間幾何與優(yōu)化動力學(xué)之間的內(nèi)在聯(lián)系展開，關(guān)注模型在不同結(jié)構(gòu)約束與訓(xùn)練條件下的泛化行為與可組合性問題。整體上，他強調(diào)以可分析的理論框架刻畫模型的能力邊界與機制來源，從結(jié)構(gòu)與原理層面理解深度模型為何有效、何時失效。

姚權(quán)銘，清華大學(xué)電子工程系副教授。長期致力于數(shù)據(jù)高效學(xué)習(xí)與智能體系統(tǒng)研究，在少樣本學(xué)習(xí)、圖學(xué)習(xí)、知識圖譜與生物醫(yī)藥智能等方向取得系統(tǒng)性成果。發(fā)表 Nature 子刊、TPAMI、JMLR、ICML、NeurIPS、ICLR 等論文 130 余篇，被引 1.4 萬余次。代表性工作包括抗噪學(xué)習(xí)算法 Co-teaching、小樣本學(xué)習(xí)綜述、自動化圖學(xué)習(xí)方法及新藥物相互作用預(yù)測模型?，F(xiàn)任 TPAMI、TMLR 編委及 Neural Networks 資深編委，多次擔任 ICML、NeurIPS、ICLR 領(lǐng)域主席，入選 IEEE Computing Top 30、IET Fellow 等。