奖励模型的得分速度有多快？ RLHF 的 C++ 和 PyTorch 推理运行时的系统研究

在 RLHF 管道中，奖励评分会阻止策略更新。缓慢的评分会成为整个循环的瓶颈，因为在每条采样轨迹都获得分数之前不会运行任何更新。然而大多数设置只是默认为 PyTorch eager 模式或 torch.compile，没有人检查这是否实际上是最快的。得分本身很小。轨迹生成消耗了更多典型的 RLHF 步骤。但评分和生成会争夺相同的 CPU 和 GPU 资源，因此更快的评分引擎本身并不会缩短步骤时间。它主要是释放产能来转用于轨迹生成。我们在 ONNX Runtime 上构建了本机 C++ 推理引擎。第一步：确认正确性。输出与 PyTorch 参考相匹配，CPU 上为 5.7 x 10^-6，GPU 上为 4.2 x 10^-3，足够接近，值得信赖。然后我们在 CPU 和 GPU 上针对 PyTorch eager 模式、torch.compile 和 FastAPI 进行了测试。 CPU是决定性的。我们的引擎超越了每条基线，置信区间甚至没有重叠。 GPU 给出了不同的观点：我们击败了 PyTorch 和 FastAPI，但 torch.compile 领先。进一步的测试将加速追溯到 ONNX 运行时本身，而不是 C++ 作为一种语言。批处理策略比语言或运行时选择更重要，超出了我们的预期。结果来自重复的独立运行，因为单次运行不够可靠，值得信赖。