Hugging Face TRL RLOO 训练器发布

性能与效率提升

与 PPO 相比，RLOO 在资源利用率和训练速度方面提供了显著提升：

  • 内存效率： 根据模型大小，RLOO 使用的显存（vRAM）约比 PPO 少 50-70%。
  • 训练速度： 对于 1B 参数模型，RLOO 的运行速度是 PPO 的 2 倍；对于 6.9B 模型，最高可达 PPO 的 3 倍。
  • 模型质量： 在响应胜率（由 GPT-4 评判）方面，RLOO 与 PPO 竞争，并且始终优于离线方法，如 Direct Preference Optimization（DPO）。

技术架构与动机

传统的 PPO 资源消耗大，因为它需要在内存中加载四个模型副本：策略模型、参考策略模型、奖励模型和价值模型。RLOO 通过去除价值模型简化了该架构，只需三个副本：策略模型、参考策略模型和奖励模型。

关键算法差异

RLOO 在处理动作和奖励的方式上与 PPO 不同：

  1. 完成级别的动作： PPO 将每个生成的 token 视为单独的动作，而 RLOO 将整个模型的完成视为一个动作。这解决了稀疏奖励的问题，通常只有 EOS（句子结束）token 会获得真实奖励；RLOO 将 EOS 奖励归因于整个完成。
  2. REINFORCE 损失： RLOO 使用 REINFORCE 损失，即将奖励与基线的差值乘以动作的对数概率。这消除了 PPO 所需的价值模型和广义优势估计（GAE）。
  3. 留一基线（Leave-One-Out Baselines）： RLOO 通过使用批次中其他所有样本的奖励作为特定样本的基线来计算基线。对于给定的提示，一个完成的基线是该提示下所有其他完成的平均奖励。

在 TRL 中的实现

RLOO 训练器基于实验性的 PPOv2Trainer 构建。虽然 RLOO 在概念上基于 REINFORCE 算法，但 TRL 实现使用 PPO 损失，因为 REINFORCE 损失是 PPO 损失的特例。

基准测试结果

在 Pythia 1B 和 6.9B 模型上进行的实验展示了 RLOO 实现的有效性：

  • Pythia 6.9B： 使用 GPT-4 作为评判，获得 78.7% 的偏好率（k=2），超过原论文报告的 77.9%（k=4）和 74.2%（k=2）。
  • Pythia 1B： 获得 40.1% 的胜率，相比 SFT（监督微调）检查点的 21.3% 胜率有显著提升。

使用 bf16 时的数值稳定性挑战

在使用 bf16 精度时，RLOO 显示出显著的数值稳定性问题。生成过程中获得的对数概率可能与训练前向传播时获得的略有差异。

在 PPO 中，剪裁系数（通常为 0.2）会在比例超过 1.2 或低于 0.8 时使单个 token 的梯度归零。而在 RLOO 中，由于整个完成被视为一次动作，这些数值差异会被放大，导致整个序列的梯度更频繁被归零。经验观察表明，PPO 在 bf16 下约有 3% 的批次数据梯度被归零，而 RLOO 在 bf16 下则有 20%~40% 的批次数据梯度被归零。

Sources