4. 混合精度训练(AMP)：FP16/BF16原理、动态损失缩放、实战配置

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

为什么要这么折腾？因为显存不够用啊。我刚开始做大模型训练时，一个7B的模型光参数就要14GB显存，加上优化器状态、梯度、激活值，一张A100 80G根本塞不下。后来用了AMP，同样的模型直接省了40%显存，训练速度还快了2倍。

嗯，这里要注意：半精度不是万能的。它的数值范围窄，精度低，搞不好就会梯度下溢或者溢出。所以我们需要一套机制来兜底——这就是动态损失缩放要做的事。

4.1 FP16 vs BF16：到底选哪个？

先看一张对比表，我直接给你结论：

特性 FP16 BF16
指数位 5位 8位
尾数位 10位 7位
数值范围 ±65,504 ±3.39×10^38
精度 高（约3.3位十进制） 低（约2.3位十进制）
硬件支持 几乎所有GPU A100/RTX30xx+
适用场景 小模型、低精度需求 大模型、梯度易下溢

我个人习惯：能用BF16就别用FP16。为什么？BF16的指数位和FP32一样，数值范围大，不容易溢出。我在项目中遇到过用FP16训练BERT-large，loss直接变成NaN，排查了半天才发现是梯度溢出了。换成BF16后，稳如老狗。

但BF16也有缺点——精度低。尾数只有7位，相当于只能精确到小数点后2位。不过对于深度学习来说，这通常够用了。你想想看，模型参数动辄上亿，少一两位精度影响不大。

核心结论：

  • 有BF16硬件（A100/H100/RTX3090+）→ 无脑选BF16
  • 只有FP16硬件（V100/T4/RTX2080）→ 用FP16 + 动态损失缩放
  • 实在不行 → 老老实实FP32，别折腾

4.2 动态损失缩放：AMP的灵魂

损失缩放（Loss Scaling）是干嘛的？说白了就是给梯度「加个放大镜」。

FP16能表示的最小正数是约6×10⁻⁸。如果梯度比这个还小，就会变成0——这叫「下溢」。训练初期梯度通常很大，但到了后期，梯度会变得非常小，尤其是那些远离输出层的层。

我刚开始用AMP时，发现模型训练到一半loss就不动了。查了半天，原来是梯度下溢了——所有小梯度都被FP16「吃掉」了，模型根本学不动。

解决方案就是损失缩放：

# 伪代码：动态损失缩放
scale = 2^16  # 初始缩放因子，约65536

for step in range(max_steps):
    # 前向传播（FP16）
    loss = model.forward(inputs)
    
    # 缩放损失
    scaled_loss = loss * scale
    
    # 反向传播（FP16）
    scaled_loss.backward()
    
    # 检查梯度是否溢出
    if any(grad.isnan() or grad.isinf() for grad in model.parameters()):
        # 溢出！缩小缩放因子
        scale = max(scale / 2, 1)
        # 跳过这一步，不更新参数
        optimizer.zero_grad()
        continue
    
    # 梯度反缩放（转回FP32）
    for param in model.parameters():
        param.grad.data = param.grad.data / scale
    
    # 更新参数（FP32）
    optimizer.step()
    
    # 如果连续N步没有溢出，增大缩放因子
    if no_overflow_steps >= 2000:
        scale = min(scale * 2, 2^24)

嗯，这里有几个关键点：

  • 初始缩放因子：我一般设2^16（65536）。太大容易溢出，太小起不到作用
  • 溢出处理：一旦检测到NaN或Inf，直接跳过这一步，不更新参数。别慌，偶尔跳一两步不影响
  • 缩放因子调整：连续2000步无溢出就翻倍，有溢出就减半。这个「2000」是我试出来的经验值

我曾经踩过的坑：

有一次我把缩放因子设成了2^24，结果前几步就溢出了，然后缩放因子被不断减半，最后变成1，等于没缩放。后来我改成2^16起步，稳了。

还有一次我忘了在反向传播前缩放loss，结果梯度直接爆炸。记住：先缩放，再反向传播，最后反缩放，顺序不能乱。

4.3 实战配置：PyTorch AMP 三步走

PyTorch从1.6开始内置了AMP支持，用起来很简单。我直接给你一个完整的训练循环模板：

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

# 1. 初始化
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler()  # 自动管理损失缩放

# 2. 训练循环
for epoch in range(num_epochs):
    for batch in dataloader:
        inputs, labels = batch
        inputs, labels = inputs.cuda(), labels.cuda()
        
        # 清空梯度
        optimizer.zero_grad()
        
        # 3. 使用autocast自动管理精度
        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, labels)
        
        # 4. 缩放 + 反向传播
        scaler.scale(loss).backward()
        
        # 5. 反缩放 + 更新参数
        scaler.step(optimizer)
        
        # 6. 更新缩放因子
        scaler.update()

看到没？就多了三行代码：

  • GradScaler()：创建缩放器
  • with autocast()：自动管理哪些算子用FP16，哪些用FP32
  • scaler.scale() 和 scaler.step()：自动缩放和反缩放

个人经验：

autocast会自动处理大部分算子，但有些操作需要手动干预：

  • BatchNorm：建议用FP32，否则训练不稳定
  • Softmax：自动用FP32，不用管
  • Embedding：建议用FP32，尤其是词表大的时候
  • 自定义算子：如果用了CUDA扩展，记得检查是否支持FP16

4.4 进阶技巧：混合精度 + 分布式训练

当你把AMP和分布式训练（DDP）结合起来时，有个坑要注意：梯度同步的顺序。

正确的做法是：

from torch.nn.parallel import DistributedDataParallel as DDP
from torch.cuda.amp import autocast, GradScaler

model = DDP(model)
scaler = GradScaler()

for batch in dataloader:
    optimizer.zero_grad()
    
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, labels)
    
    # 先缩放，再反向传播
    scaler.scale(loss).backward()
    
    # 再同步梯度（DDP自动做）
    # 最后更新参数
    scaler.step(optimizer)
    scaler.update()

我曾经犯过一个错误：在scaler.scale(loss).backward()之前手动调了model.require_backward_grad_sync = False，结果梯度没同步，模型训练出来完全不对。记住：DDP的梯度同步是在backward()中自动触发的，不要手动干预。

4.5 性能调优：到底能省多少？

我拿一个实际项目的数据给你看：

配置 显存占用 训练速度 精度（验证集）
FP32 纯精度 24.3 GB 1.0x（基准） 92.5%
FP16 AMP 15.8 GB 1.8x 92.3%
BF16 AMP 15.8 GB 1.9x 92.4%

看到没？显存省了35%，速度翻了近一倍，精度几乎没掉。这就是AMP的魅力。

但要注意：不是所有模型都适合AMP。如果你的模型对精度极其敏感（比如医学影像分割），或者训练过程中loss经常震荡，那还是老老实实用FP32吧。

总结一下AMP的适用条件：

  • ✅ 大模型（参数量 > 1B）—— 显存不够用
  • ✅ 训练速度慢 —— 想提速
  • ✅ 梯度分布均匀 —— 不容易下溢或溢出
  • ❌ 小模型（参数量 < 100M）—— 没必要，FP32更快
  • ❌ 精度要求极高 —— 比如科学计算、金融风控
  • ❌ 训练不稳定 —— 先解决收敛问题再考虑AMP

最后说一句：AMP不是银弹。它解决的是显存和速度问题，不是模型效果问题。如果你的模型本身就不收敛，加AMP只会让情况更糟。先把模型调稳了，再考虑优化。

混合精度训练(AMP)核心流程 输入数据 (FP32) 前向传播 (autocast自动管理精度) FP16/BF16算子 FP32算子 计算损失 (FP32) 损失缩放 (Loss × Scale) 反向传播 (FP16/BF16) 梯度反缩放 → 参数更新 (FP32) 关键组件 • GradScaler • autocast • 动态缩放因子 • 溢出检测 • 跳过步骤 • 连续无溢出→放大 • 溢出→缩小

想要更多资料，联系可可

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