OpenRLHF项目中Qwen2模型KTO训练Loss异常问题分析与解决

问题背景

在使用OpenRLHF项目对Qwen2 7B模型进行KTO(Knowledge Transfer Optimization)训练时，开发者遇到了Loss值变为NaN的问题。该问题出现在使用业务数据集进行SFT(Supervised Fine-Tuning)后，继续使用KTO训练时。值得注意的是，相同的训练配置在Qwen1.5模型上表现正常，这表明问题可能与Qwen2模型特性相关。

技术配置分析

训练环境配置如下：

• 使用8个GPU并行训练
• 微训练批次大小为2
• 最大序列长度8192
• 学习率5e-7
• 启用了梯度检查点和vanilla_loss选项
• 使用了DeepSpeed Zero Stage 3优化
• 软件环境包括PyTorch 2.2.2、DeepSpeed 0.14.3和Flash Attention 2.4.2

问题排查过程

经过技术分析，发现导致Loss异常的主要原因与vanilla_loss选项的使用有关。vanilla_loss是KTO训练中的一个参数选项，它控制是否使用原始版本的损失函数计算方式。

在Qwen2模型架构下，当启用vanilla_loss时，损失计算过程中可能出现数值不稳定的情况，导致梯度爆炸或NaN值的产生。这与Qwen2特定的模型结构和参数初始化方式有关，而Qwen1.5由于架构差异，对相同的损失计算方式有更好的数值稳定性。

解决方案

解决此问题的直接方法是禁用vanilla_loss选项。在OpenRLHF项目中，这可以通过以下方式实现：

1. 在训练命令中移除--vanilla_loss参数
2. 或者显式设置--vanilla_loss false

此外，作为额外的稳定性保障措施，可以考虑：

• 适当降低学习率
• 增加梯度裁剪阈值
• 使用混合精度训练时检查精度设置

技术原理深入

KTO训练中的损失函数计算涉及模型输出的logits值比较，当使用vanilla_loss时，可能会在某些模型架构下导致数值范围超出预期。Qwen2的注意力机制和前馈网络结构可能使得某些中间值变得过大，在后续的指数运算中产生数值溢出。

现代大语言模型训练中，数值稳定性是一个常见挑战。不同模型家族(如Qwen1.5和Qwen2)可能在参数初始化、层归一化位置和激活函数选择等方面存在差异，这些都会影响训练过程的数值行为。

最佳实践建议

针对OpenRLHF项目中的KTO训练，建议：

1. 对于Qwen2系列模型，避免使用vanilla_loss选项
2. 训练初期密切监控Loss曲线和梯度范数
3. 考虑使用更稳定的损失函数变体
4. 对于新模型架构，先进行小规模试验性训练

通过这些问题排查和解决过程，开发者可以更深入地理解不同LLM模型在RLHF训练中的行为差异，为后续的模型调优积累宝贵经验。