4. 混合精度训练（AMP）：FP16与BF16原理、AMP实现、损失缩放、实战配置

好，咱们今天聊聊混合精度训练。说实话，这玩意儿现在已经是炼丹的标配了。我最早接触AMP是在训练一个超大BERT模型的时候，显存死活塞不下，后来上了AMP，不仅塞下了，速度还快了将近一倍。嗯，今天就把我这些年踩过的坑和总结的经验，一次性讲清楚。

4.1 FP16与BF16：两种低精度格式的原理

先说说FP16和BF16到底是个啥。说白了，它们都是16位浮点数，但内部结构不一样。

格式	符号位	指数位	尾数位	表示范围	精度
FP32	1	8	23	~3.4e38	高
FP16	1	5	10	~6.5e4	中
BF16	1	8	7	~3.4e38	低

你看这个表，FP16的指数位只有5位，所以它能表示的最大值才6万多一点。这意味着什么？你想想看，如果梯度里有个数值超过65504，直接就溢出了。我在项目中遇到过这种情况，训练到一半loss突然变成NaN，查了半天才发现是梯度爆炸了。

BF16就不一样了，它保留了8位指数，和FP32一样。所以它的表示范围跟FP32一样大，不会溢出。但代价是尾数位只有7位，精度比FP16还低。不过实际训练中，这个精度损失通常可以接受。

核心结论：

FP16：范围小，精度中等。适合梯度值稳定的场景。
BF16：范围大，精度低。适合大模型训练，不用担心溢出。
我个人习惯：能用BF16就用BF16，省心。

4.2 AMP的实现原理

AMP全称是Automatic Mixed Precision。它不是什么黑魔法，说白了就是：该用FP16的地方用FP16，该用FP32的地方用FP32。

具体怎么做呢？我画个流程图你就明白了。

看到没？核心就两点：

前向和反向传播：用FP16/BF16计算，省显存、提速度。
权重更新和损失计算：用FP32，保证精度。

为什么权重更新要用FP32？因为低精度下的小梯度根本加不上去。你想想看，如果权重是0.001，梯度是0.00001，在FP16下这个梯度直接就变成0了。所以主权重必须用FP32维护一份副本。

4.3 损失缩放（Loss Scaling）

损失缩放是AMP里最容易被忽视，但也最关键的一环。我曾经因为没开损失缩放，训练了一个星期的模型，结果发现loss曲线跟心电图似的，上下乱跳。

为什么会这样？因为FP16的表示范围太小了。如果梯度值小于约6e-5，在FP16下就直接变成0了。这就是所谓的下溢出。

解决办法很简单：在反向传播之前，把损失值乘以一个缩放因子（比如1024），让梯度变大。等梯度计算完，再除以这个缩放因子，恢复原值。

我的经验：

初始缩放因子：推荐 2^16 = 65536
如果出现NaN或Inf，自动减半缩放因子
如果连续N步没有溢出，自动加倍缩放因子
PyTorch的AMP默认就带这个机制，不用自己写

4.4 实战配置：PyTorch AMP

好了，理论说完了，咱们直接上代码。这是我在实际项目中用的AMP配置模板。

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

# 初始化缩放器
scaler = GradScaler(init_scale=2**16)

# 模型、优化器、数据加载器
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
dataloader = get_dataloader()

for epoch in range(num_epochs):
    for batch in dataloader:
        inputs, labels = batch
        inputs, labels = inputs.cuda(), labels.cuda()
        
        # 清零梯度
        optimizer.zero_grad()
        
        # 前向传播（自动混合精度）
        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, labels)
        
        # 反向传播（带损失缩放）
        scaler.scale(loss).backward()
        
        # 梯度裁剪（可选）
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        # 权重更新
        scaler.step(optimizer)
        scaler.update()

避坑指南：

我曾经把autocast()放到了整个训练循环外面，结果所有操作都变成了FP16，权重更新也变成了FP16，模型直接不收敛。
记得autocast()只包裹前向传播和损失计算，不要包裹反向传播。
如果用了梯度裁剪，记得先调用scaler.unscale_(optimizer)，否则梯度会被缩放两次。

4.5 实战配置：Hugging Face Trainer

如果你用Hugging Face的Trainer，那就更简单了。一行代码搞定：

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./results",
    fp16=True,          # 启用FP16
    bf16=False,         # 如果GPU支持BF16，可以设为True
    fp16_full_eval=True, # 评估时也用FP16
    dataloader_pin_memory=True,
)

我个人习惯：如果用的是A100或H100，直接开BF16。如果是V100或T4，开FP16。BF16在V100上不支持，别问我怎么知道的，都是泪。

4.6 性能对比与建议

配置	显存节省	速度提升	精度损失	推荐场景
FP32（无AMP）	0%	1x	无	小模型、调试阶段
FP16 AMP	~40%	1.5x-2x	可忽略	V100/T4/3090
BF16 AMP	~40%	1.5x-2x	可忽略	A100/H100/4090

你看这个表，显存节省40%是实打实的。我训练一个7B模型，原来需要4张A100，开了AMP之后3张就够了。省下来的钱，够吃好几顿火锅了。

最后总结一下我的AMP配置清单：

检查GPU是否支持BF16（A100/H100/4090支持，V100/T4不支持）
能用BF16就用BF16，否则用FP16
始终开启损失缩放（PyTorch默认开启）
梯度裁剪前记得unscale
主权重保持FP32副本

嗯，AMP这部分就讲到这里。记住一句话：AMP不是可选项，是必选项。尤其是大模型训练，不开AMP你连模型都塞不进显存。

想要更多资料，联系可可