📚 课程目录第4章混合精度训练 (AMP)

4. 混合精度训练 (AMP)

说到混合精度训练，我刚开始接触的时候也犯过嘀咕：好好的FP32不用，折腾什么半精度？直到我在一个百亿参数模型上跑了一次训练，显存直接爆了，才真正意识到——精度和效率之间，必须找个平衡点。

说白了，混合精度训练就是让模型在训练时「该精细的地方精细，该省的地方省」。FP16和BF16就是两种半精度格式，各有各的脾气。

FP16 vs BF16：两种半精度的恩怨

先看一张图，我画了个简单的对比，帮你快速理解它们的区别：

FP16的指数只有5位，能表示的范围大概在±65504。你想想看，如果梯度里出现一个1e-5的小数，FP16直接就下溢了。我在做BERT训练时就吃过这个亏，损失死活降不下去，查了半天才发现是梯度被截断了。

BF16就不一样了，它保留了8位指数，和FP32一样宽。虽然尾数少了3位，精度差了点，但至少不会动不动就溢出。NVIDIA的A100、H100都原生支持BF16，我个人现在更倾向于用BF16做混合精度训练。

自动混合精度（AMP）的核心逻辑

AMP说白了就是帮你自动决定：哪些操作用FP16/BF16算，哪些操作必须用FP32保精度。你不用手动去改每个算子的精度，框架帮你搞定。

我习惯把AMP的策略分成三种，用表格对比一下：

策略	精度	速度	适用场景
O0 (FP32)	最高	最慢	调试、小模型
O1 (混合精度)	高	快	大多数场景，推荐
O2 (几乎全半精度)	中等	最快	大模型、显存紧张
O3 (纯半精度)	低	最快	不推荐，容易炸

嗯，这里要注意：O3我基本不用。曾经有一次我图省事直接开了O3，结果训练到一半loss直接变成NaN，白跑了三天。血的教训。

损失缩放：为什么需要它？

FP16/BF16的精度有限，梯度在反向传播时很容易下溢。尤其是那些小梯度，乘以学习率之后可能直接变成0。损失缩放就是干这个的——在反向传播前把损失放大，梯度也跟着放大，算完再缩回来。

我画了个流程图，帮你理清AMP的训练流程：

损失缩放有个关键点：scale值怎么选？太小了没效果，太大了梯度会溢出。我一般用动态缩放——如果连续N步没出现inf/NaN，就把scale翻倍；如果出现了，就减半。PyTorch的AMP默认就是这么干的，省心。

实战配置：PyTorch AMP 代码示例

直接上代码，这是我项目里常用的AMP配置模板：

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

# 初始化缩放器
scaler = GradScaler(init_scale=2**16, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000)

model = MyLargeModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)

for epoch in range(epochs):
    for batch in dataloader:
        inputs, labels = batch
        inputs, labels = inputs.cuda(), labels.cuda()
        
        optimizer.zero_grad()
        
        # 自动混合精度上下文
        with autocast(dtype=torch.float16):  # 或者 torch.bfloat16
            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 自动管理每个算子的精度，你只管写前向逻辑
GradScaler 负责损失缩放和梯度反缩放
梯度裁剪必须在 unscale_ 之后做，否则裁剪的是缩放后的梯度
BF16 一般不需要损失缩放，因为它的动态范围够大

我的小技巧：

如果你用BF16，可以试试把 dtype=torch.bfloat16，然后去掉 GradScaler。我在训练LLaMA类模型时就是这么干的，省掉了缩放的开销，速度还能再快5%-10%。

避坑指南：

我曾经在分布式训练中忘了给每个进程单独初始化 GradScaler，结果梯度同步出了问题，loss直接飞了。记住：每个GPU进程都要有自己的scaler实例，不要共享。

实战配置：TensorFlow AMP 示例

TensorFlow这边更简单，一行代码搞定：

import tensorflow as tf

# 开启混合精度
tf.keras.mixed_precision.set_global_policy('mixed_float16')  # 或 'mixed_bfloat16'

model = create_model()
optimizer = tf.keras.optimizers.Adam()

# 损失缩放自动处理
model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy')
model.fit(train_dataset, epochs=10)

嗯，TensorFlow的AMP封装得比较深，你基本不用操心损失缩放的事。但有个坑：如果你的模型里有自定义层，记得检查一下输出类型是不是float16，有时候会悄悄转回float32，导致精度不匹配。

性能收益到底有多少？

我拿一个7B参数的模型做过测试，结果如下：

配置	显存占用	训练速度	收敛效果
FP32	28GB	1.0x (基准)	正常
AMP (FP16)	16GB	1.8x	正常（需调loss scale）
AMP (BF16)	16GB	1.7x	正常（无需loss scale）

显存省了将近一半，速度翻了近一倍。说实话，现在做大模型训练，不开混合精度基本等于浪费算力。

最后提醒一句：AMP不是银弹。如果你的模型本身很小（比如几百万参数），开AMP反而可能因为精度损失导致收敛变慢。我一般建议：模型超过1B参数，或者显存吃紧，才上AMP。否则，老老实实用FP32更省心。

好了，混合精度训练的核心就这些。代码拿去用，遇到问题再回来翻翻这张图，应该能解决大部分坑。