通过在线协同训练提升大模型长文本推理速度，显著加速强化学习后训练。
Online Draft Co-Training for Speculative Decoding in Large-Scale, Long-Context RL Post-Training
- 提出端到端系统解决分支注意力与跨流水线特征传输难题
- 在1220亿参数模型上实现256K上下文下的高效扩展与内存节省
- 适合需要高速推理的大规模强化学习场景
推测解码可加速强化学习后训练中的生成过程，该过程占主导成本。在线协同训练能进一步提升草稿模型精度，带来更大加速。然而，将此方法扩展至大规模、长上下文模型面临两大挑战：(1) 标准因果上下文并行（CP）实现不支持分支注意力；(2) 目标特征跨越流水线并行（PP）阶段。本文提出端到端系统解决上述问题：针对CP，通过合并秩局部分支注意力与因果主序列注意力，扩展了打包、负载均衡的锯齿环注意力；针对PP，TapChannel通过独立路径传输中间目标特征，不影响流水线调度。实验表明，协同训练的草稿模型能紧密跟踪策略基线，在高达1220亿参数的模型上实现显著的滚动生成与端到端加速。我们的CP设计在256K标记长度下表现良好，相比先前工作具有显著内存节省；而PP传输仅引入轻微开销。代码可在 https://github.com/NVIDIA-NeMo/RL/issues/3698 获取。
原文摘要 · Abstract (English)
Speculative decoding accelerates rollout generation, which dominates the cost of reinforcement learning (RL) post-training. Online co-training can further increase the draft's accuracy, yielding greater speedups. However, scaling this approach to co-training on large models with long contexts poses two obstacles: (1) branch attention is unsupported by standard causal context-parallel (CP) implementations, and (2) target features span across pipeline-parallel (PP) stages. We address both with an end-to-end system for large-scale online draft co-training. 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 demonstrate that co-trained drafts closely track the policy baseline while delivering substantial rollout and end-to-end speedups across model scales up to 122B. Our CP design achieves strong scaling at 256K tokens 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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。