大规模长上下文 RL 后训练中面向投机解码的 Online Draft Co-Training

TL;DR 该系统首次在大规模长上下文 RL 后训练中实现在线草稿协同训练，通过扩展 CP 注意力与跨阶段特征传输，将推测解码加速扩展到 122B 模型与 256K tokens。

问题

问题背景

RL post-training 已成为构建推理与 agentic LLM 的标准范式，但 rollout 生成占据训练大部分时间，speculative decoding 是主要加速手段。

现有方法局限

• 在线 co-training 可进一步提升 draft 模型准确率、带来更大加速，但在大规模长上下文场景下面临两个技术障碍：
  1. Branch attention 无法被标准 causal context-parallel (CP) 实现支持，导致 draft 分支与主线序列的联合注意力无法高效计算。
  2. Target features 跨越 pipeline-parallel (PP) 阶段，传统 PP 调度难以在不破坏流水线的情况下传递中间特征。
• 已有 CP 方案（如 packed zigzag ring attention）未考虑 branch attention 融合，而 PP 方案缺乏轻量的跨阶段特征传输路径。

为什么这个问题难且重要

• 挑战在于：大规模模型（至 122B）与超长上下文（256K tokens）下，需同时保持 CP 负载均衡、显著降低内存开销，并让 PP 传输基本不影响原有流水线调度。
• 若无法解决，在线 co-training 只能受限于单阶段或短上下文，不能部署到当前主流的分布式 RL post-training 集群。
• 业界关注度高，因为 RL post-training 的 wall-clock time 直接决定模型迭代速度，任何显著加速都相当于降低大规模训练成本。

行业类比

类似在实时对话系统中用轻量 draft 模型快速生成候选，再由大模型验证，但这里引入在线协同训练，使 draft 与 policy 在复杂并行框架下同步进化。

核心洞察

• 在线草稿协同训练在 RL 后训练中扩展到大规模长上下文场景，核心挑战在于分支注意力的上下文并行支持与跨流水线阶段的特征传输。本文通过扩展打包的负载均衡 zigzag ring attention 将分支注意力融入因果主序列注意力，并用 TapChannel 旁路传输目标特征，从而在不增加显存压力或干扰训练流水线的前提下实现草稿模型在线训练。这与以往仅在小规模或短序列验证的在线 co-training 工作形成对比，证明了在实际分布式训练中部署推测解码加速的可行性。
• 本文独特之处在于对底层分布式注意力的定制化融合：将 rank-local branch attention 直接合并到 causal main-sequence attention 中，避免了为草稿分支单独分配 KV cache 或引入额外通信。这种方式与传统的独立分支计算或参数共享策略不同，它使得在线 co-training 的内存开销显著低于先前方法，且能随上下文长度扩展至 256K tokens。工程上这意味着可以在现有训练框架中高效加入推测解码，无需重写数据并行或张量并行逻辑。

方法

方法流程

该方法面向 RL post-training 中的 rollout 生成加速。输入为 target policy（大型 LLM）与 draft model（小型模型）组成的 speculative decoding 系统，在长上下文、大规模并行环境下进行在线 co-training。RL 训练过程中，利用 rollout 数据持续更新 draft，使其逐步逼近 target 分布，提高候选接受率，从而加速生成。

关键模块

1. CP 分支注意力：draft 的 branch attention（并行预测多个候选 token 的注意力模式）与主序列的 causal attention 需同时计算。标准 context-parallel 实现不支持这种混合计算。本工作扩展了 packed, load-balanced zigzag ring attention，将 rank-local branch attention 融合进 causal main-sequence attention，避免额外全量通信，在 256K token 场景下获得强扩展性和显著内存节省。

2. TapChannel 跨阶段特征传输：在 pipeline-parallel 下，target 模型的中间特征分散在不同 stage。为增强 draft 的条件生成，需将这些特征传递到 draft 所在 stage。TapChannel 通过独立数据通道传输目标中间特征，不干扰原有 pipeline schedule，实现低开销、非阻塞的特征注入。

输出为在线 co-trained draft 模型，其行为与 policy baseline 高度一致，并显著加速 rollout 和端到端训练。

与同类在线 co-training 方案相比，本工作首次在 CP 与 PP 双重并行下实现大规模长上下文 RL post-training 的端到端加速。

实验

实验设计叙述

论文在 RL 后训练场景下验证在线草稿共训练的效果，覆盖多种模型规模（最高 122B）。实验包含四个核心部分：策略学习在 speculative decoding 下的一致性验证；不同草稿模型与模型规模下的性能对比（含多轮工作负载）；上下文并行注意力在长上下文（256K tokens）下的扩展性分析；以及流水线并行传输开销的微基准与全运行评估。具体任务与数据集未在摘要中披露。

关键发现

共训练的草稿模型紧密跟踪策略基线，在 rollout 生成和端到端训练速度上取得显著加速。上下文并行设计在 256K tokens 长度下展现强扩展性，相比先前工作大幅节省内存。流水线并行特征传输（TapChannel）引入的开销适中，不影响整体训练效率。

与基线对比解读

与标准投机解码（无在线共训练）相比，本文方法通过扩展分支注意力上下文并行和跨阶段的 TapChannel 传输，解决了大规模长上下文场景下在线共训练的两大障碍。相比已有工作，CP 内存节省明显，PP overhead 可控，从而支撑 122B 级别的规模化 RL 后训练。

行业影响

落地场景

该技术面向大规模 RL 后训练场景，尤其适合推理模型与Agentic LLM 的迭代优化。可用于云推理服务、企业内部训练平台及 Agent 应用，例如电商智能客服的强化学习微调、金融合规问答 Agent 的长文档推理训练。

商业价值

• 降本：rollout 生成在 RL 训练中成本占比高，speculative decoding 加速可大幅降低 GPU 时费用。
• 提质增效：在线 co-training 使 draft 模型持续适配 evolving policy，保持高接受率，加速训练收敛。
• 体验提升：推理侧更低首 token 延迟和更高吞吐，直接改善用户交互体验。

与现有产品/工作流的接口

与 NVIDIA NeMo、vLLM、SGLang 等训练/推理框架集成。其 CP 扩展 packed zigzag ring attention 可接入现有 context-parallel 实现；TapChannel 作为独立特征传输路径，对 pipeline schedule 无侵入，适合在现有 PP 架构中增量部署。

具体用例：

1. 企业级代码助手：在大型代码库上下文中做 RL 微调，利用 256K 长序列 CP 加速，缩短策略迭代周期。
2. 医疗问答系统：基于长病历推理生成，通过 online co-training 保持 draft 与 target 一致，降低线上推理成本。

局限

• - **通用性限制**：论文针对 **CP + PP** 并行组合下的在线 draft 协同训练提出定制化系统设计（扩展 zigzag ring attention 支持分支注意力，TapChannel 跨 stage 传输特征）。该方案深度耦合于特定的并行策略与注意力实现，对于使用 **张量并行（TP）** 为主、或采用 **序列并行（SP）** 其他变体（如 Ulysses）的集群，以及非 Transformer 架构的 draft 模型，可能无法直接迁移。实验仅在 NVIDIA 相关框架下验证，跨硬件生态（如 TPU、AMD）的可移植性未讨论。
• - **实验覆盖与评估维度有限**：实验主要报告了加速比和策略学习跟踪效果，但未深入分析 **在线协同训练引入的额外计算/通信开销对总吞吐的影响** 在多种 RL 任务（如数学推理、代码生成、多轮对话）上的差异。同时，尽管扩展到 122B 和 256K tokens，但未对比不同 draft 模型结构（如独立小模型 vs 同一模型的部分层）的敏感度，也未提供在 **更长上下文（>256K）或更大 batch** 下 CP 设计的 scaling 极限。
• - **与现有在线 draft 训练方法的对比不足**：论文声称在线协同训练提升 draft 精度，但未将所提系统与已有的 **online draft adaptation** 方法（如 EAGLE-2、SpecTrain 等）在相同训练成本下做全面对比，包括 **收敛速度、最终加速率、内存峰值**。仅与 baseline（无在线训练）比较，难以证明该工程方案相对于其他在线训练途径的增量价值。此外，论文未讨论 draft 与 target 分布漂移过大时系统的鲁棒性，以及如何动态调整协同训练频率。