混合精度训练与断点续训：稳定与可恢复
摘要：混合精度决定训练能跑多快，断点续训决定训练敢不敢跑。前者是吞吐问题，后者是可靠性问题——千卡级训练一次故障恢复成本以小时计，这两件事做不好，集群再大也白搭。
1. 混合精度：为什么不能全用低精度
1.1 精度选择矩阵
| 精度 | 指数位 | 尾数位 | 训练用途 | 关键限制 |
|---|---|---|---|---|
| FP32 | 8 | 23 | 基准/求稳 | 显存翻倍、速度慢 |
| TF32 | 8 | 10 | 折中（Ampere+） | 精度略降，多数场景可用 |
| FP16 | 5 | 10 | 训练（需 Loss Scaling） | 动态范围小，溢出风险 |
| BF16 | 8 | 7 | 训练主力 | 尾数少，大数值下相对精度够用 |
| FP8 | 5/4 | 2/3 | 训练/推理加速 | 需额外稳定性技巧 |
核心原则：权重与梯度用高精度（FP32）保存，计算用低精度（BF16/FP16）提速，二者结合即"混合精度"。
1.2 两种主流方案的差异
FP16 方案（AMP，需 Loss Scaling）：
loss × scale → FP16 前向 → FP16 梯度 → 反缩放 → FP32 更新
问题：梯度下溢/上溢风险，scale 动态调整，实现复杂
BF16 方案（现代默认）：
BF16 前向/反向 → FP32 master weight 更新
问题：无需 scaling，动态范围与 FP32 相同，天然稳定- FP16 的动态范围小（指数 5 位），小梯度直接下溢为 0——所以需要 Loss Scaling 把梯度"放大"再更新。
- BF16 指数位与 FP32 相同（8 位），动态范围一致，不会下溢/上溢——省掉 Loss Scaling，训练稳定。代价是尾数少，但大模型训练已被验证足够。
- 工程结论：新训练项目默认 BF16；FP16 仅用于老框架兼容或特定量化场景；FP8 处于"性能香、稳定难"阶段，谨慎上生产。
2. 检查点（Checkpoint）：训练的心脏
2.1 检查点里到底该存什么
python
checkpoint = {
"model": model.state_dict(), # 权重（FP32 master + BF16 副本）
"optimizer": optimizer.state_dict(), # Adam 动量/方差（占大头！）
"lr_scheduler": scheduler.state_dict(),
"rng": torch.random.get_rng_state(), # 数据顺序、Dropout 随机性
"data_loader": dataloader.state_dict(), # epoch/sample 位置
"epoch": epoch, "step": step,
"config": config, # 超参，恢复时校验
}- 70B BF16 全参训练，检查点 ≈ 权重 140G + 优化器 420G ≈ 560 GB，千卡集群每 N 步保存一次——保存速度直接决定故障恢复时间（RTO）。
- 优化器状态是检查点里最大的一块，远超权重本身。
2.2 保存策略与异步保存
python
# 异步保存：把序列化+落盘放到后台线程，避免阻塞训练
import threading
def save_async(state, path):
def _save():
torch.save(state, path + ".tmp")
os.replace(path + ".tmp", path) # 原子替换，防半写
t = threading.Thread(target=_save, daemon=True)
t.start()
# 频率策略：小间隔+滚动保留（如每 500 步保留最近 3 个）
# 关键里程碑单独永久保存（如 pretrain 完成、SFT 各阶段）三个生产要点：
- 原子写：先写
.tmp再os.replace，避免训练中途断电产生半写检查点。 - 滚动保留：磁盘上保留最近 2～3 个 + 里程碑永久副本，防止"坏检查点覆盖好检查点"。
- 校验和：保存时记录权重 SHA256，恢复时校验——坏检查点恢复比不恢复更可怕（模型静默变差）。
3. 断点续训：恢复的不只是权重
3.1 标准恢复流程
1. 加载 config，校验超参与保存时一致（lr、batch、并行度）
2. 加载权重 + 优化器状态 + scheduler
3. 恢复 RNG 状态（数据采样顺序、Dropout）
4. 恢复 dataloader 位置（epoch/sample）
5. 恢复 step/epoch 计数（影响 LR 曲线、日志对齐）
6. 恢复通信集合（NCCL world）——重跑一个 warmup step 验证3.2 一致性验证：恢复后如何确认"续上了"
| 验证项 | 方法 | 失败含义 |
|---|---|---|
| 权重校验 | SHA256 对比保存记录 | 文件损坏 |
| 优化器状态 | 恢复后 loss 应立即贴近保存点附近水平 | 优化器状态丢失（恢复等于重训） |
| 数据顺序 | 恢复后 loss 曲线与"未中断"曲线吻合 | RNG/dataloader 未恢复 |
| 通信健康 | warmup step 无 NCCL 报错、无重传异常 | 环境变了（驱动/拓扑） |
| 吞吐对比 | 恢复后 TPS 与中断前一致（±5%） | 节点性能降级（散热/降频） |
铁律：恢复后必须跑通 warmup + 观察 3～5 步 loss 再放行全量训练，否则异常恢复会在 100 步后悄悄毁掉整个训练。
3.3 常见坑
- LR 调度器没恢复：恢复后 LR 从头开始 → 训练"回炉"，白跑。
- RNG 状态没恢复：数据顺序变了 → 相当于换了训练集顺序，loss 曲线跳变。
- 跨环境恢复：换了 GPU 型号/驱动版本恢复检查点 → 数值差异累积，先验证再全量。
- 分布式恢复失败：
torch.distributed下检查点只在 rank 0 保存，恢复时广播，网络抖动会静默失败——每个 rank 落盘 + 对比校验最稳。
4. 生产清单
- [ ] 精度方案：BF16 默认，FP16 需 Loss Scaling 已配置，FP8 经小规模验证
- [ ] Master weight 保持 FP32，更新路径正确
- [ ] 检查点含 model/optimizer/scheduler/RNG/dataloader/config 全量
- [ ] 原子写（.tmp + os.replace）+ 滚动保留 + 里程碑副本
- [ ] 保存校验和，恢复时校验
- [ ] 异步保存不阻塞训练步
- [ ] 恢复流程含 warmup step + 3~5 步 loss 观察
- [ ] 保存频率与恢复时间目标（RTO）匹配：RTO = 检查点大小 ÷ 恢复带宽 + 重启时间
延伸阅读：Loss Spike 排查（恢复后训练不稳怎么查）、NCCL 网络调优实战（通信层稳定性）。