当前位置: 首页 » 资讯 » 科技头条 » 正文

为什么BF16的FlashAttention会把训练「炸掉」?清华给出机制解释

IP属地 中国·北京 机器之心Pro 时间:2026-03-04 12:25:36



一句话总结:社区里困扰了多年的一个 “玄学” 现象终于被拆解清楚了:在 BF16 等低精度训练里,FlashAttention 不是随机出 bug,而是会在特定条件下触发有方向的数值偏置,借助注意力中涌现的相似低秩更新方向被持续放大,最终把权重谱范数和激活推到失控,导致 loss 突然爆炸。论文还给出一个几乎不改模型、只在 safe softmax 里做的极小修改,实测能显著稳定训练。



因果链总览(论文 Figure 1)



标题:Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention作者:邱海权,姚权铭机构:清华大学 电子工程系投稿:ICLR 2026 Oral关键词:低精度训练,BF16,FlashAttention,数值稳定性,舍入误差(rounding error),低秩表示(low-rank)论文链接:https://arxiv.org/abs/2510.04212代码链接:https://github.com/ucker/why-low-precision-training-fails

背景:低精度训练越来越 “刚需”,但注意力比你想的更敏感

大模型训练的现实是:显存和吞吐决定一切。工业界普遍在混合精度里使用 BF16/FP16,甚至把 FFN 推到 FP8,以换取更高的训练效率。但工程实践同样残酷:越接近 “极限精度”,训练越容易出现难以解释的不稳定。

Flash Attention 是长上下文训练的关键加速组件,几乎成了标配。问题在于,社区长期存在一个可复现却难以解释的失败案例:

用 FlashAttention + BF16 训练 GPT-2,一开始正常收敛,但在几千 step 之后突然 loss 爆炸。你可以通过回退到标准注意力、或把关键计算提高到 FP32 来 “救火”,但代价是吞吐和显存优势没了。

这类问题被报告了多年(相关 issue 在多个开源项目里反复出现),却一直缺少一条能 “从数值误差一路解释到 loss 爆炸” 的机制链。



作者的做法很工程,且足够 “可复现”:





机制解释 1:相似低秩结构,让误差变成 “持续推力” 而不是噪声







结果就是:权重更新被 “带偏”,谱范数和激活异常增长,最终把训练推到 loss 爆炸。





低秩结构相似性与偏置累积(论文 Figure 4/5)

机制解释 2:偏置从哪来?safe softmax + BF16 舍入误差里藏着一个 “离散触发器”



作者把问题追到了 FlashAttention 前向里的未归一化输出:



















检测一行 S 中最大值是否出现多次一旦出现 “重复最大值”,就动态调整safe softmax 的行移位常数 m,让最大位置的指数也变成严格小于 1

论文给出的实现(概念上)如下:





实验结果:稳定训练不再 “突然炸”

论文在 BF16 设置下验证了上述分析与修复:

GPT-2S:使用修改后的 FlashAttention,在 AdamW 与 Muon 两种优化器下,都能稳定训练到 600K stepsGPT-2M:同样能在 AdamW 下稳定训练(论文展示到 100K steps)论文还提到该现象与结论在多种硬件上保持一致(包括 A100、RTX 4090、Ascend 910B)



验证集 loss 曲线对比(论文 Figure 7)

更重要的启示:别把低精度误差当成 “零均值噪声”

这篇论文的价值不只在 “修了一个 bug”,更在于给出了一个可迁移的诊断范式:

数值误差未必是随机噪声。在特定分布与离散事件(如重复最大值、概率精确为 1)下,舍入误差可能形成系统性偏置。模型结构会放大偏置。注意力里涌现的相似低秩更新方向,让偏置误差更容易 “同向叠加”。经验修复为什么有效也能被解释:论文讨论了 attention sinks 与多最大值的关系,并给出了一个数值层面的连接;同时也指出一些稳定化技巧(如 QK normalization、Gated Attention)可能通过 “打散结构相似性” 来阻止误差同向累积。

作者介绍

邱海权是清华大学在读博士研究生,研究方向涵盖机器学习理论、表示学习与大模型机制分析。他的研究围绕模型表达能力、结构归纳偏置以及参数空间几何与优化动力学之间的内在联系展开,关注模型在不同结构约束与训练条件下的泛化行为与可组合性问题。整体上,他强调以可分析的理论框架刻画模型的能力边界与机制来源,从结构与原理层面理解深度模型为何有效、何时失效。

姚权铭,清华大学电子工程系副教授。长期致力于数据高效学习与智能体系统研究,在少样本学习、图学习、知识图谱与生物医药智能等方向取得系统性成果。发表 Nature 子刊、TPAMI、JMLR、ICML、NeurIPS、ICLR 等论文 130 余篇,被引 1.4 万余次。代表性工作包括抗噪学习算法 Co-teaching、小样本学习综述、自动化图学习方法及新药物相互作用预测模型。现任 TPAMI、TMLR 编委及 Neural Networks 资深编委,多次担任 ICML、NeurIPS、ICLR 领域主席,入选 IEEE Computing Top 30、IET Fellow 等。

免责声明:本网信息来自于互联网,目的在于传递更多信息,并不代表本网赞同其观点。其内容真实性、完整性不作任何保证或承诺。如若本网有任何内容侵犯您的权益,请及时联系我们,本站将会在24小时内处理完毕。