PyTorch混合精度训练是否有必要使用GradScaler

核心结论 #

在使用FP16精度开展混合精度训练时，GradScaler是必要的标准组件，不是可选优化项；如果使用BF16精度开展混合精度训练，则不需要配置GradScaler。

原理说明 #

梯度缩放解决的是FP16格式天生的数值范围缺陷：FP16可表示的最小正数值约为$6\times10^{-8}$，深度神经网络反向传播过程中，很多层的梯度幅值天然小于这个阈值，直接用FP16存储这些梯度会被直接截断为0（也就是文档提到的下溢问题），导致对应参数无法获得有效更新，最终模型收敛速度变慢、精度严重受损。

GradScaler的工作逻辑和NVIDIA Apex里的amp.scale_loss完全一致：

• 反向传播前先给loss乘一个动态调整的缩放因子，把整体梯度幅值抬升到FP16可正常表示的区间
• 优化器更新参数前，再把梯度除以相同的缩放因子，保证参数更新的步长和全精度训练一致，不会引入额外数值误差
• 训练过程中会自动检测梯度是否出现inf/nan，动态调整缩放因子大小，兼顾训练稳定性和数值精度

不同场景的使用规范 #

• 硬件不支持BF16（如NVIDIA图灵及更早架构显卡），只能用FP16跑混合精度：必须启用GradScaler，官方标准实现如下：

scaler = torch.cuda.amp.GradScaler()
for epoch in range(epochs):
    for input, target in zip(data, targets):
        opt.zero_grad()
        with torch.cuda.amp.autocast():
            output = net(input)
            loss = loss_fn(output, target)

        scaler.scale(loss).backward()
        scaler.step(opt)
        scaler.update()

• 硬件支持BF16（如NVIDIA安培及更新架构显卡、其他支持BF16的加速卡）：使用autocast时指定dtype=torch.bfloat16即可，BF16的动态数值范围和FP32完全一致，不存在小梯度下溢问题，不需要配置GradScaler，直接按常规训练流程执行反向传播、参数更新即可。
• Apex的梯度缩放实现是PyTorch原生AMP上线前的第三方解决方案，两者逻辑完全对齐，目前PyTorch原生AMP维护更活跃、兼容性更好，新项目优先选择原生实现即可，Apex的参考实现如下：

from apex import amp

model, optimizer = amp.initialize(model, optimizer)

loss = criterion(...)
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()
optimizer.step()

内容的提问来源于stack exchange，提问作者Deshwal