Online draft co-training keeps a small speculative-decoding draft model in sync with an evolving RL policy by training it alongside the policy under the same Context Parallelism (CP) and Pipeline Parallelism layout, delivering 1.5–1.88× end-to-end RL post-training speedups on models up to 122B.
When you post-train a large LLM with RL (say, GRPO on math or agent tasks), most of the wall-clock time is spent generating rollouts, not doing gradient updates. Speculative Decoding is the standard trick to speed that up: a small draft model guesses several next tokens, the big target model verifies them in parallel, and rejection sampling keeps the output distribution unchanged.
The catch is that the policy is moving during RL. A frozen draft trained against yesterday’s policy gets fewer of its guesses accepted as the policy drifts, so the speedup decays. Prior work (FastGRPO, ReSpec) showed you can fix this by co-training the draft online during RL. But nobody had made that work at the scale people actually train large reasoning models: 100B+ parameters, 256K-token contexts, and the usual mix of tensor, pipeline, and context parallelism. Two specific things break. Modern drafts like EAGLE-3, DFlash, and DSpark use branch attention (multiple parallel prediction branches off the main sequence), which the standard causal ring-attention kernels used for long-context training don’t support. And these drafts need to read intermediate hidden states from the target, which under pipeline parallelism live on GPUs that aren’t adjacent to the draft. This paper is the systems plumbing that closes both gaps inside NeMo-RL.
The draft is attached as a submodule on the policy’s last pipeline stage. Its loss is added to the RL loss with a stop-gradient on the target’s hidden states, so the draft learns to mimic the target without perturbing the policy update.
The first hard part is branch attention under context parallelism. In long-context training, the sequence is split across GPUs and keys/values circulate in a ring (Ring Attention). Branch attention breaks the assumption that every query attends to one causal prefix: each branch query needs to see both the causal prefix of the main sequence and a small set of branch-local keys sitting on one specific rank. The authors’ move: compute the two pieces separately, then merge. The main-sequence piece rides the existing packed Zigzag Ring Attention (which load-balances the triangular causal workload). The branch-local piece is computed on the rank that owns the branch anchor. Merging uses the same log-sum-exp trick that ring attention already uses to combine partial softmaxes between ring steps, so nothing about the numerics changes. Because branch KV never leaves its home rank, communication cost is independent of how many branches or how deep they go.
The second hard part is getting target hidden states to the draft under pipeline parallelism. Standard pipeline communication only talks between adjacent stages. The draft on the last stage needs “taps” (intermediate features) from potentially every earlier stage. Their solution, TapChannel, is a side channel that bypasses the pipeline schedule entirely: each source stage has a pre-allocated mailbox slot in the draft stage’s memory, writes its taps there after its forward pass, and the draft reads before its own forward. Cross-node sources use GPUDirect RDMA; colocated sources use CUDA IPC. A sequence-stamp handshake orders the writes and reads. Because taps only flow one way (no gradient return), the schedule doesn’t need to know about them.
# Per microbatch, on each PP stage:
policy_out, taps = policy_forward(microbatch)
for tap, dest_slot in taps:
tapchannel.write(dest_slot, tap, stamp) # side path
# On the draft stage (last PP stage):
for src in tap_sources:
tapchannel.wait_and_read(src, stamp) # rendezvous
draft_loss = draft_forward(policy_out, taps)
total_loss = rl_loss + lambda_ * draft_loss
Learning is preserved. On Qwen3-8B trained with GRPO on DAPOMath-17K and evaluated on AIME 2024, reward curves, validation accuracy, and training-inference KL divergence for the three co-trained draft variants closely track the no-speculation baseline. So the speedup isn’t bought by drifting off-policy.
End-to-end speedups across scales (Table 1): accepted tokens per verification reach 2.28–4.78, giving 1.19–2.23× rollout speedup and 1.16–1.88× end-to-end training speedup across targets from 8B up to 122B (Qwen3.5-122B-A10B and GPT-OSS-120B). DFlash and DSpark consistently beat EAGLE-3 on acceptance length. The authors note that on the large MoE targets, end-to-end speedup is smaller than rollout speedup because each verification forward pass triggers more expert compute through sparse routing.
Multi-turn agent workloads dilute the win. On NeMo Gym Workplace Assistant (a simulated tool-use environment based on WorkBench), rollout-phase speedup is 1.75–2.23× but end-to-end is only 1.25–1.43×, because rollout is 55.8% of step time and tool/environment latency inside rollout can’t be sped up by faster decoding.
The CP kernel is competitive. Against USP (from SpecForge) on the same EAGLE-3 training-time-test workload, the packed zigzag implementation is 2.9× / 2.3× / 1.5× faster at CP=2/4/8 with 2.7× lower per-GPU peak memory, largely because USP pads variable-length batches to 2.25× the real token count while the packed version doesn’t. At 256K tokens with CP=8, they report 7.5× speedup vs CP=1 (94% parallel efficiency).
TapChannel overhead is modest. Per-microbatch tap transport hits 27–39 GB/s, 4.5–8.5× faster than staging through pinned host memory, and the draft stage only waits 0.4–0.6s per policy update for taps (1.5–2.2% of optimization time). The host-staging baseline they compare against slows every rank by more than 80%.
This is systems infrastructure, so the useful decisions are mostly “do I turn this on and where.”
•
If you’re running RL post-training with rollout as the bottleneck and already using speculative decoding, online co-training with a target-feature-conditioned draft (DFlash or DSpark in their results) is worth trying: acceptance length grows monotonically during training and the RL trajectory stayed close to baseline in their runs. The prerequisite is that you can attach the draft as a submodule on the last pipeline stage of your policy learner, which requires framework support.
•
If your rollouts are multi-turn agent workloads with real tool calls, temper expectations. Their numbers show end-to-end speedup roughly halved versus single-turn, because tool latency lives inside the rollout phase and speculation can’t touch it. Profile what fraction of rollout time is actually model forward passes before committing.
•
If you’re training MoE or linear-attention targets, the paper explicitly flags weaker gains: sparse routing makes each verification forward more expensive per accepted token, and linear-attention decode is already cheap so there’s less to save. The authors call this out as future work rather than a solved case.
•
If you’re building your own long-context training stack, the CP recipe (compute causal main-sequence and branch-local attention separately, merge with the online-softmax reduction) is the reusable idea, and it’s not tied to a specific draft architecture: they show it works for EAGLE-3’s training-time-test loop, DFlash’s block-diffusion generation, and DSpark’s Markov head.
Code is linked as a GitHub issue in NeMo-RL.
All learning-preservation evidence is on Qwen3-8B with GRPO on math; the larger-scale runs (35B, 122B, GPT-OSS-120B) are reported for throughput and acceptance length, not for full reward/accuracy tracking against a no-speculation baseline. The end-to-end speedup on the 122B and 120B MoE targets drops to 1.19–1.35×, so the headline 1.88× is an 8B-target number. The multi-turn agent evaluation shows the technique’s benefit is bounded by whatever non-model latency your rollout contains. The comparison against USP is at the attention-operator level for one specific workload (EAGLE-3 TTT with variable-length batches); the 2.9× advantage partly reflects USP’s padding overhead on that workload rather than a pure algorithmic win.