MS-SWIFT项目中DPO训练时NLL_Loss为NaN问题的分析与解决

问题现象

在使用MS-SWIFT项目进行RLHF（强化学习人类反馈）训练时，特别是采用DPO（直接偏好优化）方法时，开发者遇到了NLL_Loss（负对数似然损失）持续为NaN的问题。该问题出现在使用RFT（强化微调）样本采样结合DPO训练的场景下，训练参数设置为：迭代次数4、epoch数1、batch_size为1、学习率1e-5，rpo_alpha采用默认值1.0。

问题排查过程

开发者首先检查了样本生成的数据格式，确认符合预期的['id', 'messages', 'rejected_response']结构。进一步调试发现，在训练过程中所有label值均为-100，这直接导致了NLL_Loss计算结果为NaN。

通过深入分析训练配置，发现开发者使用了以下关键参数组合：

• 模型类型：qwen2_5
• 最大长度：8192
• 训练类型：full（全参数训练）
• 精度：bfloat16
• 并行策略：deepspeed zero3
• 注意力实现：flash_attn
• 序列并行大小：4

根本原因

经过多次测试验证，最终确定问题根源在于sequence_parallel_size参数的设置。当该参数被移除后，NLL_Loss计算恢复正常。这表明在特定配置下，序列并行处理可能导致标签信息的异常处理或丢失。

此外，开发者还注意到之前为了解决另一个问题（模型推理中的张量设备不一致问题），修改了infer_engine.py文件中的代码，但确认这一修改与本问题无直接关联。

解决方案

针对这一问题，推荐采取以下解决方案：

1. 移除sequence_parallel_size参数：在当前的训练配置下，暂时不使用序列并行可以避免NLL_Loss异常的问题。

2. 数据完整性检查：虽然主要问题已定位到并行参数，但仍建议对训练数据进行全面检查，确保每条数据都包含完整的assistant响应部分，避免因数据缺失导致的其他潜在问题。

3. 梯度监控：在训练初期加入梯度监控机制，可以更早发现数值不稳定的情况。

经验总结

这一案例为使用MS-SWIFT进行RLHF训练提供了宝贵经验：

1. 并行策略的选择需要谨慎，特别是在结合特定模型架构和训练方法时，某些并行参数可能导致意外的数值计算问题。

2. 在损失函数出现NaN时，除了检查数据质量外，还应该考虑训练框架层面的配置参数影响。

3. 对于复杂的训练流程（如RFT+DPO），建议采用分阶段验证的方式，先确保基础训练正常，再逐步添加高级特性。

这一问题的解决为后续在MS-SWIFT框架下开展类似训练任务提供了重要参考，特别是在处理序列并行与损失计算交互时的注意事项。