Training Stability & Mixed Precision: 哪些算子/操作在低精度下最容易出问题？

Core Concept: 含 exp/log/除法的算子（softmax、LN、CE）、累加归约（attention、matmul）、以及需精确表示小增量的权重更新。

数学机理：低精度风险来自'<strong>范围</strong>'与'<strong>精度</strong>'两类。<strong>(1) 范围风险（上溢/下溢）</strong>——含 exp 的算子：softmax、GELU、sigmoid；FP16 下 exp(x) 在 x>11 时即超 65504 上溢。含 log 的算子：CE、KL、log-softmax；log(0) 得 −∞。含除法的算子：LN/BN（除方差）、attention 归一化；分母为 0 或极小则爆炸。<strong>(2) 精度风险（累加误差）</strong>——<strong>归约运算</strong>（sum/mean）在 FP16 下逐元素累加会累积舍入误差：累加 n 个数的误差约 O(n·u)（u 为机器精度）；attention 的 log-sum-exp 要对序列长度 L 个元素求和、LN 要对 d 个元素求和，L/d 大时误差显著。故实践中这些归约用 FP32 累加（'accumulate in FP32'）。<strong>(3) 表示风险</strong>——权重更新 θ←θ+ηΔθ：若 ηΔθ 相对于 θ 极小（如 η=1e-5、θ≈0.1、Δθ≈1），则 ηΔθ≈1e-6，而 FP16 在 0.1 附近的分辨率约 6e-5——<strong>更新量被舍入掉、参数根本不动</strong>；故必须用 FP32 master weights 做更新。这三类风险对应三条对策：<strong>减 max / 加 eps / FP32 累加 + FP32 master</strong>。

深度剖析与工程权衡：① <strong>autocast 的白名单</strong>——PyTorch 的 autocast 会自动把 matmul/conv 等'安全'算子转 FP16（吞吐收益大），而把 softmax、LN、loss、归约等'敏感'算子保持 FP32（autocast 内置白/黑名单）；理解这一点比手工管理 dtype 更可靠。② <strong>Flash Attention 的贡献</strong>——它在 kernel 内部用 FP32 做在线 softmax（维护 running max 与 running sum），既保证数值稳定又避免物化 L×L 的注意力矩阵；这是'数值稳定 + 高效'结合的典范。③ <strong>LayerNorm 的 ε</strong>——低精度下方差可能因舍入略小于真实值，ε 太小则除法放大误差；故 LN 的 ε 常取 1e-5~1e-6，且计算在 FP32 下进行。④ <strong>matmul 的累加</strong>——张量核心的 FP16 matmul 在内部用 FP32 累加（这是硬件特性），故 matmul 本身相对安全；风险主要在非 matmul 的逐元素归约。⑤ <strong>量化推理的对应问题</strong>——INT8 推理时，激活的离群值导致 per-tensor 量化误差大，故用 per-channel/per-token 量化（同'范围'问题的另一种解法）。⑥ <strong>面试要点</strong>——被问'低精度训练要注意什么'，应能按'<strong>exp/log/除法（范围）+ 归约累加（精度）+ 权重更新（表示）</strong>'三类回答，并提到 autocast 白名单与 FP32 master weights 两个具体机制；这体现'从原理到工程'的完整链条。