OpenRLHF项目中Qwen模型与PPO训练的不兼容问题分析

问题背景

在OpenRLHF项目中使用PPO算法训练Qwen模型时，发现了一个技术兼容性问题。该问题源于PPO训练过程中对序列的特殊处理方式与Qwen模型内部实现的冲突，特别是在使用flash_attention_2优化时的限制。

技术细节分析

PPO训练中的actor模块会对生成的序列进行特殊处理，具体表现在：

1. 序列处理逻辑：在生成序列后，会定位EOS(结束)标记的位置，并将该位置之后的所有标记都设置为EOS标记
2. 注意力掩码处理：同时会相应调整注意力掩码，确保模型不会关注EOS之后的无效标记

这种处理方式在大多数情况下工作正常，但当遇到以下条件组合时会出现问题：

• 使用Qwen模型
• 启用flash_attention_2优化
• 序列中存在右侧填充

Qwen模型的flash_attention_2实现中有一个严格的断言检查，要求当使用缓存(use_cache)时，如果注意力掩码存在右侧填充，则会抛出错误。这是出于对Flash Attention实现特性的考虑，确保计算正确性。

解决方案比较

目前可行的解决方案主要有两种：

1. 禁用flash_attention_2优化：

   • 优点：简单直接，能立即解决问题
   • 缺点：会降低训练速度，影响整体训练效率

2. 使用Llama3等其他兼容模型：

   • 优点：完全避免此问题，保持训练效率
   • 缺点：需要更换模型架构，可能不适合所有应用场景

技术建议

对于需要在Qwen模型上使用PPO训练的场景，建议：

1. 优先评估是否可以使用Llama3等兼容模型替代
2. 如果必须使用Qwen，则应权衡训练速度与模型性能，决定是否禁用flash_attention_2
3. 长期解决方案应考虑修改序列处理逻辑，使其与Qwen的flash_attention_2实现完全兼容

总结

这一兼容性问题揭示了深度学习框架中不同优化技术与模型架构间的微妙交互。开发者在组合使用先进优化技术(如flash_attention_2)和强化学习算法(如PPO)时，需要特别注意这种底层实现细节的兼容性。OpenRLHF项目中的这一案例为类似技术组合的应用提供了有价值的参考经验。