PyTorch 开启 FP16 混合精度后，loss 变成 NaN

同一训练配置在 FP32 下运行正常，开启 FP16 自动混合精度后，loss 很快变成 NaN，或者梯度出现 inf。先区分非有限值出现于前向损失、反向梯度还是参数更新，不能把所有情况都归结为“忘了加 GradScaler”。

本文按 PyTorch 2.8 的 CUDA AMP 接口说明，示例是普通单模型、单优化器训练。梯度累积、多优化器或分布式训练需要按自己的更新边界组织缩放。

先判断坏值发生在哪一步

用同一份受控小批次检查输入、模型输出与损失是否有限，再观察反向后的梯度。不要只在最后把 NaN 替换成零，那会掩盖产生坏值的运算。

对已有的 Tensor，可用下面的诊断写法：

print("input finite:", torch.isfinite(inputs).all().item())
print("output finite:", torch.isfinite(outputs).all().item())
print("loss finite:", torch.isfinite(loss).all().item())

这段要求 inputs、outputs、loss 是实际 Tensor；嵌套模型输出需要检查具体张量。它用于缩小故障范围，不建议对完整训练的每个中间 Tensor 都同步读回。

如果 FP32 也出现坏值，继续检查数据、除零、对数输入、标签、学习率及训练配置。FP16 的有限数值范围更小，可能更早暴露同一数值问题；“改回 FP32 后暂时没坏”也不等于数据和公式一定正确。

autocast 与缩放按各自范围使用

下面是局部训练循环。model 保持默认 FP32 参数，已经放在 CUDA；optimizer、loss_fn 和数据加载器已有定义，输入与目标设备正确：

import torch

scaler = torch.amp.GradScaler("cuda")

for inputs, targets in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast(device_type="cuda", dtype=torch.float16):
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

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

autocast 覆盖前向和损失计算，不建议把反向也包在里面。开启 autocast 后，不要再为“配合 AMP”无条件对模型和所有输入调用 .half()；它会改变参数与更多运算的数据类型，偏离上面的用法。

GradScaler 主要缓解小梯度下溢。scaler.step() 在梯度有非有限值时会跳过相应优化器更新，随后由 scaler.update() 按检查结果调整缩放。缩放不能修复前向已经产生的 NaN。也不要在这段之外再调用一次 optimizer.step()，绕过跳过更新的判断。

需要裁剪梯度时，在 scaler.scale(loss).backward() 后先执行一次 scaler.unscale_(optimizer)，再裁剪，然后交给 scaler.step()。同一个优化器同一更新周期不要重复 unscale。梯度累积时，要等本次更新所需微批全部完成，再取消缩放、更新和调整 scale。

BF16 预训练模型不一定能直接改成 FP16

FP16 最大有限值为 65504。部分 BF16 预训练模型不能在 FP16 的数值范围内正常运行，可能产生溢出；GradScaler 的 scale 甚至可能降到 1 以下，这本身不违反接口保证。

若模型明确按 BF16 训练或推荐 BF16，先核对模型支持的精度与 GPU、PyTorch 运算支持，再选择 BF16 或 FP32。BF16 的指数范围与 FP32 接近，但有效精度仍有限，不能承诺它消除所有 NaN。硬件不支持所需 BF16 路径时，优先用可支持的精度建立正确性基线，不直接强行转换。

CUDA BF16 的常规 AMP 用法通常不需要像 FP16 那样做梯度缩放；是否使用仍应服从具体训练实现。不要把 FP16、BF16 和把整个模型显式转换成低精度看成等价设置。

局部运算确需 FP32 时，进入前先转换输入

定位到某个运算在低精度下不稳定后，可以在 autocast 区域内部对那个局部禁用 autocast，并把来自低精度区域的浮点输入显式转为 FP32。仅写 enabled=False 不会自动升级已经是 FP16 的 Tensor；前面已经发生的溢出，也不会被 .float() 还原。

PyTorch 的 CUDA autocast 已让一部分损失、归一化和指数类运算采用 FP32。不要对所有层重复强制转换，或认为某层名字中有 softmax 就一定是 FP16。应查看实际 dtype 与最早的坏值位置，再做范围明确的改动。

保留原配置和检查点，在相同小批次下比较 FP32 与修正后的 AMP：输出、损失和梯度有限，参数确实按预期更新，才恢复完整训练。一个批次通过不能证明长期训练收敛；后续仍需记录损失、缩放变化和跳过更新的情况。