混合精度训练：FP16/BF16原理、AMP自动混合精度、损失缩放策略

好，咱们来聊聊混合精度训练。说实话，这玩意儿现在几乎是训练大模型的标配了。我刚开始接触的时候，也觉得不就是把精度砍一半嘛，能有多大区别？后来真上手了才发现，这里面的门道还挺深。

FP16 和 BF16：两种低精度格式

先说说 FP16 和 BF16 到底是个啥。你想想看，我们平时用的 FP32，也就是单精度浮点数，占 4 个字节，32 个比特位。这 32 位里，1 位是符号，8 位是指数，23 位是尾数。

FP16 呢，就是把它砍成 16 位。1 位符号，5 位指数，10 位尾数。好处很明显——显存占用直接减半，计算速度也能翻倍。但问题也来了：指数位从 8 位砍到 5 位，能表示的数值范围就小了很多。

BF16 是 Google 搞出来的。它也是 16 位，但分配方式不一样：1 位符号，8 位指数，7 位尾数。你看，指数位跟 FP32 一样多，所以数值范围没缩水。代价是尾数精度降低了。

核心区别一句话：

  • FP16：范围小，精度还行。容易溢出或下溢。
  • BF16：范围大，精度差点。但训练起来更稳。

我在项目中遇到过用 FP16 训练一个 NLP 模型，梯度直接下溢到零，模型死活不收敛。换成 BF16 之后，问题就解决了。嗯，这就是数值范围不够的典型场景。

AMP：自动混合精度

手动去管哪些操作用 FP16、哪些用 FP32，太累了。NVIDIA 搞了个 AMP（Automatic Mixed Precision），说白了就是帮你自动做这个决策。

AMP 的核心逻辑其实不复杂：

  • 计算密集的操作（比如卷积、矩阵乘法）用 FP16/BF16
  • 对精度敏感的操作（比如 Softmax、LayerNorm）用 FP32
  • 主权重始终维护一份 FP32 的副本

我习惯用 PyTorch 的 torch.cuda.amp，用起来特别简单：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

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

你看，就多了几行代码。但这里有个坑——autocast 只负责前向和 loss 计算，反向传播它不管。反向传播的精度控制，靠的是 GradScaler。

我的小技巧：

如果你用 BF16，可以试试把 autocast 的 dtype 参数设成 torch.bfloat16。BF16 不需要损失缩放，因为它的数值范围够大。

损失缩放策略

为什么要做损失缩放？说白了就是 FP16 的数值范围太小，梯度容易下溢成零。你想想看，梯度值可能小到 1e-6，但 FP16 能表示的最小正数是 6e-8 左右。再小就变零了。

损失缩放的做法很简单：在反向传播之前，先把 loss 乘上一个缩放因子。这样梯度也跟着被放大了，就不会下溢了。反向传播完之后，再把梯度除以这个因子，恢复原来的大小。

缩放因子的选择是个技术活。设太大，梯度可能溢出；设太小，又解决不了下溢问题。

我见过有人手动设成 1024，然后训练到一半 loss 突然炸了。嗯，这就是缩放因子太大，梯度溢出了。

GradScaler 默认用的是动态缩放策略：

  • 初始缩放因子一般是 2^16 = 65536
  • 如果连续 N 步没有梯度溢出，就把缩放因子翻倍
  • 如果某一步梯度溢出了，就跳过这一步，把缩放因子减半

这个策略其实挺聪明的。它自动适应你的模型和数据的特性，不需要你手动调参。

注意：

损失缩放只对 FP16 有效。BF16 不需要这个，因为它的指数位跟 FP32 一样多，数值范围足够大。

知识体系结构图

下面这张图是我自己整理的混合精度训练的核心逻辑，你看一眼就能明白整体脉络：

混合精度训练核心逻辑 FP16 / BF16 格式 AMP 自动混合精度 损失缩放策略 FP16：5位指数，范围小，精度高 BF16：8位指数，范围大，精度低 计算密集型 → FP16/BF16 精度敏感型 → FP32 主权重维护 FP32 副本 前向：loss × 缩放因子 反向：梯度 ÷ 缩放因子 动态调整：无溢出则翻倍 目标：显存减半 + 速度翻倍 + 精度不掉

实际使用中的注意事项

我总结几个实际踩过的坑，你注意一下：

  1. Batch Size 别太大。混合精度训练下，显存省下来了，很多人就拼命加大 batch size。但 batch size 太大，梯度更新方向会变钝，模型反而不好收敛。
  2. 监控梯度统计。我习惯在训练脚本里加几行代码，打印梯度的均值和方差。如果发现梯度突然变成零或者突然爆炸，赶紧停。
  3. 验证集精度要盯紧。混合精度训练偶尔会导致精度轻微下降。如果验证集精度掉了超过 0.5%，建议检查一下是不是某些层对精度太敏感。

避坑指南：

我曾经在训练一个 7B 参数的模型时，用了 FP16 混合精度，结果训练到一半 loss 突然变成 NaN。排查了半天，发现是某个自定义的激活函数在 FP16 下数值不稳定。换成 BF16 之后，问题就消失了。所以，如果你遇到 NaN，先试试 BF16。

好了，混合精度训练的核心内容就这些。说白了就是三件事：选对格式、用对工具、调好缩放。你动手试一次，就能体会到显存省下来、速度提上去的快感了。

想要更多资料，联系可可

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