If you have been running reinforcement learning (RL) post-training on a language model for math reasoning, code generation, or any verifiable task, you have almost certainly stared at a progress bar while your GPU cluster burns through rollout generation. A team of researchers from NVIDIA proposes a precise fix by integrating speculative decoding into the RL training loop itself, and do it in a way that preserves the target model’s exact output distribution.
The research team integrated speculative decoding directly into NeMo RL v0.6.0 with a vLLM backend, delivering lossless rollout acceleration at both 8B and projected 235B model scales.The latest NeMo RL v0.6.0 release officially ships speculative decoding as a supported feature alongside the SGLang backend, the Muon optimizer, and YaRN long-context training.
To understand the problem, it helps to know how a synchronous RL training step breaks down. In NeMo RL, each step consists of five stages: data loading, weight synchronization and backend preparation (prepare), rollout generation (gen), log-probability recomputation (logprob), and policy optimization (train).
The research team measured this breakdown on Qwen3-8B under two workloads — RL-Think, which continues training a reasoning-capable model, and RL-Zero, which starts from a base model and learns reasoning from scratch. In both cases, rollout generation accounts for 65–72% of total step time. Log-probability recomputation and training together take only about 27–33%. This makes generation the only stage worth targeting for acceleration, and the one that determines the ceiling for any rollout-side optimization.
Speculative decoding is a technique where a smaller, faster draft model proposes several tokens at once, and the larger target model (the one you are actually training) verifies them using a rejection sampling procedure. The key property and why it matters for RL, is that the rejection procedure is mathematically guaranteed to produce the same output distribution as if the target model had generated those tokens autoregressively. No distribution mismatch, no off-policy corrections needed, no change to the training signal.
This is important because in RL post-training, the training reward depends on the policy’s own samples. Methods like asynchronous execution, off-policy replay, or low-precision rollouts all trade some amount of training fidelity for throughput. Speculative decoding trades nothing: the rollouts are identical in distribution to what the target model would have generated on its own, just produced faster.
Adding a draft model to a serving backend is straightforward. Adding one to an RL training loop is not. Every time the policy updates, the rollout engine must receive new weights. The draft model must remain aligned with the evolving policy. Log-probabilities, KL penalties, and the GRPO policy loss must all be computed against the target (verifier) policy not the draft or the optimization target is silently corrupted.
The NVIDIA research team handles this in NeMo RL with a two-path architecture. The general path uses EAGLE-3, a drafting framework that works with any pretrained model without requiring native multi-token prediction (MTP) support. A native path is also available for models that ship with built-in MTP heads. When online draft adaptation is enabled, the hidden states and log-probabilities from the MegatronLM verifier forward pass are cached and reused to supervise the draft head via a gradient-detached pathway, so draft training never interferes with the policy gradient signal.
On 32 GB200 GPUs (8 GB200 NVL72 nodes, 4 GPUs per node), EAGLE-3 reduces generation latency from 100 seconds to 56.6 seconds on RL-Zero — a 1.8× generation speedup. On RL-Think, it drops from 133.6 seconds to 87.0 seconds, a 1.54× speedup. Because log-probability re-computation and training are unchanged, these generation-side gains translate to overall step speedups of 1.41× on RL-Zero and 1.35× on RL-Think. Validation accuracy on AIME-2024 evolves identically under autoregressive and speculative decoding throughout training, confirming that the lossless guarantee holds in practice.
The research team also tests n-gram drafting as a model-free speculative baseline. Despite achieving acceptance lengths of 2.47 on RL-Zero and 2.05 on RL-Think, n-gram drafting is slower than the autoregressive baseline in both settings — 0.7× and 0.5× respectively. This is a critical finding for practitioners: a positive acceptance length is necessary but not sufficient. If the verification overhead is high enough, speculation makes things worse.
The research team isolates three operational choices that practitioners must get right.
Draft initialization matters more than generic drafting ability. An EAGLE-3 draft initialized on the DAPO post-training dataset achieves a 1.77× generation speedup on RL-Zero, while a draft initialized on the general-purpose UltraChat and Magpie datasets achieves only 1.51× at the same draft length. The draft must be aligned with the actual rollout distribution encountered during RL, not just a broad chat distribution.
Draft length has a non-obvious optimum. At draft length k=3, RL-Zero achieves 1.77× speedup and RL-Think achieves 1.53×. Increasing to k=5 raises the acceptance length but drops speedup to 1.44× on RL-Zero and 0.84× on RL-Think — the latter already slower than autoregressive. At k=7, RL-Zero drops further to 1.21× and RL-Think to 0.71×. The contrast matters: RL-Zero’s rollouts are generated from a base model starting with short outputs, making them easier for the draft to predict even at high k. RL-Think’s fully developed reasoning traces are harder to speculate over, so the overhead of longer drafts erases the benefit sooner. More speculative work per step can erase the benefit of higher acceptance entirely, especially in harder generation regimes.
Online draft adaptation — updating the draft during RL using rollouts generated by the current policy helps most when the draft is weakly initialized. For a DAPO-initialized draft, offline and online configurations perform nearly identically (1.77× vs. 1.78× on RL-Zero). For a UltraChat-initialized draft, online updating improves speedup from 1.51× to 1.63× on RL-Zero.
Interaction with asynchronous execution was also tested directly at 8B scale not just in simulation. The research team ran RL-Think at policy lag 1 in a 16-node non-colocated configuration, with 12 nodes dedicated to generation and 4 to training. In asynchronous mode, most of rollout generation is already hidden behind log-probability re-computation and policy updates, so the relevant quantity is the exposed generation time that remains on the critical path. Speculative decoding reduces that exposed generation time from 10.4 seconds to 0.6 seconds per step and lowers effective step time from 75.0 seconds to 60.5 seconds (1.24×). The gain is smaller than in synchronous RL — expected, since asynchronous overlap already hides much of the rollout cost — but it confirms that the two mechanisms are genuinely complementary rather than redundant.
Using a proprietary GPU performance simulator calibrated to device-level compute, memory, and interconnect characteristics, the research team projected speculative decoding gains at larger scales. For Qwen3-235B-A22B running synchronous RL on 512 GB200 GPUs, draft length k=3 with an acceptance length of 3 tokens yields a 2.72× rollout speedup and a 1.70× end-to-end speedup.
At the most favorable simulated operating point — Qwen3-235B-A22B on 2048 GB200 GPUs with asynchronous RL at policy lag 2 — rollout speedup reaches approximately 3.5×, translating to a projected 2.5× end-to-end training speedup. Speculative decoding and asynchronous execution are described as complementary: speculation reduces the cost of each individual rollout, while asynchronous overlap hides the remaining generation time behind training and log-probability computation.
Check out the Full Paper and Nemo RL Repo. Also, feel free to follow us on Twitter and don’t forget to join our 130k+ ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。