最近在跑一个 7B 左右的模型做指令微调，用 PyTorch 2.1 + AMP 混合精度，单机 8 卡 A100。训练到 300 多 step 的时候 loss 突然变成 NaN，前面都挺正常的。我试过调小学习率、加 gradient clipping，但还是会偶发。看有人说是 fp16 溢出，建议换 bf16，也有人说要检查 loss scaling 那块。我现在有点懵，不知道该从哪排查，有没有遇到过类似情况的老哥指点一下？框架层面到底该怎么定位这种问题？
楼主
10小时前
PyTorch 2.x 训练大模型时 loss 突然 NaN，是混合精度的问题吗？
请 登录 后发表回复
全部回复
共 1 条
2楼
7小时前
300多step才炸，大概率不是单纯fp16溢出，更像是某个batch的梯度异常把scale撑爆了。建议先把GradScaler的unscale后的梯度norm打出来看看，定位到具体是哪个参数层的梯度炸了。bf16确实值得换，A100原生支持，省掉loss scaling这层就少一个坑，但换之前最好先确认下NaN到底出在前向还是反向。我之前遇到过类似的是embedding层附近数值不稳，加了点weight decay和更保守的warmup就压下去了。