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

Analysis

Why This Paper Matters

Reinforcement Learning from Human Feedback (RLHF) is the backbone of aligning large language models, but its training loop is notoriously slow. The reward model scoring step is a critical bottleneck because policy updates cannot proceed until all rollouts are scored. While many practitioners assume PyTorch eager mode or torch.compile is sufficient, this paper systematically investigates whether a custom C++ inference engine on ONNX Runtime can accelerate scoring. The findings are practically important: on CPU, the C++ engine significantly outperforms all baselines, but on GPU, torch.compile remains faster. This nuanced result challenges the default assumption that a single runtime is optimal for all hardware.

Moreover, the paper reveals that batching strategy—how many rollouts are scored together—has a larger impact on throughput than either the programming language or the inference runtime. This insight is actionable for RLHF engineers: tuning batch sizes may yield bigger gains than switching frameworks.

Technical Contributions

• C++ ONNX Runtime Engine: A native C++ inference engine for reward scoring, validated against PyTorch with numerical errors of 5.7e-6 (CPU) and 4.2e-3 (GPU).
• Comprehensive Benchmarking: Compared against PyTorch eager mode, torch.compile, and FastAPI on both CPU and GPU, with repeated independent runs for statistical rigor.
• Ablation on Runtime vs Language: Isolated the effect of ONNX Runtime versus C++ by comparing against PyTorch-based baselines, showing that ONNX Runtime (not C++) drives the CPU speedup.
• Batching Strategy Analysis: Demonstrated that batching strategy has a larger effect on throughput than language or runtime choice.

Results

• CPU: The C++ engine beat all baselines (PyTorch eager, torch.compile, FastAPI) with non-overlapping confidence intervals, indicating a statistically significant speedup.
• GPU: torch.compile outperformed the C++ engine, while the C++ engine still beat PyTorch eager and FastAPI.
• Key Insight: The speedup on CPU is due to ONNX Runtime, not C++ as a language. Batching strategy was the most influential factor overall.

Significance

This paper provides concrete guidance for RLHF practitioners: for CPU-based reward scoring, ONNX Runtime with C++ is the fastest option; for GPU, torch.compile is superior. More importantly, it emphasizes that batching strategy should be optimized first, as it dominates performance. The work also highlights that resource contention between scoring and generation means faster scoring primarily frees capacity for generation, rather than directly reducing step time. These findings can help reduce RLHF training time and improve resource utilization in production systems.