Mixed Precision
混合精度训练使用 16 位运算和 FP32 主权重训练神经网络,从而减少内存占用和训练时间。内容涵盖损失缩放、BF16 和 AMP。
混合精度是一种训练技术,它以 16 位浮点格式执行神经网络的大部分计算,格式可以是 FP16 (半精度) 或 BF16,同时保留一份 32 位 (FP32) 的模型权重主副本,并以 FP32 执行对数值敏感的操作。16 位数值占用的内存和带宽只有 FP32 的一半,而现代张量核心在GPU上对这些数值进行乘法运算的速度快得多,因此混合精度可以缩短训练时间,并支持在内存中使用更大的批次大小,同时保持与 FP32 相当的准确率。2018 年,Micikevicius 等人描述的方案包含 FP32 主权重、FP16 运算和损失缩放;如今,深度学习框架会以自动混合精度 (AMP) 的形式自动执行这套方案。
混合精度的工作原理#
以下三个要素使 16 位训练中的反向传播和权重更新能够达到与 FP32 训练相当的准确率:
- FP32 主权重:优化器会将每次更新应用到 FP32 权重副本上。若直接将微小更新加到 FP16 权重上,通常会因 FP16 仅保留大约三位有效十进制数字而被舍入掉。
- 逐操作精度控制:autocast 上下文会以 16 位执行矩阵乘法和卷积运算(Tensor Cores 在这些运算中能带来最大加速),同时以 FP32 执行归约、softmax 和损失计算。
- 损失缩放:FP16 无法表示小于约 6 × 10⁻⁸ 的数值,因此许多较小的梯度会下溢为零。在反向传播之前将损失乘以缩放因子,可以把每个梯度移到 FP16 的表示范围内;权重更新之前,再将梯度除以相同的因子。
在 PyTorch AMP中,torch.autocast负责逐操作类型转换,GradScaler负责损失缩放:
import torch
from torch import nn
# Use FP16 on an NVIDIA GPU and BF16 on CPU
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.bfloat16
model = nn.Linear(512, 10).to(device) # master weights stay in FP32
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
scaler = torch.amp.GradScaler(device, enabled=dtype == torch.float16) # loss scaling for FP16 only
x = torch.randn(64, 512, device=device)
with torch.autocast(device_type=device, dtype=dtype):
out = model(x) # the matrix multiply runs in 16-bit
loss = out.float().pow(2).mean()
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
print(f"Output: {out.dtype}, weights: {model.weight.dtype}, gradients: {model.weight.grad.dtype}")在 CPU 上,这段代码会为层输出打印 torch.bfloat16,为权重及其梯度打印 torch.float32;在 NVIDIA GPU 上则会打印 torch.float16。在一次训练步骤中,16 位运算和 FP32 主权重会同时存在。
PyTorch 的 GradScaler会动态调整缩放因子。它从 65,536 开始;当缩放后的梯度包含 inf 或 NaN 值时,就将缩放因子减半并跳过该步骤;连续 2,000 步没有溢出后,则将缩放因子加倍,因此该因子最终会稳定在梯度溢出临界值略低的位置。
FP16、BF16 与更低精度#
FP16 用 5 位表示指数、10 位表示尾数,因此精度较高,但数值超过 65,504 时就会溢出。BF16 保留了 FP32 的 8 个指数位,数值范围相同,最高约为 3.4 × 10³⁸,但只有 7 个尾数位,因此无需损失缩放,代价是舍入更粗糙。NVIDIA Ampere 及后续架构的 GPU、Google TPU以及许多较新的 CPU 都支持 BF16;它也是大语言模型常用的训练格式。同一思路还通过 NVIDIA 的 Transformer Engine扩展到了 Hopper 和 Blackwell GPU 上的 FP8,以及 Blackwell 上的 NVFP4 等 FP4格式。
混合精度实践#
使用 AMP 的主要实际风险是模型或 GPU 在 16 位下产生 NaN 损失。Ultralytics YOLO26 训练默认启用 AMP,使用 amp=True,首先在 CUDA GPU 上比较 YOLO26n 模型的 FP32 与 AMP 预测;如果两者不一致,或者 GPU 属于已知会在 AMP 下产生 NaN 损失的系列(例如 GTX 16 系列),则回退到 FP32。设置 amp="bf16" 会切换到 BF16 autocast,并关闭损失缩放。
只要训练受内存或时间限制,混合精度就能广泛派上用场。开发自动驾驶汽车的团队会在数百万张高分辨率画面上训练目标检测模型;激活值内存减半后,就能使用更大的图像和批次。在医学图像分析中,3D CT 和 MRI 体数据通常无法以 FP32 放入 GPU 内存,而 16 位激活值使得在完整体数据上训练分割模型成为可能。










