Training Stability & Mixed Precision: 解释为何混合精度下 LayerNorm / softmax 要保持 FP32。

Core Concept: 两者都含'归约 + 除法/指数'：FP16 累加误差大、exp 易溢出；保 FP32 计算可保证稳定性，开销可忽略。

数学机理：<strong>LayerNorm</strong> 计算 μ=(1/d)Σx_i、σ²=(1/d)Σ(x_i−μ)²，都是<strong>对 d 个元素做归约</strong>。低精度下的两个问题：(a) <strong>累加误差</strong>——FP16 逐元素累加 d 个数，相对误差约 O(d·u)（u≈1e-3），d=4096 时可达 4 的量级，μ 与 σ² 严重失真；(b) <strong>抵消误差（catastrophic cancellation）</strong>——计算 x_i−μ 时，若 x_i≈μ，两个相近数相减会损失有效位（FP16 尾数仅 10 位），使 (x_i−μ) 的<strong>相对</strong>误差被放大。<strong>softmax</strong> 的问题：(a) exp(z) 在 FP16 下易上溢（需减 max，但减 max 前的中间值仍需安全）；(b) 分母的求和是归约，同 LN 的累加误差。<strong>为什么保 FP32 几乎无代价</strong>：LN 与 softmax 的计算量相对 matmul 极小（matmul 是 O(d²) 或 O(L·d²)，而归约是 O(d) 或 O(L²) 但非矩阵乘）；张量核心的算力集中在 matmul，故把这些'小算子'保 FP32 对总吞吐影响 <5%，却换来数值稳定性。这正是 PyTorch autocast 把它们列入 FP32 白名单的原因。

深度剖析与工程权衡：① <strong>autocast 的白名单机制</strong>——PyTorch 自动把 matmul/conv/linear 转 FP16（收益大），而 softmax、LN、loss、归约、指数类算子保持 FP32（收益小、风险大）；用户无需手工管理，但应知道哪些算子在 FP32。② <strong>与 Flash Attention 的关系</strong>——Flash Attention 在 kernel 内用 FP32 累加与在线 softmax，等价于'把 LN/softmax 的 FP32 要求内化到 kernel 中'，因此它既省显存又保精度。③ <strong>RMSNorm 的敏感性</strong>——RMSNorm 只算均方根（不含减均值），少了一次抵消误差，故对低精度更宽容；这也是它在大模型中流行的原因之一（除省算力外）。④ <strong>LN 的 ε 与低精度</strong>——ε 太小则除法放大噪声，太大则归一化失真；低精度下常取 1e-5~1e-6。⑤ <strong>诊断方法</strong>——若怀疑低精度导致的不稳，可做对照实验：把 LN/softmax 强制 FP32 后再训练，若稳定则确认为精度问题。⑥ <strong>面试要点</strong>——被问'混合精度下哪些算子要保 FP32'，应给出'<strong>含归约、指数、除法的算子</strong>'这一判据（而非死记名字），并解释'归约累加误差 + 抵消误差 + exp 溢出'三个机制；同时指出'开销可忽略'是可行性前提。