奖励模型评分能有多快？——针对RLHF的C++与PyTorch推理运行时的系统研究

摘要

在RLHF流程中，奖励评分会阻塞策略更新。缓慢的评分成为整个循环的瓶颈，因为只有所有轨迹都获得评分后才能执行更新。然而，大多数配置只是默认使用PyTorch的eager模式或torch.compile，没人验证这究竟是不是最快的方案。评分本身计算量不大——典型RLHF步骤中，轨迹生成消耗要大得多。但评分和生成争用相同的CPU和GPU资源，因此更快的评分引擎本身并不会缩短每步时间，它主要释放出生成可以转而利用的计算能力。我们基于ONNX Runtime构建了一个原生C++推理引擎。第一步是确认正确性：输出与PyTorch参考结果在CPU上匹配误差为5.7×10⁻⁶，在GPU上为4.2×10⁻³，足够接近，可以信赖。随后我们在CPU和GPU上，与PyTorch eager模式、torch.compile以及FastAPI进行了对比测试。CPU结果是决定性的：我们的引擎击败了所有基线，置信区间甚至毫无重叠。GPU则呈现出不同图景：我们超越了PyTorch和FastAPI，但torch.compile表现更优。进一步测试发现，加速归功于ONNX Runtime本身，而非C++语言。而批处理策略的重要性超过了语言或运行时选择，甚至超出了我们的预期。上述结果均来自重复独立的多次运行，因为单次运行的可靠性不足以支撑结论。

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.