XTuner训练过程中RuntimeError问题的分析与解决

问题现象

在使用XTuner项目训练InternLM2模型时，用户遇到了一个典型的训练不稳定性问题：模型在前100多个迭代步骤中loss正常下降，但突然出现NaN值，并伴随"RuntimeError: probability tensor contains either inf, nan or element < 0"的错误提示。

可能原因分析

1. 浮点精度问题：使用float16(fp16)进行训练时，数值范围较小(约6e-5到65504)，容易出现数值下溢或上溢，导致NaN的出现。

2. 学习率设置不当：过大的学习率可能导致参数更新幅度过大，使模型参数进入不稳定的数值区域。

3. 数据质量问题：训练数据中包含异常值或不符合预期的输入，可能导致模型计算出问题。

4. 梯度爆炸：在训练过程中梯度突然增大，导致参数更新后产生无效数值。

解决方案

1. 调整学习率：适当降低学习率是解决训练不稳定问题的首选方案。可以先尝试将学习率减半，观察训练过程是否稳定。

2. 使用bf16替代fp16：如果硬件支持(可通过torch.cuda.is_bf16_supported()检查)，bf16具有与fp32相同的数值范围(约1e-38到3e38)，但精度与fp16相当，能有效减少数值溢出问题。使用bf16时需要配合deepspeed_zero2配置。

3. 梯度裁剪：实现梯度裁剪可以防止梯度爆炸问题，保持训练过程的稳定性。

4. 数据检查：仔细检查训练数据集，确保输入数据在合理范围内，没有异常值或格式错误。

5. 混合精度训练：可以考虑使用自动混合精度(AMP)训练，让框架自动管理不同层的精度，平衡训练速度和稳定性。

实践经验

在实际训练过程中，用户发现重新开始训练后问题没有复现，这表明问题可能具有一定的随机性。这种间歇性出现的问题通常与数值稳定性相关，而非确定性的数据或模型结构问题。

对于大模型训练，建议：

• 在训练初期密切监控loss变化
• 设置合理的checkpoint保存频率
• 使用学习率warmup策略
• 考虑使用更稳定的优化器如AdamW

结论

深度学习模型训练过程中的NaN问题通常与数值稳定性相关。通过调整训练参数、使用更合适的数值精度以及仔细检查数据，可以有效解决这类问题。XTuner作为训练框架，支持多种精度和优化策略，用户可以根据具体硬件条件和任务需求选择最适合的配置方案。