← 研踪 PaperTrailarXiv:2609.07108 · v2
Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
首次提交 2026-09-07 14:48
版本更新 2026-11-03 16:51
雷达观测 2026-11-04 23:39
摘要
Speculative decoding accelerates rollout generation, which dominates the cost of reinforcement learning (RL) post-training. Online co-training keeps the draft aligned with the evolving policy, further increasing acceptance length and speedup. However, training large models on long contexts requires context parallelism (CP) and pipeline parallelism (PP), and co-training drafts under this layout poses two obstacles: (1) the draft's branch attention is unsupported by causal CP, and (2) the target features that the draft consumes are spread across PP stages. In this paper, we present the design and implementation of online draft co-training in NeMo-RL. For CP, we extend packed, load-balanced zigzag ring attention by merging rank-local branch attention with causal main-sequence attention. For PP, TapChannel transports intermediate target features across stages via a separate path, leaving the pipeline schedule unaffected. Experiments show that co-training keeps the policy's reward and accuracy close to the no-draft baseline while delivering $1.48$--$2.25\times$ rollout speedups and $1.19$--$1.88\times$ end-to-end training speedups across model scales up to 122B. Our CP design achieves strong scaling at 256K sequence length with significant memory savings over prior work, and our PP transport incurs modest overhead. Code can be found at https://github.com/NVIDIA-NeMo/RL/issues/3698.
作者备注：Technical Report
观测到的版本
相邻记录可能跳过中间版本；首次发现 v2 及以上也会标为版本更新。
v2 · 2026-11-03 16:51 · 最新观测版本
Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
Speculative decoding accelerates rollout generation, which dominates the cost of reinforcement learning (RL) post-training. Online co-training keeps the draft aligned with the evolving policy, further increasing acceptance length and speedup. However, training large models on long contexts requires context parallelism (CP) and pipeline parallelism (PP), and co-training drafts under this layout poses two obstacles: (1) the draft's branch attention is unsupported by causal CP, and (2) the target features that the draft consumes are spread across PP stages. In this paper, we present the design and implementation of online draft co-training in NeMo-RL. For CP, we extend packed, load-balanced zigzag ring attention by merging rank-local branch attention with causal main-sequence attention. For PP, TapChannel transports intermediate target features across stages via a separate path, leaving the pipeline schedule unaffected. Experiments show that co-training keeps the policy's reward and accuracy close to the no-draft baseline while delivering $1.48$--$2.25\times$ rollout speedups and $1.19$--$1.88\times$ end-to-end training speedups across model scales up to 122B. Our CP design achieves strong scaling at 256K sequence length with significant memory savings over prior work, and our PP transport incurs modest overhead. Code can be found at https://github.com/NVIDIA-NeMo/RL/issues/3698.
查看此版本解读 →查看此版本原文 ↗