🤖 machine learning

Accelerating RL Post-Training Rollouts via System-Integrated Speculative Decoding

本文在 NeMo-RL 中引入了一种系统集成的投机解码框架，作为强化学习后训练 rollout 的无损加速原语，在 8B 规模下实现了 1.8 倍的吞吐量提升，并预计在结合异步执行后，在 235B 规模下可实现高达 2.5 倍的端到端训练加速。

原作者： Hayate Iso, Tiyasa Mitra, Sudipta Mondal, Rasoul Shafipour, Venmugil Elango, Terry Kong, Yuki Huang, Seonjin Na, Izzy Putterman, Benjamin Chislett, Maor Ashkenazi, Joseph Guman, Gerald Shen, Tugrul Ko

发布于 2026-04-30

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Hayate Iso, Tiyasa Mitra, Sudipta Mondal, Rasoul Shafipour, Venmugil Elango, Terry Kong, Yuki Huang, Seonjin Na, Izzy Putterman, Benjamin Chislett, Maor Ashkenazi, Joseph Guman, Gerald Shen, Tugrul Konuk, Ashwath Aithal, Ritika Borkar, Ran Zilberstein, Bita Rouhani

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在训练一位才华横溢但思考缓慢的学生（即 AI 模型）来解决复杂的数学问题。为了教导他们，你必须向他们提出一个问题，等待他们将整个思维过程一步步写出来，然后检查他们是否答对了。

问题在于，“检查”部分很快，但“书写”部分却极其缓慢。这位学生写一个字，停顿思考，再写下一个字，再次停顿，如此循环。这种逐字书写的缓慢过程，是训练这些 AI 模型的最大瓶颈。

本文提出了一种名为**推测解码（Speculative Decoding）**的巧妙技巧，旨在加速这一书写过程，同时不改变学生的思考方式或所学内容。

核心理念：“起草助手”

将 AI 模型想象成一位非常精准但动作缓慢的主厨。为了加快速度，你雇佣了一位快速、精力充沛的副厨（即“草稿模型”）。

旧方法（自回归）： 主厨写一个字，停下，思考，再写下一个字，停下，思考。这耗费了漫长时间。
新方法（推测解码）： 副厨快速猜测主厨可能写出的接下来的 3 到 4 个字。副厨迅速将这些字写下来。
检查： 主厨随后快速扫视副厨的笔记。
- 如果笔记正确，主厨会说：“太棒了！”并一次性接受全部 4 个字。
- 如果笔记错误，主厨会划掉它们，写出正确的字，然后重新开始。

神奇之处： 因为副厨动作很快，主厨大多数时候都能一次性接受多个字。最终结果与主厨独自书写完全一致，但发生速度快得多。

论文实际做了什么

研究人员将这一系统构建到了一个名为NeMo-RL的真实训练框架中。他们不仅测试了简单任务，还测试了需要 AI 深入思考的“推理”任务（例如解决数学问题）。

以下是他们的主要发现，已转化为通俗易懂的表述：

无需作弊即可生效： 一些加速方法试图走捷径（例如使用低质量的数学或跳过步骤），这可能会破坏学生的学习效果。这种方法则是“无损”的。它保证 AI 的学习方式与没有助手时完全相同，只是速度更快。
速度提升：
- 在一个中等规模模型（80 亿参数）上，他们观察到书写过程的速度提升了1.5 到 1.8 倍。
- 由于书写占据了总训练时间的约 70%，因此整个训练过程的速度提升了约1.35 到 1.4 倍。
“草稿”至关重要： 副厨必须是一个优秀的猜测者。
- 如果你让副厨在学生正在学习的完全相同类型的数学问题上进行训练，速度提升将非常巨大。
- 如果你使用一个只懂通用聊天的普通副厨，速度提升则较小。
- 不要猜得太远： 如果副厨试图一次猜测 7 个字，他们会犯太多错误，导致主厨花费过多时间进行纠正。一次猜测 3 个字是“最佳平衡点”。
未来（大模型）： 他们使用了一个超精准的计算机模拟器，来预测在大规模模型（2350 亿参数）以及数千台计算机协同工作时会发生什么。
- 他们预测，对于这些巨型模型，这项技术可能使整个训练过程的速度提升2.5 倍。

为何这很重要

在 AI 领域，时间就是金钱。如果你能将模型训练速度提高 2.5 倍，你就可以在相同的时间内获得更聪明的模型，或者以一小部分成本获得相同的模型。

这篇论文证明，你无需改变 AI 的“大脑”或游戏规则就能获得这种速度。你只需要添加一个聪明、快速的助手来协助起草，并让主模型进行最终验证。这是一项系统升级，它使整个训练流程运行得更加顺畅，而不会破坏任何现有内容。

以下是论文《通过系统集成的推测解码加速强化学习后训练 rollout》的详细技术总结。

1. 问题陈述

前沿大语言模型（LLM）的强化学习（RL）后训练正日益受到自回归 rollout 生成的瓶颈限制。在面向推理的工作负载（如数学、代码）中，生成轨迹（rollouts）所花费的时间占据了总训练时间的主导地位，往往超过了梯度计算所花费的时间。

现有的效率提升方法试图通过改变训练动态来解决这一问题，例如：

异步执行： 将生成与学习重叠，但会引入策略滞后。
离线策略回放： 重用陈旧的轨迹，但需要重要性采样修正。
低精度 rollout： 减少计算量，但引入分布不匹配。

这些方法以有效性（学习信号的质量）换取吞吐量。本文识别出一种迫切需求，即需要一种能够加速生成而不改变目标策略采样分布的方法，从而保留 RL 稳定收敛所需的精确训练语义。

2. 方法论

作者提出将推测解码（Speculative Decoding）作为一种无损加速原语，直接集成到NeMo-RL训练框架（使用 vLLM 后端）中。

核心机制

草稿生成（Drafting）： 一个“草稿”模型（或头部）并行提出多个 token（ $k$ 个）。
验证（Verification）： 目标“验证器”模型（即 RL 策略）在单次前向传播中验证这些 token。
拒绝采样（Rejection Sampling）： 根据验证器的概率分布接受或拒绝 token。这保证了最终的 rollout 分布与目标策略完全相同，确保 RL 优化目标不发生任何改变。

系统集成

该集成超越了标准的推理服务，解决了 RL 特定的约束：

权重同步： 在 RL 训练期间，草稿模型必须与快速变化的策略权重保持对齐。
双重草稿路径：
- 通用路径（EAGLE-3）： 使用在策略输出上训练的外部草稿模型。这适用于任何预训练模型。
- 原生路径： 如果模型架构支持，则使用内置的多 Token 预测（MTP）头部。
梯度分离： 为防止草稿模型的训练干扰策略的梯度信号，在策略前向传播期间缓存隐藏状态和对数概率，并通过 detach() 边界路由至草稿头部。
流水线兼容性： 该系统支持同步（生成等待训练）和异步（生成与训练重叠）的 RL 流水线。

3. 主要贡献

系统集成： 首次在生产级、开源的 RL 训练栈（NeMo-RL）中实现了推测解码，同时保留了验证器精确（verifier-exact）的训练语义。
操作分析： 对决定加速效果的因素进行了全面研究，包括草稿初始化、草稿长度和在线适应。
可扩展性预测： 使用高保真性能模拟器预测前沿规模（高达 235B 参数和 2048 张 GPU）的收益，超越了当前 8B 实验设置的局限性。

4. 实验结果

实验设置

模型： Qwen3-8B（RL-Think）和 Qwen3-8B-Base（RL-Zero）。
任务： 数学推理（DAPO-Math-17K 训练，AIME-2024 验证）。
硬件： 8 个 GB200 NVL72 节点（共 32 张 GPU）。
基线： 标准自回归解码。

关键发现

吞吐量增益：
- 生成延迟： 推测解码将生成时间减少了1.54 倍（RL-Think）和1.77 倍（RL-Zero）。
- 端到端步骤时间： 整体训练步骤加速分别达到1.35 倍和1.41 倍。收益受限于非生成阶段（对数概率计算和训练），推测解码无法加速这些阶段。
- 准确率： 自回归解码与推测解码之间的验证准确率曲线无法区分，证实了该加速方法的无损特性。
操作洞察：
- 初始化： 在领域内数据（DAPO-Math）上初始化草稿模型，显著优于通用聊天领域初始化（UltraChat），产生了更长的接受长度和更高的加速比。
- 草稿长度（ $k$ ）： 与直觉相反，**较短的草稿（ $k=3$ ）**产生了最佳的端到端加速。较长的草稿（ $k=5, 7$ ）增加了验证开销，有时甚至使过程比自回归解码更慢。
- 在线适应： 仅当初始草稿对齐较差时，在线更新草稿模型才提供边际收益。对于初始化良好的草稿，离线训练已足够。
- 异步 RL： 在异步模式下，推测解码仍然具有互补性。虽然异步重叠隐藏了部分生成成本，但推测进一步减少了关键路径上“暴露”的生成时间，在特定的异步配置下实现了1.24 倍的加速。

模拟预测（部署规模）

使用模拟器针对2048 张 GPU上的235B 模型进行预测：

Rollout 加速： 在最优条件下（特定的草稿/接受长度和适度的策略滞后），可超过3.5 倍。
端到端加速： 预计达到**~2.5 倍**，这是因为在前沿规模模型中，生成所花费的时间占比很高。
规模敏感性： 在大规模部署中，更大的模型从推测中获益更多，因为它们更受长尾延迟和利用率不足的影响，而推测可以缓解这些问题。

5. 意义

本文证明了推测解码是一种可行的、无损的加速 RL 后训练策略。其主要意义在于：

保留语义： 与其他加速技术不同，它不会损害 RL 训练信号，也不需要离线策略修正。
系统级优化： 它弥合了推理加速技术与 RL 训练循环复杂需求（权重同步、梯度处理）之间的差距。
可扩展性： 在前沿规模下预测的2.5 倍端到端加速表明，推测解码可能成为未来 RL 基础设施的标准组件，显著减少训练推理能力模型所需的时间和成本。

该工作确立了一条清晰的部署路径：使用领域内初始化的草稿配合适度的草稿长度（例如 $k=3$ ），以在保持训练稳定性的同时最大化吞吐量。