奖励模型评分能有多快?对用于基于人类反馈的强化学习的C++和PyTorch推理运行时的系统研究
How Fast Can Reward Models Score? A Systems Study of C++ and PyTorch Inference Runtimes for RLHF
浏览论文内容
中文总结 AI 辅助
研究RLHF中奖励模型评分速度,构建基于ONNX Runtime的C++推理引擎,经测试,在CPU上击败所有基线,GPU上部分领先,发现加速源于ONNX Runtime,且批处理策略对速度影响更大,结果经多次独立运行得出。
中文摘要 AI 辅助
在基于人类反馈的强化学习(RLHF)管道中,奖励评分模块会推动策略更新。缓慢的评分成为整个循环的瓶颈,因为在每次 rollout 获得分数之前不会进行更新。然而,大多数设置仅默认使用PyTorch的即时模式或某些默认配置,没有人检查这是否真的是最快的。评分本身占用资源较少,rollout生成占用了RLHF步骤的更多资源。但评分和生成会争夺相同的CPU和GPU资源,因此更快的评分引擎本身不会缩短步骤时间,主要是释放生成可使用的容量。我们基于ONNX Runtime构建了一个原生C++推理引擎。首先确认正确性,在CPU上输出与PyTorch参考的误差为5.7 x 10^-6 ,在GPU上为4.2 x 10^-3 ,足够接近可信任。然后在CPU和GPU上与PyTorch即时模式、某些默认配置和FastAPI进行测试。CPU上我们的引擎击败了所有基线,置信区间不重叠。GPU上我们击败了PyTorch和FastAPI,但某些默认配置领先。进一步测试表明加速源于ONNX Runtime本身,而非C++语言。批处理策略比语言或运行时选择更重要。结果来自重复的独立运行,因为单次运行不够可靠。
英文摘要
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.