大模型混合精度训练：原理、选型与实战笔记

混合精度训练，简单说就是用更少位数实现更快的计算速度和更低的显存占用。混合精度训练的核心是用 FP16 和 BF16 这两种半精度浮点数格式替代传统的FP32单精度格式。

1. 为什么需要混合精度训练？

• 痛点：GPT-3级别模型训练耗时数月、花费数百万美元，FP32全精度训练在显存和算力上均不可持续。
• 收益 ：在保证模型精度的前提下，训练速度提升 1~2倍 ，显存占用减少 50%。
• 精度无损：通过"高精度存储+低精度计算"的策略，最终模型收敛效果与FP32基本一致。
• 现状：几乎所有大模型训练都在使用该技术。

2. 核心概念：FP16 vs BF16 vs TF32

所有浮点数均由 符号位(S) + 指数位(E) + 尾数位(M) 组成。**指数位决定数值范围（防溢出），尾数位决定有效精度（防舍入误差）。**这三种格式都遵循 IEEE 754 标准，但位分配策略不同：

格式	位宽分配 (S/E/M)	动态范围	精度	核心特性	推荐场景
FP16	1 / 5 / 10	±65,504	高	精度高，范围窄，极易溢出/下溢，必须配合Loss Scaling	老旧GPU、推理部署、科学计算
BF16	1 / 8 / 7	±3.4×10³⁸ (同FP32)	中	范围大，截断式舍入，精度稍低。无需Loss Scaling，训练极稳定	现代大模型预训练(LLaMA/GPT)、Google TPU
TF32	1 / 8 / 10	±3.4×10³⁸	中高	NVIDIA Ampere+专属，硬件自动转换，兼顾范围与精度	NVIDIA A100/H100/B200 训练默认首选核心口诀 ：符号位定正负，指数位定大小（范围），尾数位定精度。

💡 形象理解

• FP16 = 小户型公寓：房间装修精细（尾数多），但总面积小（指数少），东西稍多就放不下（溢出）。
• BF16 = 大平层毛坯：面积巨大（指数同FP32），但装修粗糙（尾数少），能装下所有家具但细节不够精致。
• TF32 = 精装改善房：既有大面积又有不错装修，但仅限NVIDIA"小区"业主专享。

选型原则 ：NVIDIA新卡用TF32；TPU用BF16；老卡用FP16。大模型训练首选BF16，因为"宁可牺牲精度，也要保证训练稳定"。

3. 混合精度训练的"五步法"核心机制

本质是 "低精度前向/反向 + 高精度权重更新" 的流水线：

1. 📉 权重降精度：维护一份FP32主权重副本，每轮迭代开始时转换为FP16/BF16用于计算（显存占用减半）。
2. ⚡ 低精度计算：前向传播与反向传播均使用半精度，激活Tensor Core矩阵乘加速。
3. 📈 梯度升精度 ：关键步骤！ 将半精度梯度立即转回FP32，防止反向传播链式求导中的累积舍入误差导致发散。
4. 🔧 FP32权重更新 ：weight_fp32 -= lr * grad_fp32。因为 lr × grad 通常极小（如1e-7），FP16/BF16无法表示该量级变化，会被直接置零导致训练停滞。
5. 🔄 循环迭代：重复以上过程直至收敛。

4. 面试高频考点与工程陷阱

Q1: 为什么大模型训练首选BF16而非FP16？

• 根本原因：大模型训练过程中梯度分布跨度极大，FP16的动态范围（±65504）不足以覆盖，频繁触发溢出(NaN)和下溢(0)。
• 工程代价 ：FP16必须引入 Loss Scaling（损失缩放）机制------前向时放大Loss，反向后缩小梯度，还需动态调整scale factor，增加了实现复杂度和调试难度。
• BF16优势 ：指数位与FP32相同，天然覆盖训练所需全部数值范围，开箱即用、无需调参，大幅简化分布式训练工程栈。

Q2: Loss Scaling 的原理是什么？

• 问题：FP16最小正数为 6×10−56×10−5 ，小于此值的梯度被置零（下溢）。
• 解法：将Loss乘以缩放因子S（如1024），使梯度同步放大S倍进入FP16可表示区间；更新权重前再将梯度除以S恢复真实值。
• 注意 ：BF16因范围充足，通常不需要Loss Scaling。

Q3: TF32 和 BF16 如何选择？

• NVIDIA Ampere及以上：优先开启TF32（PyTorch默认开启），它在不修改代码的情况下自动获得比BF16更高的吞吐，且精度优于BF16。
• 跨平台/TPU：使用BF16，保证代码在不同硬件间行为一致。
• 极致精度敏感任务：仍可回退到FP16+Loss Scaling或FP32。

5. 实战最佳实践清单

建议	说明
✅ 始终保留FP32主权重	永远不要只用半精度存储权重，否则累积误差必然导致发散
✅ 梯度裁剪(GRAD CLIP)	混合精度下更应配合梯度裁剪，防止偶发大梯度引发溢出
✅ 使用框架原生AMP	PyTorch torch.cuda.amp / DeepSpeed / Megatron-LM 已封装完整流程，避免手写
⚠️ 谨慎对待归约操作	Softmax、LayerNorm、Loss计算等建议保持FP32，这些操作对精度敏感
⚠️ 验证数值对齐	切换精度后，应对比FP32基线的loss曲线，确保无异常偏移

一句话总结

混合精度训练的本质是用"可控的精度损失"换取"成倍的算力效率"。在现代大模型时代，BF16/TF32凭借充足的动态范围取代FP16成为训练标准，而"低精度计算 + FP32更新"的五步法是理解一切AMP实现的基石。