当训练一个大模型陪着它一起长大的小助手，会遇到什么麻烦

2024 年前后，强化学习后训练成了打造推理型大模型的标配做法。你可以把它理解成，模型先学会说话，再通过不断试错、拿反馈来学会把话说得更好、更聪明。

但这里有个几乎所有做这件事的人都会撞上的墙：训练时间大部分不是花在算梯度上，而是花在生成上。模型要先自己写出一堆候选答案，也就是所谓的"rollout"

推理式强化学习的采样阶段，模型需要自回归地一个词一个词往外蹦，写出完整的候选回答，这个过程通常占掉训练总时长的大头

写这些候选答案的过程慢得让人心疼。原因也简单：主流大模型都是自回归生成，写第 100 个字之前必须先写完前 99 个字，没法并行。模型越大，长文本任务越常见，这个瓶颈就越致命。

于是投机采样这个点子被搬了出来。它的思路挺讨巧：找一个小模型打前站，快速吐出好几个候选词，再让大模型一次性并行验证这几个词对不对。对的就留下，错的就扔掉重写。因为验证是并行的，整体速度就能飞起来，而且理论上不会改变最终生成的内容分布，你不会为了速度牺牲效果。

但这里有个隐藏的坑，如果小模型（论文里叫"draft"，草稿模型）一直用训练开始时的老版本，跟不上大模型（目标模型，也叫 policy）越训越强的脚步，它猜中的概率就会越来越低，投机采样的红利也就越来越薄。于是很自然的想法是：让草稿模型跟着目标模型一起训练，边训练边磨合，这就是论文说的在线协同训练。

思路是对的，但真正做起来，工程上冒出来两个非常具体的麻烦。这篇来自 NVIDIA 的论文，就是奔着解决这两个麻烦去的。

草稿模型的注意力机制，和现有的并行训练框架，其实是两套互不兼容的语言

大规模训练一个模型，光靠一块 GPU 是不够的，业界普遍要把模型和数据切开摊到几十甚至几百块卡上同时算。切的方式主要有几种：把一个长序列切成几段分给不同卡（这叫上下文并行，CP），或者把模型按层切开一层一层往下传（这叫流水线并行，PP）。

问题出在，像 EAGLE-3、DFlash、DSpark 这类比较先进的草稿模型架构，它们的注意力计算方式跟目标模型不太一样。目标模型算注意力时是标准的"因果"结构，说白了就是每个词只能看它前面的词，规规矩矩排队。但草稿模型为了训练时能一次性模拟多步预测，会长出"分支"结构，一个词除了要看前面的正常序列，还要额外看它自己那条训练分支里刚生成的一堆词。这种分支结构，现成的上下文并行代码完全不认识，直接套用就会算错或者算不了。

另一个麻烦更绕。草稿模型需要参考目标模型中间层的隐藏状态（论文管这个叫"tap"，特征抽头）才能做出准确预测，可问题是，在流水线并行下，目标模型的层被拆到了好几块不同的 GPU 上，草稿模型自己被安放在最后一个阶段，它需要的中间特征散落在别的、离它很远的阶段上。而标准的流水线通信只负责相邻阶段之间传数据，压根没设计过怎么把中间某一站的东西直接送到终点站。

这两个问题合起来就是一句话：草稿模型没法简单地"寄居"在现有的并行框架里，它得要一套专门的搬运和计算方案。

这就好比你在一栋已经按标准流程装修好的写字楼里，突然要给某个楼层加装一套完全不同规格的空调系统，而且这套系统还得从楼上和楼下好几层同时抽气过来循环。如果硬套原来的中央空调管道，气流根本对不上接口，你必须单独拉一条新风管才能跑通。

给分支注意力开一条"绿色通道"：把它拆成两半分开算，再拼起来

论文给出的第一个解决方案，针对的是上下文并行下算不了分支注意力的问题。

核心思路是拆分。研究者观察到，草稿模型里每个分支查询词，其实只需要看两类东西：一类是主序列里排在它前面的正常内容（这部分是标准的因果结构，可以按老办法切分到各卡上算），另一类是它自己那条分支里的局部内容（这部分数据量小，而且天然待在生成它的那张卡上，不需要跨卡传输）。

论文的具体做法是，先用现成的"打包式、负载均衡的锯齿环形注意力"（一种优化过的、把长序列切成之字形分布到各卡、通过环形通信传递 KV 键值对的计算方式）算出主序列部分的结果，再在本地单独算一遍分支部分的结果，最后把两部分结果按照在线 softmax 的标准数学公式合并成最终输出。

打包*：把长短不一的多条数据紧凑排列在一起处理，不用像传统做法那样为了对齐长度而填充大量无意义的空白数据，省内存也省计算。

锯齿环形注意力*：一种上下文并行技术，把长序列切给多张卡，各卡的 KV 数据按环形依次传递，同时把序列切块方式设计成之字形排列，让每张卡分到的计算量尽量均衡，不会有的卡算得多有的卡算得少。

在线 softmax*：一种分批次累加计算注意力权重的数学技巧，不需要一次性把所有分数都算出来再归一化，可以边算边合并，这对于分块并行计算长序列是必需的。

这个设计有个很讨喜的副作用：因为分支数据完全留在本地不跨卡传输，所以不管草稿模型的分支有多少条、多深，跨卡通信的开销都完全不受影响，一直保持恒定。这对三种不同架构的草稿模型（EAGLE-3、DFlash、DSpark）都通用，等于用一套底层机制服务了三种不同的上层设计。

反过来想一下，如果不这么拆分会怎样。之前的 SpecForge 项目也尝试过在并行环境里跑 EAGLE-3 的分支注意力，但它用的是顺序环形分片，这会导致各卡算力负载不均衡，有的卡忙死有的卡闲着，而且对草稿模型的 KV 头数量还有额外限制，扩展性打了折扣。论文里的方案通过把两种注意力彻底解耦再合并，绕开了这个限制。

实测数据也印证了这套方案确实划算。跟 USP（一种现有的并行注意力实现，靠填充对齐批次数据）比起来，在 CP 度数为 2、4、8 的情况下，论文的打包式方案分别把延迟压低了 2.9 倍、2.3 倍、1.5 倍，每张卡的峰值显存还省了 2.7 倍。而且序列长度拉到 256K token 的极限场景下，TTT（EAGLE-3 用来模拟多步自回归生成、暴露多步误差的训练时测试机制）注意力的延迟从 CP=1 时的 17.7 秒降到 CP=8 时的 2.35 秒，并行效率达到 94%，这个数字放在长上下文训练里已经相当接近理论满分。

给流水线开一条"专用快递通道"：TapChannel

第二个方案针对的是流水线并行下特征传输的问题，论文给它起了个名字叫 TapChannel。

它的设计思路挺干脆：既然标准流水线通信只管相邻阶段传数据，那就完全不走这条主干道，另开一条侧路。具体做法是，在草稿模型所在的最后一个阶段，给每个会产生特征的源头阶段都预留一个专属的"信箱"（缓冲区槽位）。每个源阶段算完自己那部分前向传播后，就把结果写进对应信箱，草稿模型要用的时候，就在自己开始前向计算之前直接去读这个信箱。

为了避免读写冲突，每个信箱还挂了一个"时间戳"，写的一方和读的一方各自在读写时递增这个戳，用来确认数据是不是最新的、有没有被正确消费过。同卡内的传输走的是 CUDA IPC（同一块 GPU 上不同进程间共享内存的技术），跨节点的传输则用上了 GPUDirect RDMA（能让不同服务器上的 GPU 显存直接互相读写，不用绕经 CPU 内存的高速网络传输技术），速度快很多。

这套侧路设计最巧妙的地方在于，它完全躲开了流水线原本的调度节奏。流水线本身在处理每个微批次时，会有一段自然的"空闲窗口"：源头阶段算完 tap 特征，到草稿模型真正需要用它之间，总有那么一小段时间差。只要传输能在这段窗口内完成，就完全不会拖慢原本的训练节奏，等于是白捡的时间。

这就像你在一栋大楼里管理一批需要多方会签的文件。传统做法是文件必须严格沿着一楼、二楼、三楼逐层往上递送，哪怕三楼急等一份来自一楼的资料，也得按流程排队等。TapChannel 相当于给每层楼装了个直达的传送带，一楼文件一签完字，直接滑进三楼专属的收件格里，三楼什么时候要用什么时候去取，不用等楼层间的正常公文走完流程。如果没有这条传送带，你只能让文件在每层楼之间暂存转手，多绕好几圈不说，还得占用主干道的通行时间，拖慢整栋楼的正常运转。

论文专门测了一下这条侧路和传统做法（把数据先暂存到主机内存再转运）的差距。结果相当悬殊：TapChannel 的单边写入完成整个多源汇入的速度，比经过主机内存中转快了 4.5 到 8.5 倍。更关键的是，它几乎不占用源头阶段的算力资源，干扰在误差范围内，接收端的草稿模型阶段也只多了 1.6% 的显存争用，而传统的主机中转方案会让每个阶段的计算速度都拖慢 80% 以上。

真金白银的收益：草稿模型学得像，跑得快

这套系统设计得再漂亮，最终还是要看两件事：草稿模型跟着一起训练会不会把原本的强化学习训练搞歪，以及速度到底能提升多少。

论文在 Qwen3-8B 上同时验证了三种草稿架构：EAGLE-3、DFlash、DSpark，训练数据用的是 DAPOMath-17K 数学题集，验证用 AIME 2024 竞赛题。结果显示，带草稿协同训练的几条曲线，无论是奖励值、验证准确率还是训练推理之间的 KL 散度（衡量两个概率分布差异的指标，这里用来确认训练和推理两套系统在数值上是不是保持一致），都和不做任何投机采样的基线几乎完全重合。这说明加速这件事没有偷工减料,输出的内容质量没有打折。

速度上的收益则相当可观。三种草稿模型的平均接受长度（每次验证平均能接受几个候选词）分别达到 2.28、3.45、3.63，换算成端到端训练速度提升，分别是 1.50 倍、1.88 倍、1.83 倍。这里有个挺有意思的细节：DFlash 和 DSpark 用的都是"一次吐一整块候选词"的架构，明显比 EAGLE-3 那种逐词模拟多步的方式效果更好，猜测这可能和块状生成能更充分利用并行验证的空间有关。

研究团队还把这套系统推到了更大的模型上验证通用性，包括 Qwen3.5-35B-A3B、Qwen3.5-122B-A10B、Nemotron-3.5-Lightning-30B-A3B 和 GPT-OSS-120B，覆盖从 8B 到 122B 的参数规模。数据摆在这里最有说服力：

| 目标模型 | 草稿模型 | 接受长度 | 推理加速 | 端到端加速 |

|---|---|---|---|---|

| Qwen3-8B | EAGLE-3 | 2.28 | 1.63× | 1.50× |

| Qwen3-8B | DFlash | 3.45 | 2.23× | **1.88×** |

| Qwen3-8B | DSpark | 3.63 | 2.18× | 1.83× |

| Qwen3.5-35B-A3B | DFlash | 4.58 | 1.50× | 1.46× |

| Nemotron-3.5-Lightning-30B-A3B | DSpark | 2.65 | 1.19× | 1.16× |

| Qwen3.5-122B-A10B | DFlash | **4.78** | 1.72× | 1.35× |

| GPT-OSS-120B | DFlash | 3.80 | 1.48× | 1.19× |

有个现象值得多说两句：模型规模变大之后，接受长度虽然更高了，端到端加速比反而缩水了。比如 Qwen3.5-122B 接受长度冲到 4.78，是表里最高的，但端到端加速只有 1.35 倍，还不如 8B 模型上 DFlash 的 1.88 倍。论文的解释是，越大的模型（尤其是用了 MoE 混合专家架构的模型）每次验证时激活的专家计算量本身就更重，验证这一步的成本被摊薄的空间变小了，投机采样省下来的时间占比自然就缩水。这提醒我们一件事：投机采样带来的收益不是无条件的，它取决于目标模型本身验证一步要花多少代价，验证越贵，投机省下来的相对比例反而越小。

论文还专门测了一个更贴近真实业务场景的多轮任务：NeMo Gym Workplace Assistant，一个模拟办公室场景、需要调用邮件、日历、CRM 等工具的多轮智能体环境。结果显示接受长度依然随训练稳步上升，但端到端加速比只有 1.25 到 1.43 倍，明显低于单轮任务的 1.50 到 1.88 倍。原因也很直白：多轮任务里，工具调用和环境反馈本身要占掉将近一半的步骤时间（论文给出的具体数字是纯生成阶段只占整体步骤时间的 55.8%），这部分时间跟模型生成速度快慢完全无关，投机采样再快也加速不了等外部工具返回结果的过程。

这个发现挺现实的:再厉害的加速技术,也只能加速它能触及的那部分流程。如果瓶颈本身在别处,比如在等一个外部 API 返回,那投机采样的红利就会被稀释。这提醒做系统优化的人,得先搞清楚时间到底花在哪儿,而不是拿着一把锤子到处找钉子。

最后再看一眼流水线并行这部分的额外开销。论文的统计显示，草稿协同训练本身带来的额外时间成本，块状草稿模型（DFlash、DSpark）控制在 15% 以内，EAGLE-3 因为要做多步模拟（TTT），开销高一些，达到 34%。但因为投机采样本身把生成时间砍掉了 28% 到 52%，两相抵消下来，依然是净赚，DSpark 拿到了表里最高的端到端加速比 1.85 倍。而 TapChannel 那条特征传输侧路，每次训练迭代草稿阶段等待特征到达的时间只有 0.4 到 0.6 秒，占整个优化时间的 1.5% 到 2.2%，几乎可以忽略不计。

Q&A

Q1：TapChannel 是什么，它解决了什么问题？

A：TapChannel 是论文提出的一种特征传输机制，专门用于流水线并行训练场景。它让分布在不同流水线阶段的目标模型中间特征，能通过一条独立于标准流水线通信的侧路直接送到草稿模型所在阶段，不需要经过传统的逐层相邻传递，从而不改变原有流水线调度节奏。

Q2：草稿模型跟着大模型一起训练，会不会拖慢强化学习本身的效果？

A：不会。论文在 Qwen3-8B 上对比了带草稿协同训练和不带草稿训练的曲线，奖励值、验证准确率、训练推理 KL 散度几乎完全重合，说明加速过程没有损害训练质量，投机采样理论上也不改变最终输出分布。

Q3：为什么模型规模越大，投机采样带来的加速比反而变小了？

A：因为大模型尤其是 MoE 混合专家架构，每次验证时激活的专家计算量本身很重，验证成本高，投机采样省下的生成时间占总时间比例就变小了。比如 Qwen3.5-122B 接受长度高达 4.78，但端到端加速只有 1.35 倍，低于小模型的表现。