NeMo-RL: Scalable RL with Speculative Decoding

• NeMo-RL is a reinforcement learning framework that accelerates LLM post-training by integrating a speculative decoding loop, preserving the original policy output.
• It employs both synchronous and asynchronous pipelines, leveraging the vLLM backend for efficient weight synchronization and scalable rollout generation.
• Empirical results show a 1.41× overall speedup at 8B scale with projections of up to 2.5× acceleration at 235B scale, validating its system-level efficiency.

NeMo-RL is a reinforcement learning (RL) post-training system for LLMs that integrates speculative decoding within its rollout engine to address the performance bottleneck posed by autoregressive (AR) trajectory generation. By augmenting the standard RL rollout pipeline with a system-integrated speculative decoding primitive—implemented within the vLLM backend—NeMo-RL achieves substantial acceleration in generating training rollouts without altering the output distribution of the target policy. This architecture supports diverse speculation mechanisms, pipelines (synchronous and asynchronous), and scaling strategies, and is validated through extensive empirical and simulated experiments at both 8B and 235B parameter scales (Iso et al., 29 Apr 2026).

1. System Architecture and Integration of Speculative Decoding

NeMo-RL comprises two primary components: a rollout engine built on the vLLM serving framework for token generation and a learner built on the MegatronLM policy model for advantage estimation and policy updates. Each synchronous RL step is decomposed as:

$$T_{\text{step}} = T_{\text{data}} + T_{\text{prepare}} + T_{\text{gen}} + T_{\text{logprob}} + T_{\text{train}}$$

where $T_{\text{gen}}$—the AR rollout generation stage—previously dominated compute time.

Speculative decoding is integrated at $T_{\text{gen}}$. Instead of single-token AR sampling, the rollout engine invokes a speculative loop combining a draft model $Q$ (responsible for proposing token blocks) with the primary verifier model $P$ (the full RL policy). This yields a compositional decoding routine that provably preserves $P$'s output distribution, enabling lossless acceleration.

2. Synchronous and Asynchronous RL Pipelines

NeMo-RL supports both synchronous and asynchronous RL training regimes with speculative decoding:

• Synchronous Pipeline:
  • The learner sends current policy weights ($\theta_P$) to vLLM.
  • vLLM loads $\theta_P$ and a draft model $Q$.
  • Trajectories are generated via vLLM’s speculative decoding loop.
  • The learner recomputes $\log$-probabilities, estimates advantages, and updates policy.
  • An optional hidden-state cache can be accessed for online $T_{\text{gen}}$0 adaptation.
• Asynchronous Pipeline (policy lag $T_{\text{gen}}$1):
  • vLLM rollout workers poll for the latest $T_{\text{gen}}$2, generate speculative rollouts, and queue them.
  • Learner workers asynchronously dequeue rollouts (which may be $T_{\text{gen}}$3 steps stale), recompute $T_{\text{gen}}$4-probs, and update parameters.
  • Critical path generation time is overlapped with downstream RL stages, with effective step time:

  $T_{\text{gen}}$5

  where $T_{\text{gen}}$6.

This dual support enables improved hardware utilization and scaling for large-scale deployments (Iso et al., 29 Apr 2026).

3. Speculation Mechanisms and Algorithms

The speculative decoding in NeMo-RL is general, supporting a variety of $T_{\text{gen}}$7 model configurations. The canonical loop is:

1. $T_{\text{gen}}$8 proposes a block $T_{\text{gen}}$9.

2. For $T_{\text{gen}}$0:

   • Compute acceptance ratio:

   $T_{\text{gen}}$1

- Accept $T_{\text{gen}}$2 with probability $T_{\text{gen}}$3; otherwise, revert to $T_{\text{gen}}$4's sampling.

1. Repeat until end-of-sequence or desired length.

Supported $T_{\text{gen}}$5 mechanisms:

• Native Multi-Token Prediction (MTP) Heads: Utilize auxiliary in-model heads trained during pretraining for $T_{\text{gen}}$6-token prediction. No distillation required, but flexibility is limited to pretraining configuration.

• Small External Draft Models: E.g., a 2B-parameter LLaMA distilled from an 8B policy. $T_{\text{gen}}$7 can be trained offline on $T_{\text{gen}}$8 pairs or adapted online using a gradient-detached policy cache.

• EAGLE-3 Drafting: Training $T_{\text{gen}}$9 to maximize expected block acceptance length:

$Q$0

The overall RL step speedup is bounded by:

$Q$1

where $Q$2 and $Q$3 is the empirical mean acceptance length.

4. System Implementation and Instrumentation

The vLLM backend is extended with a SpeculativeDecoding operator that jointly loads model weights $Q$4 and manages an asynchronous event loop for speculative block verification. Critical implementation features include:

• Weight Synchronization: $Q$5 is broadcast via RDMA or an in-memory object store each step (synchronous) or periodically polled (asynchronous).

• Hidden-State & Log-Prob Cache: Provides gradient-detached hidden states and $Q$6-probs from $Q$7’s forward passes, reusable for online updates to $Q$8 without contaminating policy gradients.

• Stage-Level Telemetry: Precise timing and performance metrics—$Q$9, $P$0, $P$1, $P$2 histograms—are logged and aggregated for system-level analysis and auto-profiling.

5. Experimental Validation and Scalability Results

Empirical studies at both 8B and 235B parameter scales demonstrate the performance and correctness of NeMo-RL with speculative decoding. Key results:

• Synchronous RL on 8B (32 × GH200 GPUs):
  • Baseline generation was 66% of per-step time (100.0s out of 151.2s).
  • With speculative decoding ($P$3, EAGLE-3 $P$4, fixed during RL):
  • $P$5s (1.8× speedup)
  • $P$6s (1.41× overall)
  • Validation accuracy indistinguishable from AR baseline.
  • Step-level metrics:

Policy	$P$7 AR / spec	$P$8 AR / spec
RL-Zero	1.79×	1.41×
RL-Think	1.54×	1.35×

• Scale-up Projections on 235B: (Simulator: Qwen3-235B-A22B; up to 2,048 GPUs)
  • Peak rollout speedup ≈6.49× (α ≈ k)
  • Max end-to-end speedup ≈2.22× (bounded by non-generation stages)
  • In asynchronous mode ($P$9 up to 8), large deployments maintain ≈3–3.5× rollout speedup and ≈2.5× overall at operationally relevant points (e.g., 2,048 GPUs, $P$0, $P$1).

6. Implementation Challenges and Solutions

The integration of speculative decoding in NeMo-RL presents several system-level challenges:

• Weight Synchronization: The need to provide the most recent $P$2 to production-scale vLLM rollout engines without blocking learner progress is solved via low-overhead distributed broadcast protocols or periodic polling.
• Draft Coherence & Alignment: Since $P$3’s acceptance degrades if it drifts from $P$4, seeding $P$5 on in-domain rollouts, and online adaptation leveraging a policy cache, mitigate drift and maximize throughput.
• Instrumentation: Fine-grained breakdowns of all RL and speculation stages are provided via embedded timers and telemetry channels in vLLM, enabling automatic bottleneck detection and optimization.
• Policy-Gradient Isolation: $P$6’s training leverages hidden-state and $P$7-prob caches with gradient detachment to prevent inadvertent updates to $P$8.
• Optimizing Draft Block Length ($P$9): A balance between $\theta_P$0 (proposed block size) and speculative work is sought through empirical profiling. For 8B models, $\theta_P$1 is observed to be optimal.

7. Summary and Significance

NeMo-RL provides a system-integrated, lossless speculative decoding primitive directly in the RL post-training loop of LLMs. Its architecture supports both in-model and external drafting strategies, smoothly composes with synchronous or asynchronous RL orchestration, and signals a shift toward systems-level acceleration of RL rollouts. Empirical and simulation results at multiple scales validate theoretical speedup bounds and confirm that speculative decoding achieves rollout speedups of 1.8× (8B, synchronous), end-to-end speedup of 1.4× (8B), and projected ≈2.5× (235B, async) at large deployments. Thorough instrumentation, modularity across RL regimes, and robust handling of policy/draft drift highlight NeMo-RL’s significance as a production-ready solution for scalable, efficient RL post-training of frontier LLMs (Iso et al., 29 Apr 2026).