No More Train-Inference Mismatch: Bitwise Consistent On-Policy Reinforcement Learning with vLLM and TorchTitan

发表时间: 2025-11 · Blog post by vLLM and TorchTitan Teams (blog.vllm.ai)

原文: https://blog.vllm.ai/blog/2025-11-10-bitwise-consistent-train-inference

作者/机构：Bram Wasti, Wentao Ye, Teja Rao, Michael Goin, Paul Zhang, Tianyu Liu, Natalia Gimelshein, Woosuk Kwon, Kaichao You, Zhuohan Li (vLLM and TorchTitan Teams)

主要贡献

强化学习（RL）在训练过程中会放大训练器（trainer）和采样器（sampler）之间微小的数值不匹配，这种不匹配会导致非确定性以及不稳定的训练行为。本文旨在解决这一长期存在的训练-推理不匹配问题。核心贡献在于，作者展示了一个开源的、位级一致（bitwise consistent）的同策略（on-policy）强化学习运行方案，该方案采用 TorchTitan 作为训练引擎，vLLM 作为推理引擎。基于 vLLM 近期在批次不变推理（batch-invariant inference）方面的技术积累，作者成功在 Qwen3 1.7B 模型上运行了 RL 微调，并在开源的指令中实现了训练和推理数值的位级完全匹配。

背景知识与关键观察

强化学习对数值不匹配的敏感性：以往的研究证实，强化学习会显著放大训练器和采样器之间极其微小的数值差异，这直接导致了非确定性和不稳定的训练行为【1，Defeating nondeterminism in LLM inference，He 等人】【2，Off-policy RL，Yao, Liu 等人】【3，When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch，Liu, Li 等人】。

框架间内核差异的根源：训练和推理框架由于底层工作负载特性的根本不同，通常会使用截然不同的计算内核（kernels）。即使在单一的推理框架内部，也会针对不同的应用场景选择不同的内核：例如，针对大批次（high batch sizes）的内核会在批次维度上进行高度并行化；而针对小批次（low batch sizes）的内核则更多地在单个实例内部进行并行化，以便在 GPU 的并行核心上获得更好的资源利用率。所有这些底层实现上的差异都会导致训练和推理框架之间产生数值计算差异，并最终导致劣质的强化学习结果。

方法细节

• 跨框架位级等价的审计与实现：在这项工作中，作者着手解决了两个不同框架之间的不变性（invariance）问题，即作为训练框架的 TorchTitan 和作为推理框架的 vLLM。为了实现这一目标，作者在正向传播（forward pass）过程中，对每一个计算内核的每一次调用都进行了详尽的审计，以确保它们在两个框架之间是绝对位级等价（bitwise equivalent）的。在具体实现上，作者直接利用了 vLLM 近期关于批次不变性（batch invariance）工作中的正向传播内核，并针对这些操作专门编写了简单的反向传播（backward passes）。
• 高度优化融合操作的导入与反向注册：vLLM 内部拥有许多经过高度优化的融合操作（fused operations），例如 SiLU MLPs 以及带有附加残差连接的 RMSNorms。为了在整个计算流中保持位级等价性，作者将这些精确的正向传播操作完整导入到了训练流中。由于这些是自定义操作，因此需要为它们注册自定义的反向传播路径，而这一注册与实现过程直接在编写 TorchTitan 所使用的原生 PyTorch（vanilla PyTorch）环境中即可完成。
• 同策略强化学习演示架构设计：为了进行强化学习的演示，作者利用 GSM8K 数据集和一个正确性奖励（correctness reward）编写了一个通用的强化学习脚本。在架构搭建上，作者使用了 TorchTitan 提供的实用工具来构建训练器（trainer），并自行编写了一个自定义的生成器（generator）。该生成器被命名为 VLLMRolloutEngine，它封装了包括调用生成（generate）和更新权重（updating weights）在内的基础功能。在运行机制上，系统在单一主机上同步运行所有组件，在训练器和生成器之间严格交替执行。这种设计精确地演示了同策略（on-policy）的执行逻辑，尽管这种完全同步的方式在业界的大规模运行中并不常见。

实验环境

• 模型架构：Qwen3 1.7B。
• 数据集：GSM8K（用于计算正确性奖励）。
• 软件配置：训练引擎使用 TorchTitan，推理引擎使用 vLLM；底层基于原生 PyTorch 实现；未启用 torch.compile，vLLM 强制运行在 eager 模式下。
• 硬件与执行模式：单一主机（single host），采用训练器与生成器交替执行的同步模式。

实验结果

作者通过对比实验验证了数值精度对强化学习结果的实质性影响：

• 非位级一致（batch_inv_OFF）实验：当采样器使用与训练器不同的内核运行时，实验发现在 $100$ 个步骤（steps）之后，模型获得的奖励（reward）出现了明显的下降。
• 位级精确一致（batch_inv_ON）实验：当启用位级精确训练时（此时散度 $kl\_div$ 始终严格等于 $0.0$），实验结果表明，模型不仅能够在更少的步骤内完成训练收敛，而且最终达到了更高的总奖励（total reward）。
• 性能开销：当前的位级一致 RL 运行速度比非位级一致的情况慢了 $2.4$ 倍。

结论

本文成功构建并验证了一个开源的、基于 TorchTitan 和 vLLM 的位级一致同策略强化学习框架。通过彻底消除训练阶段与推理阶段的计算内核数值差异，模型实现了散度为零的精确匹配。这不仅避免了强化学习中常见的训练崩溃与不稳定问题，还使得模型以更少的训练步数达到了更高的奖励。未来工作将聚焦于模型定义的统一、编译支持以及性能的进一步优化。

补充细节（未来工作展望）

作者将继续在位级一致的训练和推理方向上推进，具体将集中在以下几个核心方向：

• 统一的模型定义（Unified model definition）：尽管当前系统已经展示了位级等价的训练和推理结果，但在工程实现上仍然存在两份独立的模型代码：一份用于训练，另一份用于推理。这种双代码库的方式对于初步的系统集成来说较为简单，但对于长期维护而言却非常脆弱：因为对任何一份模型代码的微小修改，都会立刻破坏训练和推理之间的等价性，进而导致数值不匹配。因此，未来的目标是为训练和推理框架提供一份共享的模型代码，这将彻底消除引入意外人为错误的可能，并使得位级匹配这一重要属性更容易被长期维护。
• 编译支持（Compilation Support）：在当前的实现中，作者没有对 TorchTitan 模型使用 torch.compile，因此在 vLLM 侧也强制启用了 eager 模式。虽然消除这一限制在技术上是直接的，但这需要预先构建一个基于 torch.compile 版本的 TorchTitan 模型。vLLM 本身大量利用了 torch.compile 并且能够在编译模式下保持批次不变性，但是为了保持跨框架的兼容性，必须对训练版本的模型进行相应的修改。这一目标将在后续的工作中被持续推进。
• 强化学习性能优化（RL Performance）：根据当前的基准测试结果，位级一致的 RL 运行速度比非位级一致的情况慢了 $2.4$ 倍。作者将继续通过更好地调优批次不变内核（batch-invariant kernels），以及积极引入包括图编译在内的先进技术，来全面提升 vLLM 在该场景下的执行性能。
• 更广泛的模型支持（Wider Model Support）：作者计划将这个位级一致的强化学习框架扩展到 Qwen3 1.7B 之外，以支持更多种类的开源模型。同时，作者也将泛化现有的审计工具和反向传播的实现逻辑，以覆盖更广泛的算子（operator）类型，从而使位级训练-推理一致性从一个特定模型的实验，转变为一个高度可扩展且可复用的通用特性。

参考文献提取与汇总

• 【1】He 等人，"Defeating nondeterminism in LLM inference"。引用段落：背景知识。引用内容：指出强化学习会放大训练器和采样器之间的微小数值差异，导致非确定性和不稳定的训练行为。
• 【2】Yao, Liu 等人，"Off-policy RL"。引用段落：背景知识。引用内容：共同证实了数值不匹配对强化学习稳定性的负面影响。
• 【3】Liu, Li 等人，"When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch"。引用段落：背景知识。引用内容：揭示了训练与推理不匹配导致的 RL 崩溃现象。