简单记一次NaN loss的调试

发表于 更新于 分类于

LLM，神奇，太神奇了。

事情是这样的。

做LLM微调，同样一份代码，我重构了一下，loss就变NaN了。。。

第一次Loss NaN

选择的模型是Openbuddy-llama2-13b，使用shared-gpt数据进行微调。

现象是这样的，模型在第一个step的时候train loss是正常的，在第二个step开始就变为NaN了。

先进行了一些检查，排除了常见的一些原因：

  • 输入包含NaN

  • 权重包含NaN

  • 学习率过大

反复思考了一下，模型本身前向的过程应该是没问题的，从loss到回传梯度，到优化器优化中必有一处出现了NaN。默认的eps是1e-8。

最后，确定了Adam优化器中eps参数过小（1e-9）导致了NaN loss。

第二次Loss NaN

这一次直接第一个step就发生了loss为NaN的情况。这次定位原因更简单一些。

这一次是因为tokenizer padding side是在左侧，且pad token为eos，导致第一个token就是eos，后面的全部token都被mask掉了。