Mixed Precision Training

Mixed Precision Training 指在训练中混合使用不同数值精度，以降低显存、提高吞吐，同时尽量保持收敛稳定性。大模型训练中常见做法是：模型权重、激活和梯度使用 fp16 或 bf16，部分累加、归一化、optimizer state 或 master weights 保持 fp32。

混合精度的核心不是“把所有东西都变小”，而是在不同张量上选择合适 dtype：低精度负责节省带宽和显存，高精度负责保护数值稳定性。

常见精度

bf16 常比 fp16 更适合大模型训练，因为它保留了 fp32 的 exponent range，数值溢出风险更低，通常不需要像 fp16 那样依赖 loss scaling。

选择 dtype 时应区分三个问题：

• 存储精度：tensor 在显存中长期以什么格式保存；
• 计算精度：矩阵乘、attention、normalization 等 op 用什么格式计算；
• 累加精度：reduction、dot product、optimizer update 的中间累加用什么格式。

很多硬件上的低精度矩阵乘会用 fp16 / bf16 输入，但在更高精度 accumulator 中累加，再输出低精度结果。因此“bf16 训练”并不意味着每一步计算都只有 bf16 精度。

Loss Scaling

fp16 的数值范围较窄，小梯度可能 underflow 成 0。Loss scaling 的做法是先把 loss 放大：

反向传播得到放大的梯度后，再在 optimizer update 前除以 $s$：

如果检测到 overflow，则降低 scale；如果一段时间稳定，则可以增大 scale。这就是 dynamic loss scaling。bf16 通常不需要 loss scaling，但仍可能需要关注归一化、softmax、attention score 等局部数值稳定性。

Master Weights 与 Optimizer State

混合精度训练常保留 fp32 master weights：

• forward/backward 使用 fp16/bf16 weights；
• optimizer 在 fp32 master weights 上更新；
• 更新后再同步或 cast 到低精度训练权重。

这会增加显存。以 bf16 + AdamW + fp32 master weights 为例：

因此混合精度节省了参数、梯度和 activation 的部分显存，但 optimizer state 可能仍是 fp32 大头。详见 Optimizer State。

AMP

AMP，Automatic Mixed Precision，是框架自动选择部分 op 使用低精度、部分 op 保持高精度的机制。它通常包括：

• autocast：根据 op 类型自动选择 dtype；
• grad scaler：fp16 下动态调整 loss scale；
• fp32 fallback：对数值敏感 op 保持高精度。

AMP 降低了手写 mixed precision 的复杂度，但不代表可以忽略数值问题。训练中仍需监控 loss spike、NaN/Inf、gradient norm 和收敛曲线。

常见需要保留或回退到高精度的部分包括：

• loss reduction 与 gradient norm 统计；
• softmax 前后的稳定化计算；
• LayerNorm / RMSNorm 的部分统计；
• optimizer update；
• logits 与 cross entropy 的某些实现路径；
• 大规模 all-reduce 后的梯度检查。

具体哪些 op 使用 fp32 取决于框架、硬件和 kernel。实践中应优先使用成熟 AMP / distributed training recipe，而不是手动对所有张量强制 cast。

Switch Transformer 的 Router Precision

Switch Transformer 给出了 selective precision 的一个早期实证：纯 BF16 的 sparse router 可能导致训练 divergence，而把 router input、logits 和 softmax 局部提升到 FP32，可以在保留 BF16 expert computation 与通信的同时恢复稳定性。该策略的关键不在于“所有 router tensor 都长期用 FP32 存储”，而在于把 FP32 限制在产生离散 expert assignment 的局部计算中，离开 router 后再将 dispatch / combine tensors 转回 BF16。

这类做法适合需要先做离散选择、再进行大规模通信或稀疏计算的结构。它也提醒我们，mixed precision 的优先级不应只按算力占比决定：router 的计算量很小，却会改变 token 的整个 expert path，因此其数值误差可能比普通 GEMM 更敏感。最终 dtype policy 仍需结合硬件、kernel、loss curve、NaN/Inf 和 assignment stability 验证。

与 FSDP / ZeRO 的配合

FSDP 和 ZeRO 会切分 parameters、gradients 和 optimizer states。Mixed precision 决定这些状态的 dtype，FSDP/ZeRO 决定这些状态是否在 data parallel 维度上复制。

二者解决的是不同问题：

• mixed precision：减少每个 tensor element 的 bytes；
• ZeRO/FSDP：减少每张 GPU 持有的 tensor elements 数量。

例如 bf16 full fine-tuning 可以把 parameter/gradient 从 fp32 的 4 bytes 降到 2 bytes；ZeRO-3 可以进一步把 parameter/gradient/optimizer state 分片到多个 GPU。

分布式训练中还要区分 parameter dtype、reduce dtype 和 buffer dtype。有些配置会用 bf16 参数和 bf16 梯度通信，有些会在 fp32 中做 gradient reduction 以保护稳定性。通信 dtype 越低，带宽压力越小；但如果梯度数值范围或累加误差变大，可能引入训练不稳定。

FP8 的边界

FP8 进一步降低显存和带宽，但训练 recipe 更复杂，通常需要：

• 支持 FP8 tensor core 的硬件；
• per-tensor 或 per-channel scaling；
• amax history 或动态缩放；
• 对 attention、MLP、normalization、optimizer 等模块分别设定精度策略；
• 更严格的 loss spike 和 NaN 监控。

因此 FP8 更适合成熟训练栈中的性能优化，而不是初次训练或小规模实验的默认选择。对知识库中的通用估算，bf16/fp16 仍是更稳定的基准。

DeepSeek-V3 的 FP8 训练策略

DeepSeek-V3 的 FP8 经验说明，低精度训练需要按算子和 tensor 类型设计 mixed-precision policy，而不是把所有计算统一 cast 到 FP8：

• Linear 的 Fprop、Dgrad 和 Wgrad GEMM 使用 FP8；
• embedding、output head、MoE gating、normalization 和 attention 保持 BF16 或 FP32；
• master weights、weight gradients 和关键 optimizer state 保留更高精度；
• MoE dispatch 前的部分 activation 使用 FP8，combine 保持 BF16。

为处理 outlier，DeepSeek-V3 使用 fine-grained quantization：

• activation 按每个 token 的 1x128 tile 做 scale；
• weight 按 128x128 block 做 scale；
• scale 在当前 tile / block 上在线计算，而不是只依赖历史 maximum。

H800 Tensor Core 的 FP8 accumulation 有效精度有限，因此 DeepSeek-V3 每累计 128 个 FP8 乘积，就把 partial result promotion 到 CUDA Core 的 FP32 registers 中继续累加。这个高精度 accumulation 与 group-wise scale 配合，才能把 FP8 与 BF16 的 relative loss error 控制在报告的 0.25% 以内。

报告同时展示了一个重要失败模式：activation gradient 如果简单使用与 weight 相同的 128x128 block-wise quantization，MoE 模型可能在数百 billion tokens 后 divergence。activation gradient 的 token-correlated outliers 需要更适合反向传播方向的分组方式。因此 FP8 的验证必须包含长期 loss curve、NaN/Inf、gradient norm 和 divergence monitoring，不能只做短时间吞吐测试。

失败模式与边界

• Overflow / Underflow：fp16 下尤其常见，表现为 loss NaN、gradient Inf 或训练停滞。
• Loss Spike：低精度计算可能放大不稳定训练配置的问题。
• 不稳定 op：softmax、normalization、large reduction 等可能需要 fp32 或特殊 kernel。
• dtype mismatch：手写模块或第三方 kernel dtype 不一致，可能造成隐式 cast 和性能损失。
• 误判显存节省：即便参数是 bf16，optimizer states 仍可能是 fp32。

诊断信号

混合精度问题通常可以通过以下信号定位：

• loss 突然变为 NaN / Inf；
• gradient norm 突然爆炸或长期为 0；
• dynamic loss scale 频繁下降；
• bf16 与 fp32 小规模对照收敛曲线明显不同；
• 某个 kernel 或自定义 op 触发大量隐式 cast；
• 只在多卡通信后出现数值异常。

排查时应先用小 batch、短序列和较保守 dtype 复现，再逐步开启 fused kernel、FlashAttention、ZeRO/FSDP、FP8 或其他高性能配置。

相关概念

• Training Memory Estimation
• Optimizer State
• Checkpoint Sharding
• Loss Spike
• FSDP
• ZeRO
• Numerical Stability