4. 混合精度训练（AMP）：FP16/FP32原理、自动混合精度、梯度缩放与风险

闲聊时经常有同事问我：你们做分布式训练，显存快占满了，怎么办？我的第一个建议往往就是——上混合精度。说白了，就是把模型里一部分计算从 FP32 降到 FP16，用精度换显存和速度。这个技巧几乎成了我每张大集群训练的标配。今天咱们就把 AMP 的底牌翻个底朝天。

4.1 FP16 与 FP32 到底差在哪？

先看张对比表，直观感受一下差异：

特性 FP32 FP16
位数 32 位 16 位
指数位 8 位 5 位
尾数位 23 位 10 位
表示范围 ~3.4e-38 ~ ~3.4e+38 ~6.55e-5 ~ ~6.55e+4
最小正规格数 ~1.18e-38 ~6.10e-5
显存占用 100% 约 50%
计算吞吐（典型 GPU） 100% 2x ~ 8x

看到没？FP16 的表示范围窄得可怜，尤其最小值只能到 1e-5 量级。这意味着什么？我记着有一次跑一个大语言模型，梯度里很多值都小于 1e-5，结果全被下溢成 0 了。模型直接不收敛，我还以为是代码写错了。后来一查，FP16 根本 hold 不住这些 tiny values。所以就有了混合精度的故事。

我的习惯：先把模型里对精度不敏感的部分——比如大部分卷积、线性层——切到 FP16，而像 BatchNorm、loss 计算这些还是老老实实用 FP32。AMP 自动帮我们干的就是这件事：它不是一刀切，而是有选择地使用 FP16 和 FP32。

4.2 自动混合精度（torch.cuda.amp）——把选择交给框架

手动维护哪些层用 FP16、哪些用 FP32，碰上 100 层以上的模型简直想哭。PyTorch 从 1.6 开始推出 torch.cuda.amp，用法极其简单：

import torch

scaler = torch.cuda.amp.GradScaler()   # 梯度缩放器

for data, target in dataloader:
    with torch.cuda.amp.autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

你看，只需要三个步骤：

  1. 用 autocast() 包住前向传播
  2. 用 scaler.scale(loss) 替换 loss.backward()
  3. 用 scaler.step() 替换 optimizer.step()，然后 scaler.update()

为什么这里要引入一个 GradScaler？我们来画个流程图，看看 AMP 到底在背后做了什么。

自动混合精度（AMP）典型流程 输入数据 + 模型 autocast 前向传播 卷积/线性自动用FP16，BatchNorm/Loss用FP32 计算 loss (FP32) scaler.scale(loss).backward() 防止 FP16 梯度下溢 乘上放大因子 scaler.step(optimizer) 内部自动 unscale 梯度 再调用 optimizer.step() scaler.update() 检测是否出现 Inf/NaN 动态调整放大因子 下一个 batch

嗯，这张图应该能帮你理清 AMP 的核心脉络。简单说：autocast 负责在运行时自动切换精度，GradScaler 负责对付梯度下溢问题。

4.3 梯度缩放——为什么需要它？

FP16 能表示的最小正数大约是 6.1e-5。而深度学习里的梯度经常小到 1e-7、1e-8。直接存成 FP16 就等于被吃掉，变成 0。所以我们需要把 loss 先放大（比如乘一个初始值 2^16=65536），让梯度也放大，保证它们在 FP16 范围内不丢精度。反向传播完后再把梯度缩小回来。

我曾经踩过的坑：有一次我忘了调用 scaler.update()，结果放大因子一直不变。恰好那个 batch 的 loss 特别大，放大后梯度爆了 Inf，接着后面所有梯度都废了。模型直接发散。所以记住：scaler.update() 必须在每个 batch 都调用，不能偷懒。

AMP 的 GradScaler 还会自动检测是否有 Inf/NaN：

  • 如果某个 batch 反向传播后出现 Inf 或 NaN，scaler.step() 会跳过本次更新（相当于丢掉这个 batch）
  • 然后 scaler.update() 会稍微减小放大因子，避免下次再爆
  • 连续多个 batch 都正常，放大因子会缓慢增大，追求更高精度

这个自适应机制非常实用。我个人习惯把初始缩放因子设为 2**16，大多数模型都够用。

4.4 收益与风险——算清这笔账

收益非常诱人：

  • 显存减半：模型本身参数 + 中间激活值都切成 FP16，显存占用能降 40%~60%
  • 训练速度提升：V100 上 FP16 吞吐大约是 FP32 的 2 倍；A100 上 Tensor Core 更夸张，能到 8 倍
  • 通信量减半：分布式训练时梯度同步只需要传一半数据，节省带宽

但风险也不小：

风险 表现 应对策略
梯度下溢 小梯度变 0，模型不收敛 使用 GradScaler，或增大缩放因子
梯度上溢 梯度变成 Inf 或 NaN 检查缩放因子是否过大；确认 loss 没有突然爆炸
精度损失 部分层（如 attention softmax）在 FP16 下精度不足 用 autocast 白名单强制这些层使用 FP32
数值不稳定 训练曲线跳跃，波动大 适当增加 batch size 或采用 Adam 这类对 scale 不敏感的优化器

我曾在训练一个 70 亿参数的模型时，发现 loss 一直卡在 3.2 左右下不去。关掉 AMP 后，loss 能降到 2.7。后来我定位到是 attention 里的 softmax 在 FP16 下给截断了。解决方案很简单：在 autocast 上下文里手动指定 with torch.cuda.amp.autocast(enabled=False): 强制那些层用 FP32。你别说，同一个模型，只动了两行代码，loss 立马掉到 2.8。

我的建议：不要一开始就对整个模型无脑开启 AMP。先在小规模数据集上快速跑几个 epoch，对比 FP32 baseline。如果 loss 曲线能对齐，说明精度损失可接受。如果差很多，就用 torch.cuda.amp.custom_fwd 或 torch.cuda.amp.custom_bwd 手动指定哪些操作保持 FP32。这个过程的调试时间，往往能帮你省下好几天的训练排错时间。

4.5 实战 Checklist

最后，给你一份我在团队内部流传的 AMP 落地自查清单：

  1. 确认 GPU 支持 FP16（V100 以上都支持）
  2. 确认 torch.backends.cudnn.allow_tf32 等设置合理（不管 AMP 时最好也关注）
  3. 开启 autocast 包裹前向，尽量把 loss 也包在里面（但 loss 自身计算必须是 FP32）
  4. GradScaler 初始化放在训练循环外 scaler = torch.cuda.amp.GradScaler()
  5. backward 前先用 scaler.scale(loss) 替代 loss
  6. optimizer.step 前先 scaler.step(optimizer)
  7. 每个 batch 结尾务必执行 scaler.update()
  8. 监控 loss 值变化，如果连续 5 个 batch 出现 NaN/Inf，立刻检查缩放因子和模型是否发散
  9. 如果模型用了 Gradient Checkpointing，AMP 和它配合良好，但要注意 checkpoint 的 forward 同样会被 autocast 影响
  10. 保存 checkpoint 时，scaler 的状态也要一并保存（scaler.state_dict()）

我个人把上面的 checklist 写成了一个小脚本，每次开始训练时自动 log。如果某个参数的梯度在 FP16 下出现了 100% 的下溢比例，我会弹个警告。好多次它帮我拦截了潜在的不收敛问题。你说，是不是很香？

混合精度训练已经是分布式训练里最成熟、最有效的显存优化手段之一。当然它不是银弹——当模型特别大、梯度分布极其不均匀时，你可能还需要结合 ZeRO 优化器、梯度累积等手段。但我们先把 AMP 稳住，剩下的事情后面再聊。

← 上一章 📖 返回目录 下一章 →