大规模长上下文 RL 后训练中的在线草稿协同训练:面向投机解码的系统设计

HuggingFace Daily Papers(社区热门论文)·2026-09-07 08:00·2天前
AI 导读

针对投机解码在大规模长上下文 RL 后训练中的扩展难题,研究者提出一套端到端在线草稿协同训练系统。该系统通过扩展 zigzag ring attention 支持分支注意力,并借助 TapChannel 跨流水线并行阶段传输中间特征。实验显示,协同训练的草稿在最高 122B 参数规模上紧密跟踪策略基线,并在 256K token 下实现强扩展和显著内存节省。

HuggingFace Daily Papers(社区热门论文)
41AI 编辑部评分,满分 100

大规模长上下文 RL 后训练中的在线草稿协同训练:面向投机解码的系统设计

2026-09-07 08:00· 2天前
AI 导读

针对投机解码在大规模长上下文 RL 后训练中的扩展难题,研究者提出一套端到端在线草稿协同训练系统。该系统通过扩展 zigzag ring attention 支持分支注意力,并借助 TapChannel 跨流水线并行阶段传输中间特征。实验显示,协同训练的草稿在最高 122B 参数规模上紧密跟踪策略基线,并在 256K token 下实现强扩展和显著内存节省。

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.

来源:HuggingFace Daily Papers(社区热门论文)· arxiv.org