优化奖励模型推理速度，用C+++ONNX提升训练效率
How Fast Can Reward Models Score? A Systems Study of C++ and PyTorch Inference Runtimes for RLHF
- 自研C++引擎基于ONNX Runtime，保证与PyTorch结果一致
- CPU上速度超越PyTorch和FastAPI，GPU上torch.compile仍领先
- 批处理策略比语言或运行时更重要，实测结果可靠
在强化学习人类反馈（RLHF）流程中，奖励评分是推动策略更新的关键环节。评分慢会拖慢整个训练循环，因为每次迭代必须等所有轨迹完成评分才能继续。然而当前多数系统默认使用PyTorch的急切模式或torch.compile，未验证是否最优。尽管评分本身计算量小，但其与轨迹生成争夺相同计算资源，因此更快的评分可释放更多资源供生成使用。我们构建了基于ONNX Runtime的原生C++推理引擎，首先验证其正确性：在CPU上输出误差为5.7×10⁻⁶，GPU上为4.2×10⁻³，与PyTorch参考结果一致，可信度高。随后在CPU和GPU上对比测试该引擎与PyTorch急切模式、torch.compile及FastAPI的表现。结果显示，CPU上本引擎全面胜出，置信区间无重叠；而GPU上虽优于PyTorch与FastAPI，但torch.compile仍最快。进一步分析表明，加速主要来自ONNX Runtime，而非语言本身。此外，批处理策略的影响远超语言或运行时选择，超出预期。所有结果均基于多次独立重复实验，确保可靠性。
原文摘要 · Abstract (English)
In RLHF pipelines, reward scoring blocks policy updates. Slow scoring bottlenecks the entire loop, since no update runs until every rollout gets a score. And yet most setups just default to PyTorch eager mode or torch.compile, no one checks if that's actually fastest. Scoring itself is small. Rollout generation eats far more of a typical RLHF step. But scoring and generation fight over the same CPU and GPU resources, so a faster scoring engine doesn't shrink step time on its own. It mainly frees up capacity generation can use instead. We built a native C++ inference engine on ONNX Runtime. First step: confirm correctness. Output matched the PyTorch reference to 5.7 x 10^-6 on CPU and 4.2 x 10^-3 on GPU, close enough to trust. Then we tested it against PyTorch eager mode, torch.compile, and FastAPI, on both CPU and GPU. CPU was decisive. Our engine beat every baseline, confidence intervals didn't even overlap. GPU gave a different view: we beat PyTorch and FastAPI, but torch.compile came out ahead. Further testing traced the speedup to ONNX Runtime itself, not C++ as a language. And batching strategy mattered more than either the language or the runtime choice, more than we expected. The results are from repeated, independent runs, since single runs just aren't reliable enough to trust.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。