混合精度训练 AMP：FP16/BF16 加速与稳定性

什么是混合精度训练？

混合精度训练（Automatic Mixed Precision, AMP）是一种在深度学习训练中同时使用半精度（16位浮点）和单精度（32位浮点）的技术。其核心思想是：用更快的低精度计算来加速训练过程，同时在关键环节保留高精度以保证模型收敛和数值稳定性。

通常混合精度训练会使用以下两种低精度格式之一：

• FP16：IEEE 754 半精度浮点，指数位 5 位，尾数位 10 位
• BF16：Brain Floating Point，指数位 8 位，尾数位 7 位（与 FP32 指数位相同）

两种格式相比 FP32 都节省了一半的内存占用和带宽，但数值特性差异极大，直接影响训练的便捷性和稳定性。

为什么需要混合精度？

现代 GPU（如 NVIDIA Volta、Turing、Ampere、Hopper 架构）以及 TPU 都配备了专门的低精度计算单元（Tensor Cores），它们在 FP16/BF16 下的吞吐量是 FP32 的数倍。具体收益包括：

• 训练速度提升：矩阵乘法和卷积可以跑到 2~6 倍的理论算力
• 显存占用减半：模型参数、梯度、优化器状态都可用 FP16/BF16 存储，可以增大 batch size 或训练更大模型
• 通信带宽降低：分布式训练中传输低精度梯度数据量减半，减少通信瓶颈

但单纯将所有运算都切换为 FP16 几乎一定会导致训练崩溃，原因在于半精度能表示的动态范围极其有限。FP16 能表示的最小正规范化数约为 6e-8，最大约 65504，一旦梯度值过小就会下溢为 0，过大则上溢为 Inf；而 BF16 虽然指数范围与 FP32 相同，不会轻易上溢，但尾数只有 7 位，舍入误差更大。

FP16 vs BF16：如何选择？

选择建议：

• 如果使用 NVIDIA Ampere/Hopper 或更高架构，优先选择 BF16：无需损失缩放，训练稳定，代码改动最小。
• 对于较老的 GPU（V100、T4）或需要严格控制精度敏感模型的场景（如某些强化学习任务），FP16 + 损失缩放仍是可靠选择。

FP16 混合精度的三大核心挑战

在 Volta/Turing 时代，AMP 以 FP16 为主，面临三大经典问题：

1. 梯度下溢：绝对值小于约 6e-8 的梯度在 FP16 中直接变成 0，阻止权重更新。
2. 损失值超出范围：FP16 最大可表示 65504，训练初期或大型模型可能出现超出该范围的 loss 值，产生 NaN。
3. 部分运算必须保持 FP32：例如 softmax、归一化、求 loss 的平均值等，这类运算对精度极为敏感，在 FP16 中容易累积误差。

解决方案：损失缩放与前向-反向精度管理

为了解决 FP16 的数值范围不足，损失缩放（Loss Scaling） 是最关键的技术。

静态损失缩放

基本思想：将损失值在反向传播之前乘以一个大的常数（如 1024），从而将微小的梯度值也放大到 FP16 可表示的范围。反向传播完成后，再将梯度除以相同常数，恢复真实梯度，用于 FP32 的参数更新。

loss = model(input)
scaled_loss = loss * scale_factor   # 放大
scaled_loss.backward()
for param in model.parameters():
    if param.grad is not None:
        param.grad.data /= scale_factor   # 缩小
optimizer.step()

手动实现的缺点是缩放因子固定，无法适应训练过程中 loss 的变化。因子太小，梯度仍会下溢；因子太大，梯度可能上溢，导致 Inf。

动态损失缩放

现代 AMP 实现（如 NVIDIA apex、PyTorch 的 torch.cuda.amp）通常会采用动态损失缩放：从初始高缩放因子开始，每次迭代检查梯度中是否出现 Inf/NaN。若出现，则跳过本次权重更新并降低缩放因子；若连续多次迭代没有 Inf，则尝试增大因子。这种自动调整机制极大简化了使用。

PyTorch 中的自动混合精度

PyTorch 从 1.6 版本开始内置了 torch.cuda.amp，提供完整的 AMP 支持。

基本使用流程

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

model = ...                   # 定义模型
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)
scaler = GradScaler()        # 用于 FP16 的动态损失缩放（BF16 可不启用）

for data, target in dataloader:
    optimizer.zero_grad()

    # 前向传播自动插入低精度运算
    with autocast(device_type='cuda', dtype=torch.float16):
        output = model(data)
        loss = loss_fn(output, target)

    # 反向传播使用损失缩放
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

关键组件说明

• autocast 上下文管理器：在其作用域内，PyTorch 会根据运算类型自动选择精度。大部分矩阵乘法、卷积会用 FP16 执行，而 softmax、loss 计算等敏感操作自动回退到 FP32。该自动调度列表由框架维护，开发者无需手动干预。
• GradScaler：封装了动态损失缩放逻辑。
  • scaler.scale(loss) 返回缩放后的 loss，对其调用 .backward() 会生成缩放后的梯度。
  • scaler.step(optimizer) 会在内部执行梯度反缩放，并检查是否存在 Inf/NaN。若存在 Inf，会自动跳过 optimizer.step()。
  • scaler.update() 根据检查结果调整缩放因子。

BF16 的简化用法

如果使用 BF16 且硬件支持（如 A100、H100），可以不使用 GradScaler，因为 BF16 的指数范围与 FP32 相同，下溢风险极低：

with autocast(device_type='cuda', dtype=torch.bfloat16):
    output = model(data)
    loss = loss_fn(output, target)

loss.backward()
optimizer.step()

此时训练过程如同全精度一样简单，但享受了显存和速度红利。

混合精度下的模型保存与加载

权重和优化器状态通常仍以 FP32 保存，保证精度和可复现性。加载时同样恢复为 FP32，直到再次进入 autocast 作用域时部分转换为低精度计算。因此保存和加载代码无需特别改动，除非显式想用 FP16/BF16 保存（不推荐，可能导致精度丢失）。

常见问题与最佳实践

1. 什么模型适合混合精度？

几乎所有常见模型（CNN、Transformer、RNN）均可受益。但有些任务对数值极为敏感（如某些对抗训练、强化学习的值函数逼近），建议先用小规模实验验证收敛性。

2. 推理可以用混合精度吗？

可以，推理时同样可使用 .half() 或 .bfloat16() 将模型转换为半精度，配合 torch.no_grad() 和 autocast 实现性能提升。推理通常不需要损失缩放。

3. 验证或测试阶段需要 AMP 吗？

一般不需要。可关闭 autocast，或直接在全精度下运行验证，避免精度影响评估指标。如果显存压力大，也可开启 autocast，但要注意评估指标的稳定性。

4. 如何检查是否产生了 Inf/NaN？

当 GradScaler 发现 Inf/NaN 时会跳过 optimizer.step() 并减少缩放因子。你可以打印 scaler.get_scale() 来观察缩放因子的变化。如果因子持续下降甚至降到 1，说明训练不稳定，需检查学习率、权重初始化或损失函数。

5. 混合精度与梯度累积结合

使用梯度累积时，每次 backward() 调用都会累积缩放后的梯度。应在所有 backward() 完成后再调用 scaler.step(optimizer) 和 scaler.update()。确保每次 step 之前执行一次 scaler.unscale_(optimizer)（scaler.step 会自动执行反缩放，但如果需要在 step 之前裁剪梯度，必须手动调用）：

scaler = GradScaler()
for i, (data, target) in enumerate(dataloader):
    with autocast():
        output = model(data)
        loss = loss_fn(output, target) / accumulation_steps

    scaler.scale(loss).backward()

    if (i + 1) % accumulation_steps == 0:
        scaler.unscale_(optimizer)  # 反缩放后裁剪梯度
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

6. 多 GPU 分布式训练

在 DistributedDataParallel (DDP) 中，AMP 的使用与单卡一致：每个进程独立创建 autocast 和 GradScaler。PyTorch 的 DDP 会在反向传播时自动 allreduce 梯度，缩放后的梯度在通信前由 DDP 内部处理（通过 hook），无需额外改动。

性能调优技巧

• 开启 cuDNN 的自动优化：torch.backends.cudnn.benchmark = True 可以让 cuDNN 在输入形状固定时选择最快的低精度卷积实现。
• 避免在 autocast 块内进行频繁的 CPU 操作或数据转换，减少与 GPU 的同步开销。
• 选择高效的优化器：对于 BF16 训练，推荐使用有 BF16 优化的优化器（如 apex.optimizers.FusedAdam、torch.optim.AdamW 在 CUDA 下的实现通常已足够），确保更新步骤尽量在 FP32 下进行。
• Profiling 确认加速效果：使用 torch.profiler 或 nvprof 确认 Tensor Cores 是否被激活，理论加速比可能因内存带宽、模型结构等打折扣。

总结

混合精度训练是现代大模型训练的必备手段，它将训练速度、显存效率和数值稳定性结合了起来。从笨重的手动 FP16 损失缩放到如今一键 autocast + GradScaler（或 BF16 无缩放），框架已经让实现门槛降至最低。根据你的 GPU 架构选择 FP16 或 BF16，遵循本文的流程和最佳实践，你可以在几乎不损失模型精度的前提下，获得显著的训练加速。

延伸阅读：

• NVIDIA 混合精度训练指南
• PyTorch AMP 官方文档
• bfloat16 浮点格式详解