4、混合精度训练：什么是混合精度训练、AMP自动混合精度、手动控制精度

说到混合精度训练，我得先坦白一件事。

我刚入行那会儿，总觉得训练模型就是堆算力。显存不够？换更大的卡呗。直到有一次，我接了个超大模型的项目，A100 80G 都撑不住。那时候我才意识到——显存优化不是选择题，是生存题。

混合精度训练，就是那个帮我「活下来」的关键技术。

4.1 什么是混合精度训练？

先问个问题：你训练模型时，默认用的什么精度？

大部分人回答是 FP32（32位浮点数）。没错，PyTorch 默认就是它。但你想过没有——真的每一步都需要这么高的精度吗？

我在项目中观察过，很多层的梯度其实很小，用 FP16（16位浮点数）完全够用。FP32 占 4 字节，FP16 只占 2 字节。显存直接砍半！

但问题来了：FP16 的动态范围太小，容易溢出。怎么办？

混合精度训练的思路很简单：

前向和反向传播：用 FP16 计算，省显存、加速
权重更新：用 FP32 保存主权重，保证精度
损失缩放：把损失值放大，避免梯度下溢

说白了，就是「该省的地方省，该保的地方保」。我习惯把 FP32 比作「保险柜」，FP16 比作「零钱包」。日常花销用零钱包，大额交易才开保险柜。

核心收益：显存占用减少 40%-50%，训练速度提升 1.5-3 倍。

4.2 AMP 自动混合精度

手动管理精度切换？太累了。PyTorch 1.6 之后推出了 torch.cuda.amp，也就是 AMP（Automatic Mixed Precision）。

AMP 的核心是两个组件：

GradScaler：自动做损失缩放，防止梯度下溢
autocast：自动选择每步用 FP16 还是 FP32

我刚开始用 AMP 时，其实有点不放心。总觉得「自动」的东西不靠谱。后来在 BERT 训练中试了一次，效果出奇的好——显存从 12G 降到 7G，速度还快了 40%。

来看个标准用法：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

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

嗯，就这么几行代码。你只需要把前向计算包在 autocast() 里，然后用 scaler 处理反向传播和优化器更新。其他都不用管。

我的习惯：AMP 默认对大部分操作都有效。但像 softmax、layer_norm 这些操作，AMP 会自动用 FP32 计算，保证数值稳定性。

4.3 手动控制精度

AMP 虽好，但有时候你得亲自下场。为什么？

我在做 NLP 模型时遇到过一个问题：某些自定义的 attention 层，AMP 自动选择了 FP16，结果 loss 直接变成 NaN。排查了半天，发现是 FP16 精度不够，导致 softmax 溢出。

这时候就需要手动控制精度了。

手动控制主要有两种方式：

4.3.1 模型级别的精度控制

model.half()  # 整个模型转 FP16
model.float() # 整个模型转 FP32

简单粗暴，但不够灵活。我一般只在推理时用 model.half() 加速。

4.3.2 层级别的精度控制

更精细的做法是：

class MixedPrecisionModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3).half()   # 这个层用 FP16
        self.conv2 = nn.Conv2d(64, 128, 3).float() # 这个层用 FP32
    
    def forward(self, x):
        x = self.conv1(x.half())   # 输入转 FP16
        x = self.conv2(x.float())  # 输入转 FP32
        return x

你想想看，这样就能针对每个层做定制。比如：

卷积层：对精度不敏感，用 FP16 省显存
BatchNorm 层：必须用 FP32，否则训练不稳定
Embedding 层：建议用 FP32，避免精度损失

我曾经踩过的坑：手动控制精度时，一定要注意输入输出的类型匹配。FP16 和 FP32 混用，很容易出现类型不匹配的错误。建议在 forward 函数开头统一做类型转换。

4.4 知识体系图

下面这张图，是我自己总结的混合精度训练核心逻辑。你看一眼就能明白整个流程：

4.5 实战建议

最后，分享几个我在项目中总结的经验：

场景	推荐做法	原因
小模型（<100M 参数）	AMP 自动混合精度	省心，效果明显
大模型（>1B 参数）	AMP + 手动调整关键层	避免精度瓶颈
推理阶段	model.half()	极致加速
自定义算子	手动控制精度	AMP 可能不兼容

小技巧：如果你不确定哪些层该用 FP16，可以先用 AMP 跑一次，然后打印出每个层的精度选择。PyTorch 提供了 torch.cuda.amp.autocast.list_funcs() 来查看。

混合精度训练，说白了就是「用更少的资源，做更多的事」。我个人觉得，这是每个深度学习工程师必须掌握的技能。不管你是做 CV 还是 NLP，显存永远不够用，速度永远想更快。

嗯，今天就聊到这儿。记住：AMP 是默认选项，手动控制是进阶技能。先用好 AMP，再根据需求做手动优化。

📚 课程目录

4、混合精度训练：什么是混合精度训练、AMP自动混合精度、手动控制精度

4.1 什么是混合精度训练？

4.2 AMP 自动混合精度

4.3 手动控制精度

4.3.1 模型级别的精度控制

4.3.2 层级别的精度控制

4.4 知识体系图

4.5 实战建议