CANN/cann-recipes-train:SAM投机推理:长序列强化学习训练加速利器

发布时间:2026/7/11 6:19:26

CANN/cann-recipes-train:SAM投机推理:长序列强化学习训练加速利器 SAM投机推理长序列强化学习训练加速利器【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train针对大语言模型强化学习训练RL-training中的海量交互式采样耗时瓶颈传统的投机解码Speculative Decoding技术依赖辅助模型而且存在策略更新导致的分布漂移风险。本文提出在 RL 训练中引入一种基于后缀自动机SAM的无模型Model-Free投机解码方案该方法无需任何辅助模型利用 RL 数据如数学推理、代码生成中固有的结构化重复特性生成候选序列检索。本文结合自适应 batch 调度与向量化拒绝采样等工程优化在Qwen3系列模型32B/235B上完成了 SAM 投机解码在 RL 后训练场景的端到端验证。实践表明SAM 在保证精度严格无损的前提下显著降低了 Rollout 推理延迟特别是在长尾阶段获得了超过 35% 的加速收益相关代码已在cann-recipes-train 全部开源。1. 背景介绍1.1 SAM 原理SAMsuffix automaton后缀自动机是一个能够高效解决许多字符串问题的数据结构。直观上字符串的 SAM 可以理解为给定字符串的所有子串的压缩形式。SAM 主要维护两个重要的集合结束位置endpos考虑字符串 S 的任意非空子串 T我们记endpos(T)为字符串 S 中 T 的所有结束位置的集合。例如对于字符串 ABCBC 我们有endpos(BC){2,4}。在 SAM 中所有满足endpos集合相同的字串被归入同一个状态也被称为endpos等价类。endpos等价类的划分使得 SAM 可以以 O(|S|) 的空间复杂度存储所有子串信息。后缀链接link对于 SAM 中任意非初始状态 uu 中的最长子串为 w则link(u)指向的状态对应于 w 的后缀中与它的endpos集合不同且最长的那个即 u 所代表所有子串的最大真后缀所在状态。后缀链接构成了 SAM 的核心树形结构被称为后缀链接树它使得 SAM 可以在 O(|S|) 的线性时间内构造完成因为每次添加新字符时只需要通过后缀链接快速定位和更新相关状态。下图展示了基于字符串 ABCBC 构建的 SAM其中左图next是 SAM 的状态转移函数表示的是 SAM 中的不同状态如何按前序进行转移是 SAM 的自动机部分右图link则是上文所述的后缀链接是 SAM 的结构部分。SAM 采用线性复杂度的增量构造法每次向当前 SAM 中加入一个新字符时同时更新next和link单次更新的时间复杂度为 O(1)。SAM Decoding将 SAM 用于投机解码投机解码被转化为在后缀自动机中的后缀匹配使用 SAM 能够在 O(1) 的常数时间内高效实现这个过程。原文中使用了两种后缀自动机——静态 SAMS-SAM和动态 SAMD-SAM。静态 SAM 是基于已有的语料库构建后缀自动机将所有文本拼接成连续字符串并在每个文本片段末尾添加特殊符号如 EOS然后逐 token 遍历预处理后的语料通过添加 next 边和 link 边扩展节点。对于自动机的每一个状态我们记录该状态所对应的所有 endpos 等价字符串第一次出现的位置 (min_endpos)。动态 SAM 是完全基于当前处理的 prompt 建立一个动态的后缀自动机。​该自动机随生成过程动态扩展每生成一个新 token 并被 LLM 验证通过后立即更新 D-SAM若当前 token 可通过现有节点的 next 边转移则直接更新节点状态若无法转移则新建节点并添加 next 边和 link 边确保 D-SAM 始终包含最新生成的文本序列。对于 D-SAM可以在构建自动机的同时维护 min_endpos。在本文的实践中为了实现零迁移成本以及不引入额外系统维护复杂度我们使用了纯动态 SAM。基于 SAM-Decoding 的草稿draft生成从当前序列的最后一个 token 开始通过 SAM 的 link 边回溯直至找到能通过 next 边转移到下一个 token 的节点。 根据节点的 min_endpos在当前文本序列中定位匹配后缀的起始位置提取后续连续 token 作为候选序列。单次查询的复杂度为最坏 O(|S|) 、平均 O(1)。以下是一个具体的例子。还是以字符串 ABCBC 为例下图对应上文的 link 状态图为了简洁起见我们使用每个等价类中最长的子串来表示这个节点。当前的最后一个 token 是 C我们将 SAM 的最后一个节点沿着 link 边回溯到达节点 BC。该节点能够找到 next 边下图蓝色虚箭头转移到下一个 token停止回溯。该节点中的最长子串 BC 就是我们寻找的前缀子串中出现过的最长后缀因此我们将该节点的后续序列 BC 作为本次生成的 draft token。在生成草稿后我们会用大模型进行逐 token 验证若候选 token 符合大模型的预测分布则保留不符合就触发 “拒绝”—— 从拒绝位置开始由大模型重新生成这个过程就是拒绝采样。拒绝采样确保最终输出和大模型直接生成的结果完全一致实现 “无损加速”。假设有函数 $\alpha\in[0,1]$ 表示接受概率拒绝采样从 q(x) 采样一个样本 x从 U[0,1] 中采样一个随机数 $\epsilon$若 $\epsilon\leq\alpha(x)$ 则接受该样本否则拒绝并重新按照此流程采样。1.2 和其他投机方法的比较随着大语言模型推理加速需求的增长投机解码已从最初的“大模型小模型”范式演化出多种更为精细和高效的技术路线。为了更直观地理解当前业界的前沿方案我们将重点对比包含 SAM 投机的三种代表性方法​SAM-Decoding​基于检索的非参数化方法。​EAGLE-3​基于特征外推的轻量化 auxiliary head 方法。MTP​基于模型本身的自投机方法。这三种方法分别代表了投机解码的三个不同进化方向SAM 试图完全通过算法数据结构后缀自动机挖掘上下文中的重复模式以极低的计算代价实现加速EAGLE-3 则进一步压榨模型内部特征通过极其轻量的“寄生” head 在特征层级进行树状推测代表了辅助模型路线的 SOTA 水平而 MTP如 DeepSeek-V3 所采用则将投机能力内化为模型训练目标的一部分通过共享主干参数的辅助头实现原生投机。下表总结了这三种方法在草稿生成机制、架构需求及适用场景上的核心差异特性/算法SAMEAGLE-3MTP草稿生成方式借助后缀自动机进行检索从已生成的历史文本中检索“最长后缀 后续序列”作为草稿。在主模型的特征层之上利用轻量级 EAGLE head通常1层结合树状注意力Tree Attention进行特征外推生成草稿。在模型结构里加入多 token 预测头 (multi-token heads)这些预测头一次预测多个未来 token 作为草稿。额外组件无额外模型参数如果使用冷启动需要一个额外的语料库。需要训练一个依附于主模型的极小网络参数量通常 1%。采用训练时模拟测试 (Training-time Test)——在训练中直接模拟推理时的“猜测-验证”过程消除训练与推理的偏差让 draft 模型输出与主模型的更加一致。在模型结构里加入多 token 预测头 (multi-token heads)这些预测头一次预测多个未来 token 作为草稿。部署成本/难度较低即插即用无需对模型进行任何微调或附加训练。主要挑战在于工程上高效实现自动机的数据结构如果使用冷启动则需要一个覆盖面好的语料库。中等需要针对特定主模型训练对应的 EAGLE Head。虽然训练成本低但增加了模型管理的复杂度Base Model Adapter。较高需要在预训练阶段就加入 MTP 头外挂 MTP head 微调能让现有模型获得一定程度 MTP 能力但通常效果远不如原生 MTP 结构。适用场景在 prompt 复用率高或生成内容具有强结构性时表现极佳。适用于输入输出有高重叠度或上下文具有强重复模式的场景如数学、代码和部分 Agent 场景。取决于训练 EAGLE head 的数据在与训练数据分布相似的场景中有较好的表现。由于使用了原生训练的 MTP heads具有较好的泛化性能够应对更加多样化的使用场景。1.3 SAM 接受率我们使用 Qwen3-8B 模型在 Math500 数据集测试了 SAM-Decoding 的接受率具体的配置为投机 token 个数3采样参数temperature0repetition penalty1不设置 top-k 和 top-p 采样并行设置tensor parallel 4推理框架vLLM-Ascend注这里的接受率指的是 mean acceptance length包括了1个 bonus token。另外这里包括后文中的 SAM 使用的是动态 SAM不涉及额外语料库。可以看到在不使用冷启动的情况下接受率与序列长度呈正向相关。因此在长序列推理场景使用 SAM-Decoding 会带来不错的性能收益。我们设置的默认投机个数是3因为再加上1个本来就要解码的 bonus token 刚好是4个 token。此外随着投机个数增加带来的接受率提升也有边界效应不太建议将这个值设太大。1.4 SAM 在 RL 后训练场景落地的意义在LLM的强化学习后训练场景中模型需要进行大量的交互式采样来收集或评估数据。这个过程通常涉及数百万次的生成迭代成为推理成本的主要瓶颈。SAM-Decoding 因其无模型Model-Free和基于检索Retrieval-Based的特性为这一关键阶段的加速提供了一种轻量化、易落地的实践路径。首先作为一种非参数化的投机采样方法SAM-Decoding不需要额外的模型参数这节省了专门训练 draft model 或者 auxiliary head 的成本同时也不必担心模型权重更新导致的分布漂移问题。并且在训练时不需要进行额外的内存管理这对于训推共卡的 Colocate 框架尤为重要对于现有系统实现了即插即用的加速。其次SAM-Decoding 对强化学习数据集具有天然的​亲和性​。在许多 RL 任务中生成的 Context 中往往包含大量的​重复或结构化元素​​。例如在 RLVR Reinforcement Learning with Verifiable Rewards场景模型被训练来生成详细、结构化的推理过程这种推理过程在数学问题求解中表现为重复的公式结构和计算步骤在代码生成中表现为高度重复的语法模式和库调用在 Agent 规划任务中表现为结构化的行动指令或观察报告。SAM-Decoding 的后缀自动机设计能够高效捕捉这些​局部重复和结构化元素​​。它能快速地从历史生成中检索最长匹配后缀作为高质量草稿从而显著提高 draft token 的​接受率​。这种能力对于需要大量生成结构化、长序列推理轨迹的 RL 后训练流程来说至关重要能有效降低高昂的采样成本加速模型在专业领域如数学、编程和自主智能体的迭代效率。月之暗面团队近期提出的 Seer 强化学习框架中也使用了基于压缩后缀树的 Model-Free 投机解码方法利用 RL 算法如 GRPO中 “同一 prompt 生成多个响应” 的分组特性通过在线上下文学习挖掘组内响应的模式相似性实现高效、高准确率的投机解码与本文的 SAM 有异曲同工之妙。接下来本文将通过 Qwen3 系列模型长序列 RL 训练实践详细展开 SAM 投机解码在RL后训练场景的精度与性能验证结论。2. RL 训练精度验证虽然拒绝采样在理论上保证了 SAM 投机的无损性但在实践中精度还受其他各种复杂因素的影响。为了确保 SAM 在 RL 场景的精度正确性我们进行了包括算子校验、下游任务评测和端到端训练的验证工作。2.1 算子精度首先是算子的精度验证。普通解码和投机解码使用的是不同的算子在 vLLM-Ascend 中我们使用torch_npu._npu_paged_attention进行 token-by-token 的解码使用torch_npu._npu_paged_attention_splitfuse进行投机解码生成 draft tokens 的 logprobs。为了验证这两个算子的等价性我们统计了对相同的输入两个算子对各个 token 预测的 logprobs 的差值结果如下所示可以看到两者之差严格为0这表明普通解码和投机解码使用的算子之间没有精度差异。2.2 下游评测我们基于 Qwen3-32B 模型在数学数据集 GSM8k 以及 AIME2024 进行了纯推理的测评pass1 采样参数为 temperature0.6repetition penalty1不设置 top-k 和 top-p。 测试结果如下数据集不开启SAM开启SAMGSM8k88.988.7AIME202466.566.8可以发现开启 SAM 投机并不影响这些推理任务的准确率。2.3 端到端后训练最后我们通过端到端的 RL 训练直接检验 SAM 的精度。我们使用 DAPO 算法基于verl在 DAPO-MATH-17k 数据集上训练 Qwen3-32B 模型我们固定了 vLLM 的随机种子开启和不开启 SAM 每步的平均回复长度曲线和 reward 曲线如下所示可以看到开启 SAM 的训练曲线在数值和趋势上都与不开启 SAM 的基线基本吻合结合之前的算子精度和下游评测的验证可以断定 SAM 投机解码在 RL 后训练场景的精度是无损的。3. RL 训练性能验证3.1 性能优化投机解码的性能取决于两个因素第一是接受率第二是投机算法本身带来的额外耗时draft token 的生成和验证、拒绝采样。为了能够最大化 SAM 无损投机在RL训练中的收益我们需要尽可能地提高接受率并减少额外耗时。为此我们做了如下两大优化——关于 batch size 的自适应开关以及拒绝采样加速相关代码已在cann-recipes-train 全部开源。3.1.1 自适应开关我们在vllm_ascend/worker/model_runner_v1.py中适配了无损投机关于 batch size 的自适应开关当开关开启时只有当 vLLM engine 当前处理的请求数量小于等于阈值投机解码才会生效。设置投机自适应开关既是为了减少额外耗时也能够提高接受率。首先从减少额外耗时角度考虑。我们当前使用的对draft token进行验证的 Page Attention 算子在输入shapebatch size sequence length增加时会有比较严重的劣化。在大 batch size 长序列的输入下draft token 验证带来的额外耗时会超过投机本身的收益。因此我们希望在小 batch size 的情况下才开启 SAM 投机以平衡加速收益与计算开销。其次是提高接受率的角度。由于我们目前使用的是动态构建 SAM 的模式即把当前序列当作后缀自动机的全部语料在 Rollout 初期后缀树的体量较小因此接受率也就较低这一点在上文 1.3 章节中也有体现。在 Rollout 后期由于在同步 RL 框架下存在 Rollout 生成长度不均的问题请求排队和抢占机制会导致长输出请求的调度延迟从而导致少数请求主导最终执行阶段。我们把单步处理的请求总数当作 Rollout 进程的指示物来提示是否开启 SAM 投机。这样一来可以在长尾请求中得到显著加速。3.1.2 拒绝采样加速拒绝采样是投机推理引入的额外操作想要最大化投机推理的收益也需要对这个模块的性能做尽可能的优化。vllm-ascend原生的投机推理实现性能较差单个推理 step 会带来超过50ms的额外耗时完全掩盖了投机推理带来的性能收益。我们使用了纯 Pytorch 层面的代码优化将拒绝采样的耗时减少到原本的十分之一左右。本次我们主要优化了vllm_ascend/sample/rejection_sampler.py中的三个函数expand_batch_to_tokens把采样参数temperature、top-k、top-p从“每个请求一个值”的 batch 向量展开成“每个 draft token 一个值”的 token 向量使后续拒绝采样能逐 token 正确应用这些参数。sample_recovered_tokens_pytorch在拒绝采样中为每个被否决的草稿位置即时生成一个 recovered token保证序列能继续推进。rejection_random_sample_pytorch按接受-拒绝规则逐个比对草稿与目标概率决定用 draft token、recovered token 还是终止并在全接受时追加 bonus token最终拼出随机采样模式下的输出序列。以sample_recovered_tokens_pytorch为例原生实现中含有嵌套的for循环操作。这种逐个元素的操作方式无法利用NPU的并行计算能力并且在循环内部反复创建新张量如.clone(),torch.full()导致效率极低。我们优化的核心思想是​**向量化 (Vectorization)**​去掉所有Python循环使用PyTorch的张量操作来一次性处理所有数据。我们使用张量操作消除了循环内重复的张量创建并通过高效的索引和广播实现了关键变量的计算。具体实现参考我们开源的代码。后续可以使用 Triton 对这个模块做进一步的优化。3.2 RL 训练实测收益我们在 DAPO 长序列 RL 场景测试了 SAM 投机解码的加速收益。至于模型选择我们选用了 Qwen3-32B dense 模型以及 Qwen3-235B-A22B MoE 模型这两个比较有代表性的模型。下文介绍具体的配置和性能数据。3.2.1 Qwen3-32B相关配置数据集 DAPO-MATH-17k最大回复长度max_response_length34816训练 batch 大小train_batch_size32生成 batch 大小gen_batch_size96Rollout 最大请求数max_num_seqs128Rollout 模型张量并行tensor_model_parallel_size8SAM相关配置自适应 batch size 开关阈值为8投机 token 个数为3个。910C卡数2机16卡32die运行模式eager性能数据如下表所示其中“单轮平均推理时间”指每个 DP Data Parallel域的单轮DAPO 算法由于其动态过滤机制可能会有多轮生成推理平均耗时“单步总推理时间”指单个 step 的 Rollout 总耗时单轮平均推理时间 / s单步总推理时间 / s单步总时间 / s不开启SAM3904.224159.064730.86开启SAM3548.623793.294374.99收益10.02%9.64%8.13%3.2.2 Qwen3-235B-A22B相关配置数据集 DAPO-MATH-17k最大回复长度max_response_length34816训练 batch 大小train_batch_size128生成 batch 大小gen_btach_size128Rollout 最大请求数max_num_seqs64Rollout 模型张量并行tensor_model_parallel_size4SAM相关配置自适应 batch size 开关阈值为8投机 token 个数为3个。910C卡数8机64卡128die运行模式eager性能数据单轮平均推理时间 / s单步总推理时间 / s单步总时间 / s不开启SAM7102.9214287.5415441.45开启SAM6467.5212811.9813960.98收益9.82%11.52%10.60%3.2.3 客户场景验证我们在客户的训练场景部署了 SAM 投机解码帮助达成了 A3 集群 128die Qwen3-32BverlDAPO 强化学习性能提升的目标。相关配置最大回复长度max_response_length34816训练 batch 大小train_batch_size512生成 batch 大小gen_btach_size1536Rollout 最大请求数max_num_seqs256Rollout 模型张量并行tensor_model_parallel_size8SAM相关配置自适应 batch size 开关阈值为4投机 token 个数为4个。910C卡数8机64卡128die运行模式eager具体性能数据如下单步总推理时间 / s单步总时间 / s不开启SAM909411870开启SAM822311000收益10.59%7.91%3.2.4 Rollout 长尾分析此外我们特别统计了 SAM 投机在 Rollout 长尾阶段当前处理的请求数在投机阈值以下即下图中紫色部分的表现性能达到非投机1.35X。如果后续对算子进行优化支持更大投机阈值的话这一收益将更加明显。4. 总结与展望SAM-Decoding 的成功落地证明了在 LLM 的训推加速领域传统算法和数据结构层面的创新如后缀自动机可以作为​模型参数方法强有力的补充。它为 LLM 在长序列、高并发、训推共存的复杂 RL 场景中提供了一个即插即用并且易于集成的加速方案。这对于迫切需要降低海量交互采样成本的 RL 后训练流程具有重大的实践意义。后续我们会探索更高效的投机策略提高接受率以及尝试与图模式结合提高算子执行效率。通过持续的技术深耕和工程优化我们相信以 SAM 投机解码为代表的 Model-Free 投机技术会在 LLM 的强化学习训练中发挥更大的价值。【免费下载链接】cann-recipes-train本项目针对LLM与多模态模型训练业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-train创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考

相关新闻