4. 混合精度训练 (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还低。不过实际训练中，这个精度损失通常可以接受。

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），让梯度变大。等梯度计算完，再除以这个缩放因子，恢复原值。

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()

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