PyTorch从1.6版本开始引入的混合精度训练工具torch.cuda.amp，让模型在几乎不损失精度的情况下获得显著的显存节省和速度提升。然而很多人在实际使用时都会遇到一个棘手的问题：GradScaler的scale值不断波动，反向传播过程中频繁出现梯度溢出警告，训练loss不下降甚至直接变成NaN。这个现象的本质是loss scale与梯度数值范围之间的动态平衡被打破，需要从缩放机制和参数调整两个层面来系统解决。
理解GradScaler的缩放与自动回退机制
混合精度训练的核心思路是在前向传播时使用FP16或BF16存储激活值和权重，反向传播时梯度也以低精度计算，从而减少显存占用和提升吞吐。问题在于FP16的数值表示范围非常有限，最小正规数约为6.1e-5，最大约为65504。当梯度值小于最小正规数时会被下溢为0，当梯度值大于65504时则会上溢为inf。loss scale的作用就是在反向传播之前将loss乘以一个较大的系数，使原本很小的梯度值放大到FP16可表示的范围内，避免下溢。
GradScaler内部维护一个当前的scale值，每次调用scaler.scale(loss).backward()时，它会将loss乘以该scale，然后正常进行反向传播。反向传播完成后，执行scaler.step(optimizer)之前，GradScaler会检查本次迭代的梯度中是否存在inf或NaN。如果没有发现溢出，则正常进行优化器更新，并按照growth_interval的设定决定是否增大scale值；如果发现了溢出，则跳过本次参数更新，同时将scale值乘以backoff_factor进行缩小。这个自动检查回退的机制保证了即使scale设置得过大导致溢出，也不会破坏模型参数，只是浪费了一次前反向计算。
这里需要明确一个关键细节：GradScaler的scale是全局状态，针对所有参数统一缩放。当模型包含某些对数值范围特别敏感的网络层时，例如带有指数运算的自注意力分数或者softmax之前的logits，即使整体梯度没有溢出，局部梯度可能已经接近FP16的边界。GradScaler只能检测最终梯度中的inf和NaN，无法定位具体是哪一层产生的溢出，因此调试时需要结合梯度统计信息来定位问题层。
FP16与BF16的溢出差异及Loss Scale选择
NVIDIA从Ampere架构开始对BF16提供了硬件支持，PyTorch中可通过autocast(dtype=torch.bfloat16)启用。BF16与FP16最大的区别在于指数位宽：FP16有5位指数，而BF16有8位指数，与FP32相同。这意味着BF16的数值表示范围与FP32几乎一致，可以表示到约3.4e38，完全不会出现上溢问题，只有下溢风险。因此使用BF16时，通常不需要较大的loss scale，甚至可以将scale固定为1。这也是为什么在训练大语言模型时，很多框架默认使用BF16加constant loss scale，而不会出现GradScaler频繁回退的情况。
对于FP16混合精度训练，初始scale的选取非常关键。如果初始scale设置过小，训练初期梯度下溢严重，模型收敛慢；如果初始scale设置过大，前几个迭代就会出现梯度溢出，触发回退，浪费计算资源。PyTorch默认的初始scale为65536.0，这个值对于大多数卷积网络和中等规模的Transformer是可行的。但对于一些梯度本身就比较大的场景，比如使用了大学习率、梯度裁剪阈值设置过高、或者模型包含大量未归一化的残差连接时，65536的初始值很容易导致第一次反向传播就溢出。此时可以将初始scale降低到2048或1024，让GradScaler从较小的值开始动态探索。
动态调整中的两个核心参数是growth_interval和backoff_factor。growth_interval表示连续多少个迭代没有溢出才将scale翻倍，默认值为2000。这个默认值对于大规模数据集的训练来说比较保守，在训练初期scale增长非常缓慢，可能导致很长一段时间内梯度都处于下溢状态。backoff_factor表示溢出时scale缩小的比例，默认值为0.5。当训练中出现偶发溢出时，scale折半后基本能恢复；但如果模型梯度极不稳定，连续溢出会导致scale指数级下降，很快变成很小的数，此时loss scale就失去了防止下溢的作用。针对训练不稳定问题，建议将backoff_factor设置得稍微大一些，例如0.7或0.8，让scale回退时衰减更加温和。
溢出排查流程与手动Loss Scale调优实践
第一步是确认溢出发生的频率和位置。可以在每个optimizer.step之后记录scaler.get_scale()的返回值，同时统计连续溢出的次数。下面这段代码展示了如何在训练循环中监控scale变化和梯度统计信息：
import torch
from torch.cuda.amp import GradScaler, autocast
scaler = GradScaler(init_scale=4096, growth_interval=100, backoff_factor=0.7)
model = MyModel().cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.cuda(), target.cuda()
optimizer.zero_grad()
with autocast(dtype=torch.float16):
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
# 手动检查梯度是否包含inf或NaN（不经过scaler内部检查）
total_norm = 0.0
inf_count = 0
for p in model.parameters():
if p.grad is not None:
grad = p.grad.float()
if torch.isinf(grad).any() or torch.isnan(grad).any():
inf_count += 1
total_norm += grad.norm(2).item() ** 2
total_norm = total_norm ** 0.5
# 记录当前scale和梯度范数
if batch_idx % 50 == 0:
print(f"batch {batch_idx}, scale={scaler.get_scale():.1f}, "
f"grad_norm={total_norm:.4f}, inf_layers={inf_count}")
# 先手动取消这次step，改用无条件缩放更新测试
if inf_count == 0:
scaler.step(optimizer)
scaler.update()
else:
print(f"Overflow detected, skipping step. Current scale={scaler.get_scale():.1f}")
scaler.update() # 只更新scale，不执行step
这段代码的关键意义在于：它绕过了GradScaler内部自动跳过step的逻辑，而是先检查梯度中是否存在inf或NaN，再决定是否执行优化器更新。这样做的目的是在调试阶段获得更多信息，知道哪些层产生了溢出。当inf_count大于0时，输出对应的层名和梯度范数，可以快速定位问题。定位后，可以针对该层使用更宽的数值类型，比如对该层的输入输出单独转换为FP32，或者在autocast上下文中排除该模块。
第二步是根据监控结果调整scale参数。如果训练初期scale持续增长并在一定范围后出现偶发溢出，且溢出后scale回退后训练能继续稳定进行，说明当前设置基本合理，只需要适当增大growth_interval或者降低学习率即可。如果scale一直在低位徘徊，频繁触发溢出回退，说明模型梯度数值本身可能过大，此时单纯调整loss scale无法解决根本问题，需要检查模型初始化、归一化层配置以及是否遗漏了梯度裁剪。相反，如果scale增长到上限（通常是2的24次方左右）且从未出现溢出，但训练loss下降缓慢甚至不收敛，则可能是loss scale过大导致反向传播中大量梯度被clamp到FP16的最大值，失去了梯度方向信息，此时可以尝试降低初始scale或者手动缩小scale。
第三步是针对特定模型实施精细控制。GradScaler支持对部分参数禁用缩放，通过register_full_backward_hook或者直接控制梯度的方式实现。例如在训练包含语言模型头的网络时，logits层之前的隐藏状态数值可能很大，将其转换为FP16后softmax的输入溢出，此时可以在模型前向传播中手动将该部分转换为FP32再计算。另一种更直接的方法是修改GradScaler的缩放策略，在每次调用scaler.scale前根据当前batch的loss动态调整scale值，但这种方式会破坏GradScaler自动增长机制，更推荐的做法是保持自动机制并配合梯度裁剪：torch.nn.utils.clip_grad_norm_手动执行后，再手动除以scale以抵消GradScaler的缩放影响。
进阶方案：恒定Scale与逐层Scale的取舍
对于已经完成调优的模型，动态GradScaler在推理阶段部署时会引入额外的控制流开销，而且在多卡训练中不同device上的scale同步需要额外通信。如果训练达到稳定阶段，可以考虑将GradScaler设置为恒定模式，即初始化后禁用自动增长和回退。PyTorch中可以通过继承GradScaler并重写update方法实现，或者直接使用enabled=False关闭autocast的缩放功能，手动对loss进行缩放。恒定scale的优势是代码逻辑简单，在多GPU环境下行为一致，缺点是必须人工找到合适的scale值，否则容易出现过拟合或下溢。
逐层Scale是另一种更精细的方案，针对不同网络层使用不同的缩放因子。这需要对每个参数单独应用scale和unscale操作，PyTorch没有直接提供官方API，但可以通过手动缩放梯度来实现：在backward之后，对每一层的梯度乘以不同的系数后再调用optimizer.step。这种方案适用于网络深度极大且各层梯度数量级差异明显的场景，例如视觉Transformer中浅层patch embedding的梯度与深层attention block的梯度可能相差几个数量级。不过实现复杂度较高，且容易出现微分计算错误，一般只有在标准GradScaler反复调优无效后才考虑使用。
最后要强调的是，虽然BF16混合精度训练基本不受梯度上溢影响，但在某些GPU架构（如V100、T4等）上BF16的计算速度反而比FP16慢，因为硬件不支持原生BF16运算。因此在实际项目中选择FP16还是BF16，需要综合考虑硬件支持、模型梯度范围以及调参成本。对于需要快速迭代且不追求极致精度的场景，直接使用BF16加constant loss scale是最稳妥的方案；对于需要充分利用旧款GPU计算能力的场景，则必须针对FP16的特点仔细调整GradScaler参数，才能既保证训练稳定又获得性能收益。
混合精度训练Loss Scale调整GradScaler修改时间：2026-08-26 22:23:17