显存不够、训练太慢，是深度学习工程里最高频的两个抱怨，而它们各有便宜的解药：混合精度（AMP）用 FP16 换速度和省显存，梯度累积用时间换 batch size。两者正交、可叠加，PyTorch 原生支持。本文给出开启方式、度量对比与防坑清单。

前置准备

PyTorch 1.10+ 原生 AMP，无需额外依赖；确认 GPU 支持 Tensor Core（Volta 及以后）：

python -c "import torch; print(torch.__version__, torch.cuda.get_device_name(0))"
  • FP16：指数位少，单值范围小但矩阵运算快、占内存减半
  • 混合精度：前向/反向的大矩阵运算用 FP16，权重主副本与归约保持 FP32

步骤一：开启 AMP

autocast 负责按算子选精度，GradScaler 负责防止 FP16 梯度下溢：

import torch
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for x, y in loader:
    opt.zero_grad()
    with autocast():                 # 算子级自动选 FP16/FP32
        loss = criterion(model(x), y)
    scaler.scale(loss).backward()    # 放大 loss 防梯度下溢
    scaler.step(opt)                 # 反缩放后更新权重
    scaler.update()

不要手动把整个 model.half()——那会连 BatchNorm、LayerNorm 一起压到 FP16，是发散的头号原因。

步骤二：梯度累积扩大等效 batch

小显存跑不了大 batch，就攒 K 步梯度再更新一次：

ACCUM = 8   # 等效 batch = micro_batch * ACCUM
for i, (x, y) in enumerate(loader):
    with autocast():
        loss = criterion(model(x), y) / ACCUM   # 梯度按累积步数平均
    scaler.scale(loss).backward()
    if (i + 1) % ACCUM == 0:
        scaler.step(opt)
        scaler.update()
        opt.zero_grad()
  • 学习率按等效 batch 线性缩放（batch ×8，lr ×8 起步再微调）
  • 累积期间不要 zero_grad，更新后才清

步骤三：度量三方对比

加速不是信仰，用同一配置跑 200 step 对比三组数字：

import time
def bench(tag):
    torch.cuda.reset_peak_memory_stats()
    t0 = time.time()
    for step, (x, y) in enumerate(loader):
        if step >= 200: break
        train_step(x, y)
    print(tag, round(time.time()-t0, 1), "s |",
          torch.cuda.max_memory_allocated()/1024**3, "GB")

bench("fp32")        # 基线
bench("amp")         # 期望:速度 1.5-3x,显存 -30~50%
bench("amp+accum8")  # 期望:显存再降,等效 batch 变大

记录：step 耗时、峰值显存、验证集指标——三者一起看才算结论。

步骤四：对齐 loss 曲线

加速手段必须不伤收敛：画 FP32 基线与 AMP 的 loss 曲线对比：

  • 前几百 step 允许 AMP 略高，但趋势应平行收敛
  • 最终验证指标差距应在随机种子波动范围内（同种子跑 2 次估噪声）
  • 差距超标 → 检查是否漏了 GradScaler、norm 层是否误入 FP16

常见问题

  • loss 出现 NaN：GradScaler 会自动跳过该步；持续 NaN 则把 softmax/norm 相关层显式包到 autocast(enabled=False) 里
  • 加速不明显：瓶颈在数据加载（CPU/IO）而非计算——先开 num_workers、pin_memory 再看 AMP
  • 累积后效果变差：lr 没按等效 batch 缩放，或累积步数过大导致更新次数太少
  • 分布式叠加：DDP + AMP + 累积三者可共存，注意 scaler 各卡独立、梯度同步在 step 前完成

小结

AMP 与梯度累积是显存-速度-精度三角里的两枚免费杠杆：autocast+GradScaler 拿速度与显存，累积拿等效 batch 与收敛稳定性。正确姿势是”基线 → 单开 AMP → 叠加累积 → 三方度量 + 曲线对齐”，每一步都有数字背书；任何一步指标跑偏，回到该步的防坑清单而不是盲目加 epoch。