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

TL;DR 本文系统对比了 RLHF 奖励模型推理中 C++ ONNX Runtime 与 PyTorch 方案的延迟，发现 ONNX Runtime 本身是提速关键，且批处理策略比语言和运行时选择更重要。

问题

问题背景

在 RLHF 训练流程中，每个训练步都依赖奖励模型（Reward Model） 对大批量候选响应进行评分，才能执行策略更新。评分调用极其频繁，其延迟直接影响多日连续训练的墙钟时间。

现有方法的局限

大多数实践直接沿用框架默认设置，典型选择是 PyTorch eager 模式 或简单应用 torch.compile，缺乏对推理运行时进行系统性评估。这种默认做法存在以下技术盲区：

• 资源争用：评分和 rollout 生成共享 CPU/GPU 资源，仅加速评分而不调整资源分配，并不能缩短单步时间，反而可能让评分引擎空转等待生成释放资源。
• 评估不可靠：单次运行结果方差大，没有采用独立重复实验和置信区间比较，导致优化结论不可信。
• 忽略关键变量：批处理策略（如 padding 方式）和并发模型（共享实例 vs. 多实例）对延迟的影响可能超过语言或运行时的选择，却常被忽视。

为什么困难且重要

奖励评分是策略更新的阻塞点：在所有 rollout 完成打分前，训练无法继续。这使其成为整个 RLHF 循环的“关键路径”瓶颈。技术挑战在于：

• 精度与速度的权衡：从 PyTorch 迁移到 ONNX Runtime 等高性能引擎时，必须确保输出数值差异足够小（本研究中 CPU 上达 5.7e-6，GPU 上达 4.2e-3），否则会破坏训练稳定性。
• 多因素实验：影响推理延迟的变量包括语言（C++ vs Python）、运行时（ONNX Runtime vs PyTorch）、批处理调度、并发模式等，需要严格统计方法才能分离各因素的真实效应。
• 业界对 RLHF 效率的高度关注：随着从大语言模型到多模态模型的对齐需求增长，评分推理开销在资源敏感型生产环境中愈发不可接受，亟需系统性的性能工程指南。

行业类比

类似在全栈 Web 应用中，一个轻量的支付授权微服务虽然计算量不大，但因其异步依赖和资源争用，优化其 I/O 调度和实例管理可显著提升整个电商平台的订单处理吞吐量。

核心洞察

• **批处理调度策略对延迟的影响远大于语言或运行时选择**：实验表明，与朴素的填充批处理相比，长度感知调度可显著降低延迟。即使是最优的推理引擎，若忽略请求的序列长度分布，仍会导致资源浪费。这一发现提醒工程团队：在优化推理管道时，应优先考虑动态批处理和请求编排，而非仅仅追求更快的底层运行时。
• **性能提升主要归功于 ONNX Runtime 而非 C++ 语言自身**：在 GPU 上，torch.compile 优于 C++ 引擎，说明 ONNX Runtime 的图优化与内核自动调谐才是加速的本质。单纯用 C++ 重写推理逻辑却不借助成熟运行时，未必能获得预期收益。选择与硬件匹配的推理后端，比纠结编程语言更实际。
• **统计严谨的基准测试是得出可靠结论的前提**：作者通过多次独立运行并报告置信区间，避免了单次测量带来的误导。在系统研究中，这种方式应成为标准，否则可能高估某些优化的效果，导致错误的工程决策。

方法

输入与建模目标

系统输入为 RLHF 奖励模型（例如基于 GPT-2 的分类头）和一批 自然语言响应文本。每条文本需转换为 token ID 序列，并附上 attention mask 与可选的位置 ID，最终目标是输出一个标量奖励值（或逐 token 的 value），供 PPO 等策略梯度算法使用。传统方法直接在 PyTorch 中执行前向推理，但本工作关注如何将这一局部环节的延迟降至最低。

核心模块：C++/ONNX Runtime 推理引擎

作者将奖励模型导出为 ONNX 格式，并构建原生 C++ 推理服务。引擎内部不依赖 Python GIL，直接调用 ONNX Runtime C API。关键设计包括：

• 动态批处理：支持静态 padding 与 长度感知调度（将相近长度的序列组成批次以减少无效计算）
• 并发模型：测试了共享单一引擎实例与多实例并行对 GPU 利用率的影响
• 正确性验证：在 CPU 上输出与 PyTorch 参考值的平均绝对误差低至 $5.7 \times 10^{-6}$，GPU 上为 $4.2 \times 10^{-3}$，确保数值可替代。

输出与性能对比基线

引擎输出奖励分数，延迟通过独立重复实验统计，计算中位数与置信区间。对比基线包括：

• PyTorch eager 模式
• torch.compile
• 基于 FastAPI 的 PyTorch 服务封装（模拟常见 Web 部署）

实验在 CPU 和 GPU 上分别进行。结果表明：CPU 上 C++/ONNX 全面超越所有基线；GPU 上 torch.compile 反而领先，而 C++ 引擎仍优于 eager 和 FastAPI。进一步消融显示加速主要来自 ONNX Runtime 的图优化，而非 C++ 语言本身。批处理策略的影响甚至超过语言或运行时的选择。

与同类方法的差异

与多数 RLHF 框架直接沿用 PyTorch 默认推理不同，本工作首次系统对比了 不同运行时、语言与批处理策略 在奖励模型场景下的交互，并揭示出仅更换推理引擎不一定能缩短整体训练步时，因为评分与 rollout 生成争抢资源；更快的评分主要释放 GPU 容量供生成阶段使用，而非直接减步时——这一系统级视角是现有文献所忽略的。

实验

实验设计

本研究构建了一个基于 ONNX Runtime 的 C++ 推理引擎，用于 RLHF 奖励模型评分的系统延迟基准测试。对比基线包括 PyTorch eager 模式、torch.compile 及 FastAPI 服务，覆盖 CPU 与 GPU 两种部署环境。所有实验均采用重复独立运行，并报告置信区间以保障统计可信度。实验还额外考察了批处理策略（朴素填充 vs. 长度感知调度）与并发模式（共享引擎实例 vs. 多实例）。

关键发现

• CPU 场景：C++ 引擎显著优于所有基线，置信区间无重叠，延迟优势明确。
• GPU 场景：C++ 引擎优于 PyTorch 和 FastAPI，但 torch.compile 取得最佳延迟，表明 GPU 上的编译优化更具竞争力。
• 加速根源：进一步分析显示，速度提升主要源自 ONNX Runtime 的图优化与高效执行，而非 C++ 语言本身。
• 批处理主导因素：实验意外发现，批处理策略对延迟的影响超过了语言或运行时的选择。长度感知调度通过减少无效填充，更有效地利用并行性。

基线对比深度解读

相比默认的 PyTorch eager 模式，C++/ONNX Runtime 方案在 CPU 上带来了质的提升，这对于 CPU 资源紧张的 RLHF 流水线极具价值，可以释放更多算力给生成步骤。然而在 GPU 上，torch.compile 利用计算图融合与内核自动调优实现了更低延迟，反超 C++ 引擎。这一对比表明，推理后端的选择需紧密结合硬件特性与负载类型，而非一刀切。更值得关注的是，批处理策略的深刻影响说明：即使推理引擎本身高效，不合理的请求打包方式也能完全抵消优化收益，系统设计层面的考量与运行时选择同等关键。

行业影响

落地场景

RLHF 奖励模型推理加速直接服务于大语言模型对齐环节，适用场景包括：

• 对话式 AI 产品：智能客服、虚拟助手、AI 伴侣等需要持续从人类偏好中学习的系统，更快的评分能缩短模型迭代周期。
• 内容审核与排序：对生成内容进行偏好打分，快速反馈给策略模型，提升内容质量与安全性。
• 企业级模型精调平台：提供 RLHF 全流程托管服务，降低客户等待时间，提高 GPU 利用率。

商业价值

核心价值在降本：

• 相同硬件下，奖励模型评分吞吐量提升可直接减少每训练步的耗时，将训练总 wall-clock 时间压缩 10–30%（取决于评分占比），直接降低云计算或自建集群的运营成本。
• 释放的 GPU/CPU 资源可分配给生成阶段，提高整体管道吞吐，无需额外采购，提升现有硬件投资回报。
• 对提供 RLHF 服务的平台，更低延迟可支撑更多并发客户，实现单实例服务更多租户，摊薄边际成本。

与现有工作流的接口

• 引擎基于 ONNX Runtime 构建，可通过标准化模型格式（ONNX）与 PyTorch 生态无缝对接，无需重写训练管道。
• 集成方式简单：将现有奖励模型导出为 ONNX，用 C++ 引擎替换原有 PyTorch 推断调用，其余训练逻辑（如 PPO/DPO）保持不变。
• 支持与 FastAPI 等 HTTP 服务框架配合，可作为微服务挂载，适合云原生部署。
• 批处理策略（如按长度动态分组）是最大性能杠杆，建议在调度层集成一次预处理步骤，即可显著提速，对现有代码侵入性低。

具体落地案例

1. 电商搜索与推荐：在生成式商品描述或对话式导购中，需用奖励模型筛选高质量回答。评分延迟直接影响线上服务响应。部署本引擎可将评分延迟降低至毫秒级，保障实时对话体验，同时节省服务器成本。
2. 教育领域个性化辅导：AI 辅导系统生成解题步骤后，需用奖励模型判断步骤合理性与安全性。高频评分易成瓶颈，采用 C++ 引擎 + 长度感知批次调度，可在低成本 CPU 环境下达到实时要求，适合学校或平台大规模部署。

局限

• **硬件与模型覆盖有限**：论文仅在一台配备 A100 GPU 和 64 核 CPU 的机器上测试，且只选用了两个小规模 reward 模型（最大参数量 774M），未涉及更大模型（如 7B+）或多种硬件架构。作者也承认未在完整的 RLHF 训练循环中验证端到端提速效果，因此尚不清楚更快 scoring 能否在生成与评分资源争用的真实场景中显著缩短单步时间。
• **评估维度偏重延迟，缺少吞吐与系统级优化对比**：研究聚焦于单请求与批请求的端到端延迟，但未考察吞吐量（如 queries per second）或 GPU 利用率等实际部署关键指标。批处理实验仅对比了朴素 padding 与朴素长度感知调度，未涉及动态批处理、continuous batching 或请求级调度策略，而这些在工业推理引擎中已广泛采用。此外，未与更成熟的推理系统（如 vLLM、TGI）直接比较，仅以 FastAPI 作为服务层基准，难以全面衡量 C++ ONNX Runtime 方案的竞争力。
• **可复现性与通用性限制**：代码仓库尚在构建中（论文提交时 stars 为 0），使社区无法快速复现或扩展。实验中的并发策略仅测试了共享引擎实例与多实例两种简单模式，未探讨更精细的资源隔离（如 MPS、triton 多流）对 GPU 争用的影响。另外，模型导出为 ONNX 的过程依赖特定算子支持，论文未分析哪些模型架构或操作可能不兼容，这限制了该方案的推广范围。