这类问题我又碰上了一次。用 bf16 训练到两百多步时 loss 变成 NaN,前面几步 loss 还很正常,重启后同样的步数又复现。
第一反应是「损失函数里有除零或对零取对数」,验证下来不成立。
这类报错背后常见的原因有几处,逐个排除就行。学习率过高导致参数更新爆炸;数据里有超长样本,注意力计算出溢出;梯度累积时没有做归一化。
第一件事是把学习率降一个量级重跑
python train.py --learning_rate 5e-6
再往下开启梯度裁剪抑制异常更新
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
过滤超长样本,限制最大长度
python -c "print(max(len(x) for x in samples))"
末了在损失里加极小值避免除零
loss = -(y * torch.log(p + 1e-8)).mean()
排查用的命令都是临时敲的,没有写成脚本,事后才整理。
恢复的判定其实很直接。训练连续跑完三个 epoch 没有出现 NaN,loss 曲线平滑,验证指标正常。
bf16 的数值范围比 fp16 宽,出现 NaN 时更可能是学习率或数据问题,把精度当唯一原因往往会浪费很多时间。如果同一现象隔几天又出现,说明根因没找到,别当偶发处理。