4、混合精度训练（AMP）：FP16/BF16原理、GradScaler、实战配置

混合精度训练，说白了就是让模型在训练时「脚踩两只船」——一部分计算用半精度（FP16或BF16），另一部分用单精度（FP32）。

我刚开始接触这个技术时，心里也犯嘀咕：精度降了一半，模型还能收敛吗？后来踩过几次坑才发现，只要用对工具，这玩意儿不仅能省显存，还能把训练速度提上去一大截。

4.1 为什么需要半精度？

先看一组直观数据。一张RTX 3090显卡，显存24GB。如果你用FP32训练一个7B参数的模型，光参数就要占28GB——直接爆显存。但换成FP16，只需要14GB，瞬间就塞进去了。

你想想看，这还只是参数。还有梯度、优化器状态、激活值……每一项都能省一半。我有个项目，原本用FP32跑不动一个12层的Transformer，切到混合精度后，不仅跑起来了，batch size还翻了一倍。

  核心收益：
  显存占用降低40%-50%
计算吞吐量提升1.5-3倍（取决于硬件）
通信带宽减半（分布式训练时尤其明显）

4.2 FP16 vs BF16：到底选哪个？

很多人以为FP16和BF16只是精度不同，其实它们的「脾气」差很多。

特性	FP16	BF16
指数位	5位	8位
尾数位	10位	7位
表示范围	±65,504	±3.39×10³⁸
精度	较高（约3.3位十进制）	较低（约2.4位十进制）
硬件支持	几乎所有GPU	Ampere架构及以上

我个人习惯是：能用BF16就用BF16。为什么？因为BF16的指数范围和FP32一样，这意味着它几乎不会出现「溢出」问题。FP16经常遇到梯度变成NaN的情况，BF16就很少。

我在项目中遇到过一件事：用FP16训练一个BERT-large，loss突然跳到NaN，查了两天才发现是某个层的梯度超过了65504。换成BF16后，问题直接消失。嗯，这就是指数位多的好处。

注意：BF16的尾数位只有7位，精度确实比FP16低。如果你的任务对数值精度极其敏感（比如某些科学计算），还是得用FP32兜底。

4.3 GradScaler：为什么需要它？

FP16有个致命弱点：下溢。梯度值如果小于2⁻²⁴（约5.96×10⁻⁸），在FP16里就直接变成0了。你想想看，梯度都归零了，参数还怎么更新？

GradScaler就是来解决这个问题的。它的思路很简单：

前向传播时，把loss乘以一个缩放因子（比如1024）
反向传播时，梯度也跟着被放大了
更新参数前，再把梯度除以这个缩放因子

说白了，就是先把数值「撑大」，避免小梯度被吃掉，算完再「缩回去」。我刚开始用AMP时，总觉得这步是多余的，直到有一次忘了开GradScaler，模型loss死活降不下去……

我的经验：缩放因子初始值设为2¹⁶（65536）比较稳妥。如果连续200步没有出现inf/NaN，AMP会自动把因子翻倍；如果出现了，就减半。PyTorch的GradScaler默认就是这么干的。

4.4 实战配置：PyTorch AMP 三步走

PyTorch从1.6开始内置了AMP支持，用起来非常方便。我一般按这个模板来：

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

# 1. 初始化
scaler = GradScaler()
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    
    # 2. 清零梯度
    optimizer.zero_grad()
    
    # 3. 用autocast包裹前向传播
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    # 4. 用scaler反向传播
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

这里有几个细节要注意：

autocast只包前向传播，反向传播会自动处理
loss不要手动缩放，scaler.scale()会自动做
optimizer.step()必须用scaler.step()替代

我曾经犯过一个错：在autocast外面又手动乘了个缩放因子，结果梯度被放大了两次，loss直接飞了。嗯，这种低级错误希望大家别重蹈覆辙。

4.5 知识体系：一张图看懂AMP

下面这张SVG图，把AMP的核心流程画清楚了。我建议你保存下来，写代码时对照着看：

4.6 避坑指南

最后分享几个我踩过的坑，希望能帮你省点时间：

坑1：不要对loss进行手动缩放。我见过有人写 loss = loss * 1024，然后又用 scaler.scale(loss)，结果梯度爆炸。

坑2：BN层在FP16下表现不稳定。如果模型里有BatchNorm，建议把BN层留在FP32下计算。PyTorch的autocast会自动处理，但如果你自己写训练循环，记得检查。

坑3：分布式训练时，每个进程都要有自己的GradScaler实例。我一开始图省事，只创建了一个，结果梯度同步全乱了。

混合精度训练，说白了就是「用一半的精度，换双倍的速度和显存」。只要把GradScaler和autocast用对，大部分模型都能无缝切换。我个人现在所有新项目都默认开AMP，除非遇到精度敏感的特殊场景。

想要更多资料，联系可可

📚 课程目录

4、混合精度训练（AMP）：FP16/BF16原理、GradScaler、实战配置

4.1 为什么需要半精度？

4.2 FP16 vs BF16：到底选哪个？

4.3 GradScaler：为什么需要它？

4.4 实战配置：PyTorch AMP 三步走

4.5 知识体系：一张图看懂AMP

4.6 避坑指南