3. 混合精度训练：FP16与BF16原理、AMP自动混合精度、Loss Scaling策略、硬件支持

混合精度训练，说白了就是让模型在训练时「该快的地方快，该准的地方准」。

我刚开始接触这个技术时，心里也犯嘀咕：用低精度不会把模型训崩吗？后来踩过几次坑才明白——只要策略得当，混合精度不仅能提速，还能省显存。今天咱们就把这块掰开揉碎了讲。

3.1 FP16 与 BF16：两种低精度格式的「性格」差异

先看两张图，这是两种格式的位分布：

FP16 和 BF16 都是 16 位浮点数，但「性格」完全不同：

• FP16：5 位指数 + 10 位尾数。范围窄（最大 65504），但精度稍高。我早期在 NVIDIA V100 上做实验，经常因为梯度下溢（underflow）导致 loss 变成 NaN。
• BF16：8 位指数 + 7 位尾数。范围跟 FP32 一样大（最大 3.4e38），但精度低一些。说白了，它牺牲了尾数位，换来了跟 FP32 一样的动态范围。

3.2 AMP 自动混合精度：让框架替你操心

手动管理精度太累了。AMP（Automatic Mixed Precision）就是帮你自动决定：哪些操作用 FP16，哪些用 FP32。

我习惯用 PyTorch 的 torch.cuda.amp，代码写起来很清爽：

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

model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()  # 自动处理 loss scaling

for data, target in dataloader:
    optimizer.zero_grad()
    
    with autocast():  # 自动选择精度
        output = model(data)
        loss = loss_fn(output, target)
    
    scaler.scale(loss).backward()  # 缩放 loss，防止梯度下溢
    scaler.step(optimizer)         # 反缩放梯度，更新参数
    scaler.update()                # 动态调整缩放因子

AMP 的决策逻辑其实不复杂：

• 计算密集型操作（卷积、矩阵乘法）→ FP16，加速明显
• 精度敏感操作（Softmax、LayerNorm、Loss 计算）→ FP32，保证数值稳定
• 参数更新（权重、偏置）→ 始终维护 FP32 主副本

3.3 Loss Scaling 策略：对抗梯度下溢的「护身符」

FP16 的指数范围只有 5 位，能表示的最小正数是 2⁻²⁴ ≈ 5.96e-8。如果梯度比这个还小，就会直接变成 0——这就是梯度下溢。

Loss Scaling 的原理很简单：在反向传播前把 loss 放大，梯度也跟着放大，更新参数前再缩回来。

常见的策略有两种：

策略	原理	适用场景
固定缩放	手动设置一个常数（如 2¹⁶ = 65536）	模型稳定，梯度分布已知
动态缩放	根据梯度是否溢出自动调整	大多数情况，推荐使用

动态缩放的具体流程：

1. 初始缩放因子设为 2¹⁶
2. 如果连续 N 步没有梯度溢出 → 放大 2 倍
3. 如果出现梯度溢出 → 缩小 2 倍，跳过当前步
4. 重复直到稳定

3.4 硬件支持：你的 GPU 能跑混合精度吗？

不是所有 GPU 都支持混合精度。我整理了一张表，方便你对照：

硬件	FP16 支持	BF16 支持	Tensor Core
NVIDIA V100 (Volta)	✅ 原生	❌	✅ 第一代
NVIDIA A100 (Ampere)	✅ 原生	✅ 原生	✅ 第三代
NVIDIA H100 (Hopper)	✅ 原生	✅ 原生	✅ 第四代
AMD MI250X	✅ 原生	✅ 原生	✅
Intel Gaudi 2	✅ 原生	✅ 原生	❌

这里有个关键点：Tensor Core 是混合精度加速的核心。它能在单个时钟周期内完成 FP16 矩阵乘法并累加到 FP32，速度比普通 CUDA Core 快 8-16 倍。

我建议你跑一下这个命令，看看自己的 GPU 支持什么：

# 检查 CUDA 设备属性
python -c "import torch; print(torch.cuda.get_device_capability())"
# 返回 (8, 0) 表示 Ampere 架构，支持 BF16
# 返回 (7, 0) 表示 Volta 架构，只支持 FP16

3.5 实战建议：从入门到调优

最后分享几个我踩过的坑和总结的经验：

• 起步阶段：直接用 PyTorch AMP 的默认配置，90% 的模型都能跑通
• 遇到 NaN：先检查 loss scaling 是否正常，再检查模型是否有数值不稳定的操作
• 性能调优：用 torch.cuda.amp.autocast(enabled=True) 配合 torch.profiler 分析瓶颈
• 多卡训练：DDP 下每个进程独立维护 scaler，但要注意梯度同步的顺序

嗯，混合精度训练这块其实不难，关键是理解「为什么」和「什么时候用」。你只要记住：FP16 省显存但容易下溢，BF16 范围大但精度低，AMP 帮你自动平衡，Loss Scaling 是最后的保险。

下次训练大模型时，试试这些技巧。你会发现，显存占用直接砍半，速度还能快 2-3 倍。