獎勵模型評分能有多快？——針對RLHF之C++與PyTorch推理運行時的系統研究

摘要

在RLHF（基於人類回饋的強化學習）管線中，獎勵評分環節會阻礙策略更新。評分速度慢會拖累整個循環，因為必須等到所有生成結果都獲得評分後，才能執行更新。然而，多數設定仍直接採用PyTorch的即時模式或torch.compile，沒有人確認這是否真的最快。評分本身的運算量並不大，生成生成內容在典型的RLHF步驟中耗費的資源遠多於評分。但評分與生成會競爭同一批CPU與GPU資源，因此單純加快評分引擎並不能直接縮短步驟時間，主要效益是釋放出更多資源供生成環節使用。我們基於ONNX Runtime建構了一個原生C++推理引擎。第一步：驗證正確性。其輸出在CPU上與PyTorch參考結果的差異為5.7 × 10⁻⁶，在GPU上為4.2 × 10⁻³，精準度足以信賴。接著我們在CPU與GPU上分別與PyTorch即時模式、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.