How Fast Can Reward Models Score? A Systems Study of C++ and PyTorch Inference Runtimes for RLHF

📄 arXiv: 2607.19712v1 📥 PDF

作者: Venkata Naga Sai Vishnu Rohit Pulipaka, Anish Katta, Deva Rohit Reddy Peddireddy

分类: cs.LG

发布日期: 2026-07-22

💡 一句话要点

提出C++推理引擎以提升RLHF中奖励模型评分速度

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 强化学习 奖励模型 C++推理引擎 ONNX Runtime 性能优化 批处理策略 机器学习

📋 核心要点

  1. 现有RLHF方法中，奖励评分的速度成为政策更新的瓶颈，影响整体效率。
  2. 论文提出了一种基于ONNX Runtime的C++推理引擎，以提高奖励模型的评分速度。
  3. 实验结果表明，CPU上新引擎超越所有基线，GPU上虽然torch.compile表现更佳，但整体性能提升显著。

📝 摘要（中文）

在RLHF管道中，奖励评分是政策更新的瓶颈，慢速评分会拖慢整个循环。尽管大多数设置默认使用PyTorch的急切模式或torch.compile，但并未验证其是否真的最快。我们构建了一个基于ONNX Runtime的原生C++推理引擎，首先确认其正确性。随后，我们在CPU和GPU上与PyTorch急切模式、torch.compile和FastAPI进行了测试。结果显示，CPU上我们的引擎超越了所有基线，而在GPU上，虽然超越了PyTorch和FastAPI，但torch.compile表现更佳。进一步测试表明，速度提升主要归因于ONNX Runtime，而批处理策略的重要性超出了我们的预期。

🔬 方法详解

问题定义：论文旨在解决RLHF管道中奖励评分速度慢的问题，现有方法如PyTorch急切模式和torch.compile未能有效优化评分过程，导致政策更新延迟。

核心思路：通过构建一个基于ONNX Runtime的原生C++推理引擎，论文希望提高奖励模型的评分效率，从而加速RLHF的整体流程。

技术框架：整体架构包括数据输入、推理引擎、结果输出等模块。首先确认引擎的正确性，然后在不同环境下进行性能测试，比较与现有方法的差异。

关键创新：最重要的创新在于使用ONNX Runtime作为推理引擎，显著提升了评分速度，并且发现批处理策略对性能的影响超过了语言和运行时的选择。

关键设计：在设计中，采用了高效的批处理策略，并在参数设置上进行了优化，以确保在CPU和GPU上都能获得最佳性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，CPU上新引擎的性能超越了所有基线，且置信区间没有重叠。在GPU上，虽然torch.compile表现更佳，但我们的引擎仍然优于PyTorch和FastAPI，验证了ONNX Runtime的有效性。

🎯 应用场景

该研究的潜在应用领域包括强化学习、自然语言处理和计算机视觉等领域，尤其是在需要快速反馈和迭代的场景中。通过提升奖励模型的评分速度，可以加速模型训练过程，提高整体效率，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.