報酬モデルはどの程度高速にスコアリングできるか？：RLHFのためのC++およびPyTorch推論ランタイムのシステム研究

要旨

RLHFパイプラインでは、報酬スコアリングがポリシーの更新をブロックする。スコアリングが低速だとループ全体がボトルネックになる。なぜなら、全てのロールアウトのスコアが得られるまでは更新が実行されないからだ。それにもかかわらず、ほとんどの設定は単にPyTorchのeagerモードかtorch.compileをデフォルトとし、それが実際に最速かどうかは誰も確認していない。スコアリング自体は小規模な処理である。典型的なRLHFのステップではロールアウト生成の方がはるかに多くのリソースを消費する。しかし、スコアリングと生成は同一のCPUおよびGPUリソースを奪い合うため、スコアリングエンジンを高速化しても、それだけではステップ時間は短縮されない。主に、生成が代わりに利用できるキャパシティを解放する効果がある。我々はONNX Runtime上にネイティブC++推論エンジンを構築した。第一段階として、正確性を確認した。出力はPyTorchのリファレンスとCPU上で5.7×10^-6、GPU上で4.2×10^-3の誤差に一致し、信頼できる範囲であった。次に、CPUとGPUの両方でPyTorch eagerモード、torch.compile、FastAPIと比較テストを行った。CPUでは明確な差が出た。我々のエンジンはすべてのベースラインを上回り、信頼区間も重複しなかった。GPUでは異なる結果が得られた。PyTorchとFastAPIには勝ったが、torch.compileが優位に立った。さらなるテストにより、高速化の要因はC++という言語自体ではなく、ONNX Runtimeにあることが判明した。また、バッチ処理の方針が言語やランタイムの選択以上に重要であり、その影響は想定以上であった。これらの結果は、単一の実行では信頼性が不十分であるため、繰り返し独立した実行に基づいている。

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.