Papers

054 · Speculative Decoding Training

SpecForge:用推理引擎训练 EAGLE-3 draft model

SpecForge 怎样拆训练:冻结的大 target model 交给 SGLang 做高吞吐 prefill,待训练的小 draft model 则使用 FSDP/ZeRO-2。稀疏树 attention 和原地梯度进一步压低 EAGLE-3 Training-Time Test 的显存占用。

本文目录 · 18 节
  1. 1. 为什么 speculative decoding 还没有普及
  2. 2. baseline 的系统错配
  3. 3. Target–draft 解耦与混合并行
  4. 4. 2 个显存优化
  5. 5. 完整例子:Qwen3-235B-A22B 的 1 轮训练
  6. 6. 训练性能证据
  7. 7. SpecBundle:框架产出的 draft 是否真有用
  8. 8. 贡献、限制及与推测解码系统的关系
  9. 9. 推测解码的数学前置与接受长度
  10. 10. EAGLE-3 的特征、树与训练时测试
  11. 11. Monolithic DeepSpeed 为什么在大目标模型上失效
  12. 12. 1 轮分布式训练的数据流与同步
  13. 13. 稀疏树注意的内存推导
  14. 14. 原地损失反传逐式解释
  15. 15. 端到端训练表如何拆解
  16. 16. SpecBundle 数据配方与推理评测
  17. 17. 消融、开销与适用边界
  18. 18. 复现清单与结论边界

1. 为什么 speculative decoding 还没有普及

推测解码让小 draft 1 次提议 \(\gamma\) 个 token,大 target 1 次并行验证;接受前缀越长,target forward 次数越少,且拒绝时由 target 校正,理论上保持 target 的输出分布。瓶颈却转到 draft:许多 target 没有同家族小模型;公开 draft 常只用 ShareGPT/UltraChat 等小数据训练,数学、代码接受率差;EAGLE-3 需要 hybrid feature、dynamic tree 与 Training-Time Test(TTT),训练实现复杂。

论文给出直接例子:LLaMA-3.1-Instruct draft 用 ShareGPT 120K + UltraChat 200K 训练,在 Math500 的 acceptance length 为 2.82;换 1.4M conversation 的 Perfect-Blend 后达到 3.48,对应额外 1.17× inference speedup。只实现算法不足以形成可部署 draft。

2. baseline 的系统错配

官方 EAGLE-3 实现把 target 和 draft 包在同一个 DeepSpeed module。ZeRO-2 时冻结 target 参数仍在每卡完整复制,大模型很快 OOM;ZeRO-3 虽切参数,却在每次 target prefill 前频繁 all-gather,通信主导。更根本的问题是 target 只做 inference,draft 要 backprop:用同一个训练后端无法同时优化 2 种路径。

target prefill 是训练每个 batch 的固定大头,普通 training stack 缺少 chunked prefill、FlashAttention/FlashInfer、CUDA Graph 和成熟 MoE kernel。SpecForge 因此把“训练 speculative model”重写成 1 个 inference service 与 1 个 training worker 的协同问题。

3. Target–draft 解耦与混合并行

target 由 SGLang model runner 执行,沿用 tensor/expert/pipeline parallelism 和推理 kernel;draft 通常只有 target 的 3%–5%,只切 optimizer state 与 gradient,等价 ZeRO Stage 2,避免对小模型做昂贵 tensor/pipeline parallel。两者可 colocate 或 disaggregate,论文端到端实验采用 colocate。

  1. 1 个 batch 的完整 token 先送 target,SGLang 产生 EAGLE-3 所需 feature/logit。
  2. feature 传给每卡 draft replica,draft 预测多步候选。
  3. TTT 在训练时模拟多步 speculative rollout,让 draft 见到自身预测造成的分布偏移。
  4. 损失只反传 draft;target 冻结且留在 inference-optimized backend。

4. 2 个显存优化

Sparse Tree Attention

TTT length=7 时会同时维护多条候选树路径,naive SDPA 把 attention logits 保存为 activation;profile 中它占总 activation memory 的 80%。SpecForge 用 PyTorch FlexAttention 的 BlockMask 表达 causal prefix 与 TTT suffix 的可见关系,FlashAttention 式流计算不物化完整 score matrix,并跳过全空 block。每个 TTT step 只把新 K/V 放入 cache。

\[ M=M_{\mathrm{causal}}\lor M_{\mathrm{suffix}}. \]

\(M_{\mathrm{causal}}\) 允许 query 查看历史且 key 未 padding;\(M_{\mathrm{suffix}}\) 只允许同一 TTT 分支的对应位置相连。

In-place log-softmax backward

forward 完成后 logits 不再需要。自定义 Triton kernel 直接把梯度覆盖回 logits buffer,不另申请同尺寸 gradient tensor:

\[ \begin{aligned} s &= \sum(p\odot g),\\ \pi &= \operatorname{softmax}(z),\\ z &\leftarrow -\left(p\odot g-\pi s\right). \end{aligned} \]

最后一行将 \(\partial\mathcal{L}/\partial z\) 写回 \(z\) 的原内存。

该优化按 context length 和词表大小再降低 30%–40% 显存。二者作用不同:BlockMask 解决 attention activation,原地 backward 解决大词表 loss buffer。

5. 完整例子:Qwen3-235B-A22B 的 1 轮训练

8 张 H200 上,target 是 Qwen3-235B-A22B MoE,SGLang 用 TP=8 常驻 target;draft 用 ZeRO-2 分布训练。长度 4096 的 batch 到来,target 通过优化 prefill 产出 feature,draft 做 7 步 TTT。每步 BlockMask 只计算候选树真正可见的块,新 K/V 追加到 cache;loss backward 时覆盖 logits 存梯度。下一 batch 不需要 ZeRO-3 为冻结 target 反复 all-gather。

官方 EAGLE baseline 的 ZeRO-2 在这个规模 OOM,ZeRO-3 step time 11.2 秒、吞吐 2025.7 token/s;SpecForge step time 1.62 秒、吞吐 20227.2 token/s,即 9.99×。加速来自系统解耦,不代表 draft 算法本身改变。

6. 训练性能证据

所有端到端实验使用 8 张 NVIDIA H200、sequence length 4096,并为各系统选择可容纳的最大 batch。LLaMA-3.1-8B 即使 target 都不切分,SGLang prefill 仍让吞吐从 63015.4 提到 126639.6 token/s,2.01×。LLaMA-3.3-70B 从 ZeRO-3 的 14827.1 到 20608.8,1.39×。Qwen3-30B-A3B 从最佳 baseline 29257.1 到 126030.8,4.31×;235B MoE 达上述 9.99×。

后端消融中 SGLang 相对 Hugging Face/custom backend 最多 6.8×,说明 target prefill backend 是主因。attention microbenchmark 在 sequence length 4096、TTT=7、单张 H200 上,相对 SDPA 将执行时间降低 62.1%、peak memory 降低 93.5%,而且上下文越长差距越大。

7. SpecBundle:框架产出的 draft 是否真有用

作者用 Open-PerfectBlend 1.4M 对话,覆盖 chat、math、coding、instruction following;先以 temperature 0.8 用 target 重生成 assistant response,再从头训练 draft 2 个 epoch,learning rate 1e-4、cosine schedule。评测由 SGLang 在 H200 上运行,对 speculative steps、top-k、draft tokens 多组配置取最高吞吐。

通用表中 Llama-4-Scout + SpecBundle 在 FinanceQA 从 288.9 提到 1189.6 token/s,4.12×;Qwen-235B-A22B 在 MTBench 从 529.9 到 814.5,1.54×。数学/代码表中 LLaMA-3.1-8B 的 Math500 从 191.0 到 638.0,3.34×。跨全部实验,SpecBundle 相对无 speculative decoding 最多 4.48×,相对已有 draft 最多 1.35×。数据配比会造成 domain trade-off,某些通用任务并非每格都超过旧 draft。

8. 贡献、限制及与推测解码系统的关系

方法贡献:Target 按 inference workload 处理,draft 按 training workload 处理,两边因此可以使用各自成熟的后端。TTT 的 BlockMask 和原地 loss backward 让这套架构能够处理长 context;SpecBundle 还验证了训练产物可以进入 SGLang,评测没有停在训练吞吐。

局限:目前核心支持围绕 EAGLE-3,其他 draft 架构需适配 feature 与 loss。colocation 会让 target inference 和 draft training 争显存/带宽;disaggregation 虽支持,论文主实验未系统量化跨节点传 feature。最快结果依赖 H200、SGLang 和特定 kernel。SpecBundle 的最高吞吐从多组 speculative 参数中选出,部署仍需按模型与 workload 调参;接受率低时 draft 开销可能抵消收益。

合集关系:与 ReSpec 都涉及 speculative decoding,但 SpecForge 解决离线 draft training 与发布,ReSpec 优化 RL 系统中不断变化 policy 的 speculative rollout。和 DAWN/DSB 不同,它针对自回归 target,用 draft 提议并由 target 验证;diffusion LLM 的并行 mask 更新没有同一验证流程。

9. 推测解码的数学前置与接受长度

自回归大模型每生成 1 个 token 都要读取大量权重,单请求解码常受显存带宽限制。小草稿模型先连续提出若干候选,大目标模型把这些 token 作为 1 个序列并行计算概率;从左到右接受与目标分布一致的前缀,在第 1 个不接受处由目标模型采样纠正。1 次目标 forward 因而可能推进多个 token,同时标准校正规则保持目标模型分布。

设草稿每轮提议的 token 数为 \(\gamma\),单个位置在目标下被接受的平均概率为 \(\alpha\),草稿 1 次 forward 成本相对目标为 \(c\)。论文给出的理论分析把期望接受长度与 \(\alpha\) 的几何级数联系起来;\(\alpha\) 越高,验证 1 次平均推进越长,但 \(\gamma\) 增大也会增加草稿生成和目标验证的计算。实际速度还受 batch、并行度、树宽和 kernel 影响,不能只用接受率预测。

Acceptance length 是每次目标验证平均确认的 token 数,是 draft 质量的直接指标。它比草稿单独 next-token accuracy 更贴近部署:前几个位置错误会截断后续前缀,即使后面 token 单独正确也无法接受。EAGLE 系列不直接训练 1 个完整小语言模型,而是利用目标模型中间 feature 加轻量自回归 head,减少草稿难度并提高连续接受。

10. EAGLE-3 的特征、树与训练时测试

EAGLE-3 从目标模型多个层级取得混合 feature,让 draft 同时看到较低层词法信息和较高层语义信息,再预测目标 token。推理时 draft 可以按树展开多个候选分支,目标模型 1 次验证整棵树;分支数越多,覆盖目标选择的机会越高,也增加 attention 与显存。

普通 teacher forcing 训练时,draft 每一步都条件在真实历史;推理时它条件在自己上一轮预测,早期小错误会改变后续输入,产生暴露偏差。Training-Time Test 在 1 个训练 forward 中模拟多步草稿生成,把 draft 自己产生的 token/feature 接回后续步骤,使训练路径接近推理。长度 7 意味着同一 batch 要展开 7 步,保存多轮 attention、隐藏状态和 logits,这正是 SpecForge 显存问题的来源。

树中不同分支不能互相看见,否则某分支会利用实际推理不存在的兄弟 token。注意 mask 既要允许每个节点看 causal prefix,又只允许它沿自己的祖先链读取 suffix。论文的 BlockMask 算法用 query、key 索引、prefix 长度和序列真实长度判断属于 causal 区还是同分支对角关系,把逻辑编码成可跳块的稀疏掩码。

11. Monolithic DeepSpeed 为什么在大目标模型上失效

原实现把冻结 target 与可训练 draft 包成同一个模块,再统一使用 ZeRO。ZeRO-2 切 optimizer state 与 gradient,但 target 没有 optimizer/gradient,最占空间的冻结参数仍在每卡复制;7B 模型尚可,70B 或混合专家模型很快超显存。ZeRO-3 连参数也切,可容纳大 target,却要求每次 target forward 前把所需参数 all-gather,反复 prefill 的通信量巨大。

训练框架还把 target 当普通可反传模块构建执行图,即使关闭梯度,也缺少服务引擎针对只读权重的分页、CUDA Graph、专用 attention 和混合专家路由。Draft 相反,需要优化器状态、激活和 backward,却很小;给它套 target 所需的 tensor/pipeline parallel 会让通信超过计算。统一并行策略无法同时满足两边。

SpecForge 的解耦边界是 feature 接口:target runner 接收 token batch,输出 draft 所需 feature;draft trainer 消费 feature、执行 TTT、计算损失并反传。两边不共享 1 个分布式模块,所以 target 可用 SGLang 的张量、专家、流水并行,draft 用数据并行与 ZeRO-2。只要接口和模型层映射正确,新 target backend 不要求重写训练器。

12. 1 轮分布式训练的数据流与同步

  1. 数据加载器形成长度 4096 的对话 batch,并按 target tokenizer 得到 token、标签和 padding 信息。
  2. SGLang target runner 在多卡上执行 prefill,提取 EAGLE-3 所需的多层 feature 与目标 logits;target 全程冻结,不保留 backward activation。
  3. 每张 draft replica 收到对应样本 feature。草稿参数复制,optimizer state 与 gradient 按 ZeRO-2 切分,避免小模型做频繁参数聚合。
  4. Draft 做多步训练时测试。每一步新 key/value 写入局部 cache,BlockMask 保证分支可见性,FlexAttention 流式计算。
  5. 根据 target 分布形成监督,计算 masked log-softmax loss;自定义 kernel 在 backward 把 logits buffer 原地覆盖成梯度。
  6. 各数据并行 replica reduce-scatter gradient,更新参数,进入下一 batch。Target 仍驻留 SGLang,不随优化器 step 改变。

论文实验 colocate 两部分,即同一组 8 张 H200 同时放 target shard 与 draft replica。这样 feature 不跨网络,但两者竞争显存。架构也允许 disaggregate:target 独占推理卡,draft trainer 在另一组卡;此时 feature 传输可能成为瓶颈,论文没有把它作为主结果,因此不能假设分离一定更快。

13. 稀疏树注意的内存推导

Naive attention 先物化分数矩阵 \(S=QK^\top\),再 softmax 并与 \(V\) 相乘。序列长度为 \(n\) 时,每个头每层保存平方级分数;训练时为了 backward 还要保留或重算。TTT 把多个 suffix 分支拼接,虽然每个节点只能看祖先,dense 实现仍为不可见位置分配矩阵。论文 profile 中这些 attention logits 占 activation memory 的 80%。

FlexAttention 采用类似 FlashAttention 的分块流式计算:1 个 query block 只读取允许的 key block,在片上计算局部最大值、归一化和输出,不把完整分数写回显存。BlockMask 预先标出全空、全可见和部分可见块,全空直接跳过,部分块才执行元素级规则。它同时减少平方级中间量和无效计算,而不改变允许边的结果。

论文在单张 H200、训练时测试长度 7 的 microbenchmark 上,序列 4096 时相对原生 SDPA 把执行时间降低 62.1%、峰值内存降低 93.5%。这里是 attention 子算子结果,不等于端到端 9.99×;端到端还包含 target prefill、draft 前馈、loss 和通信。

14. 原地损失反传逐式解释

设 \(z\) 是词表 logits,\(p\) 是监督目标分布,\(g\) 是上游对逐 token loss 的梯度。Softmax 概率记作 \(\pi\)。标准实现会保留 \(z\),再申请同尺寸缓冲存对 \(z\) 的梯度;大词表、长序列和多步 TTT 下,这个缓冲很大。Forward 结束后若后续算子不再需要原 logits,就可复用其存储。

\[ \begin{aligned} s &= \sum(p\odot g),\\ \pi &= \operatorname{softmax}(z),\\ z &\leftarrow -\left(p\odot g-\pi s\right). \end{aligned} \]

\(s\) 是目标分布与上游梯度逐元素乘积之和,\(\pi\) 是由原 logits 计算的概率,最后一行将 \(\partial\mathcal{L}/\partial z\) 写回 \(z\) 的原内存。

自定义 Triton kernel 必须在覆盖前完成 softmax 所需的归约,否则前半块写梯度会破坏后半块读原值。实现按词表行做稳定最大值与指数和,得到 \(\pi\) 后再原地写。论文报告该优化按上下文与词表规模再节省 30%–40% 内存。它与树 attention 互补:前者压 loss/logit 缓冲,后者压 attention activation。

15. 端到端训练表如何拆解

8 张 H200、长度 4096 下,LLaMA-3.1-8B 的 baseline 每步 1.04 秒、吞吐 63015.4;SpecForge batch 可从 16 增到 64,step 2.07 秒但吞吐 126639.6,整体 2.01×。单步更慢不代表系统更差,因为 1 次处理样本更多,必须以 token 吞吐比较。

LLaMA-3.3-70B 的 ZeRO-2 baseline OOM,ZeRO-3 吞吐 14827.1;SpecForge 用 target TP 4、draft ZeRO-2 达到 20608.8,1.39×。Qwen3-30B-A3B 的混合专家模型,最佳 baseline 29257.1,SpecForge 126030.8,4.31×。

Qwen3-235B-A22B 时,ZeRO-2 baseline OOM,ZeRO-3 每步 11.2 秒、2025.7 token/s;SpecForge target TP 8,每步 1.62 秒、20227.2,9.99×。模型越大,避免 ZeRO-3 target 参数 all-gather 的收益越明显,所以最大倍数不能外推到小 target。

16. SpecBundle 数据配方与推理评测

框架快不代表 draft 好。作者用 1.4M 对话的 Open-PerfectBlend,覆盖聊天、数学、代码和指令;不直接沿用原 assistant 文本,而让对应 target 以温度 0.8 重生成,使训练答案来自实际目标分布。Draft 从头训练 2 个 epoch,学习率 1e-4 并用余弦退火。数据组成决定接受长度:缺少数学代码样本的旧 checkpoint 在这些域更容易早拒绝。

推理在 SGLang 与 H200 上运行。LLaMA-3.1-8B 并发 8,其余并发 16;对 speculative step、树宽和 draft token 数尝试多组配置,报告每个模型任务的最高吞吐。这个选择反映部署会调参,也使结果是可达上界,不是单一默认参数的无条件性能。

LLaMA-3.1-8B 在 Math500 从 191 到 638 token/s,3.34×;Llama-4-Scout 的 FinanceQA 从 288.9 到 1189.6,4.12×;整套结果最高相对无 draft 4.48×,相对已有 draft 1.35×。某些通用域旧 draft 可能接近或略优,说明平衡数据存在领域取舍。

17. 消融、开销与适用边界

Backend 消融在相同模型与并行配置下比较 Hugging Face、团队自写和 SGLang target runner,SGLang 最多快 6.8×;这隔离了推理优化 prefill 的价值。Attention microbenchmark 隔离 BlockMask/FlexAttention;端到端表再叠加混合并行、batch 容量和 backend。3 组结果不能互相相乘,因为收益重叠。

SpecForge 当前围绕 EAGLE-3 的 feature、TTT 和损失接口。换 Medusa、普通小模型 draft 或 n-gram,不一定需要同样训练图;框架的 target-draft 解耦思想可复用,kernel 与数据配方要重做。Target 新模型若 SGLang 尚不支持,可用 Hugging Face/custom backend,但论文已显示吞吐可能显著下降。

训练速度还依赖 target 与 draft colocate。目标模型占用大部分显存时,draft batch 仍受限;分离部署会增加 feature 通信。论文没有系统比较网络拓扑、容错、checkpoint 和多节点扩展。公开 SpecBundle 只覆盖给出的模型版本,target 权重更新后 feature 分布变化,旧 draft 未必兼容。

推测解码并非所有负载都加速。高并发服务可能已把 target 算力吃满,树验证增加计算;短输出无法摊销启动;域外请求接受长度下降,draft 成为额外开销。部署要按真实 prompt 分布测端到端延迟和吞吐,并允许低接受率时关闭 speculative mode。

18. 复现清单与结论边界

复现需固定 8 张 H200、序列 4096、同一数据与 EAGLE-3 超参数;分别测官方 baseline 的 ZeRO-2、ZeRO-3 并取最佳可运行值。SpecForge 要记录 target TP、draft ZeRO-2、最大 batch、step time、token 吞吐和峰值显存。Backend 实验保持并行度不变,kernel 实验固定 TTT 7 和相同 mask。

正确性测试应给小树手工列出每个 query 可见的 causal prefix 与同分支祖先,检查 BlockMask 不泄漏兄弟;原地 backward 与标准自动微分逐元素比较;target feature 在解耦前后做数值对齐。训练完成后在固定 SGLang 版本、并发和 speculative 参数网格上测 acceptance length、吞吐与目标输出质量。

先验证推测解码本身

在接入训练框架前,应以已有 draft 跑一套标准 target verification,确认拒绝位置、纠正 token 和随机采样与非推测 target 分布一致。Greedy 模式可以逐 token 比较输出;随机模式应使用相同随机数流或做统计分布检验。否则最终速度变化可能混入验证器错误,而不是 draft 质量。还要分别记录 draft proposal、target verification 和其他调度时间,接受长度高但 target 树 attention 很慢时,理论推进数不会变成真实加速。

理论公式中的相对成本 \(c\) 不是草稿参数量比例。1 个只有目标 3% 参数的 draft,若 kernel 小、batch 低、启动多,相对延迟可能远高于 3%;混合专家 target 又只激活部分专家。复现应直接测每轮 draft 与 target 时间,再结合平均接受长度估计速度,并用端到端结果校验。增加 \(\gamma\) 让几何级数的期望接受前缀增长趋于饱和,而草稿成本继续上升,所以论文推理评测搜索多组步骤和 token 数而不固定最大值。

手工检查 1 个 3 步训练时测试树

设 prefix 有 4 个真实 token,草稿第 1 步产生节点甲,第 2 步从甲分成乙、丙,第 3 步乙产生丁、丙产生戊。甲可以看 4 个 prefix;乙和丙都能看 prefix 与甲,但乙不能看丙、丙不能看乙;丁可看 prefix、甲、乙,不能看丙、戊。把这些节点按扁平序列放入 attention 时,普通 causal mask 会错误允许后排节点看前面恰好排列的兄弟。BlockMask 必须根据分支身份而非纯索引决定 suffix 可见性。

Algorithm 中的 modulo 与 diagonal 条件正是把重复展开的训练时测试步映射回同一分支位置。Padding 条件还要防止短样本的无效 token 被其他节点读取。可用极小张量将不可见位置 logits 设为显著标记,运行 attention 后检查输出不含该标记;再与显式 dense boolean mask 的结果对齐。性能 kernel 的正确性要先在这种小例子证明,不能只依赖 loss 最终下降。

Target feature 契约

EAGLE-3 draft 消费的不只是最终 logits,还包括目标模型指定层的 feature。不同模型家族的归一化、词嵌入共享、混合专家路由和层编号不同;SGLang runner 输出必须与官方训练实现使用的张量在形状、精度、位置编码和 token 对齐上相同。若错 1 位,draft 仍可能训练但学到错误条件,最终接受长度很低。框架应在一批固定输入上保存 monolithic 版 feature,与解耦 backend 逐层比较误差。

Target 采用张量并行时,某些 hidden state 是按维度切分的,需要 gather 成 draft 期望的完整 feature,或让 draft 接口理解切片;专家并行下 token 在不同 GPU 路由,输出还要恢复原 batch 顺序。论文使用成熟 SGLang runner 的价值不只是 kernel 快,也包含这些模型特定并行语义。自写 backend 能跑 forward 不等于能正确暴露训练 feature。

训练吞吐为什么要同时报告 batch 与 step

SpecForge 的内存优化允许更大 batch,表中 LLaMA-3.1-8B 从 16 提高到 64。若只报告 step time,2.07 秒看似慢于 1.04 秒;除以每步处理 token 后才得到 2× 吞吐。反过来,靠极大 batch 提高吞吐也可能改变优化动态,需要保持总 token、学习率缩放和有效 batch 可比,验证 draft acceptance 不因训练配方变化而下降。系统性能与模型质量要分别守恒。

大模型表还应注明哪个 baseline OOM。不能把 OOM 记为零吞吐再计算无限倍,而是对可运行的 ZeRO-3 基线计算 9.99×;ZeRO-2 OOM 说明扩展性失败,是定性证据。7B 与 70B 的加速差反映 target 参数通信占比,混合专家模型的 SGLang 优势还含专家 kernel 与路由实现。跨模型平均 1 个倍数会遮蔽这些原因。

数据生成和质量控制

用 target 重生成回答使训练文本与目标模型条件分布对齐,但温度 0.8 会引入多样性和错误。数据流水线应保存 prompt、目标回答、生成参数、停止原因和 tokenizer 版本,过滤截断、空回答与解析异常。聊天模板必须与部署一致;同一内容换 system token 或角色分隔符会改变 target feature,draft 接受率可能下降。

平衡 1.4M 对话并非只按样本数平均。代码和数学回答通常更长,对 token 总量贡献更大;训练 2 个 epoch 时长样本被截断的比例也影响领域。复现应报告各域样本数、token 数、长度分布和过滤规则,并分别在通用、数学、代码、金融任务测接受长度。总体吞吐提高可能由某一域主导,逐域结果才能判断生产流量是否匹配。

推理表中的吞吐与质量

表 2、3 报告同一 benchmark 运行时的 token 每秒和相对无推测 target 的倍数。推测验证设计上保持 target 输出分布,因此 draft 不应改变 benchmark 正确率;但实现错误、浮点差异或不同采样参数仍可能改变输出。复现需要同时跑任务分数,确认速度提升没有通过缩短回答、不同停止条件或错误接受获得。

对已有 draft 的比较要使用相同 target、SGLang 版本、并发与参数搜索预算。SpecBundle 在数学代码上明显领先,与 PerfectBlend 覆盖有关;若给新 draft 搜索 5 种树配置、旧 draft 只用默认值,会夸大差距。论文明确列出尝试的多组配置并取最佳,复现应对所有 draft 同样搜索并保存每组结果,而不是只公布胜者。

容错与 checkpoint

Target 冻结,checkpoint 只需保存 draft 参数、optimizer、scheduler、数据进度和随机状态,不应重复写数十亿级 target 权重。恢复时重新启动 SGLang target runner,再校验模型版本与 feature 契约。若 target service 某个 rank 失败,当前 batch 的 feature 不完整,trainer 应整体重试该 step,不能用部分旧 feature 更新。

Colocate 下 target runner 与 draft trainer 共用进程或设备,显存分配峰值必须协调。Target CUDA Graph 可能预留固定 buffer,draft activation 又随 batch 和 TTT 长度波动;启动顺序不同会导致碎片 OOM。复现实验要记录内存配置、图捕获 batch 和分配器参数,并在多步训练中监控峰值,而不是只看初始化后剩余显存。

草稿监督目标如何理解

草稿的目标不是独立回答用户,而是近似目标模型在给定前缀与内部特征下的下一步选择。训练标签来自目标 token/分布,损失只更新 draft。若用数据集原始答案而不经过目标模型重生成,答案可能是正确的,却不是该 target 在相同采样参数下会走的轨迹;draft 会学成另一个合理模型,验证时仍频繁被拒。SpecBundle 先让 target 重生成正是为了缩小这个条件分布差距。

多步 TTT 又让监督从单步近似扩展到自身 rollout。第 1 步 draft 预测偏离后,第 2 步输入已经不是训练语料真实前缀;算法仍用目标 feature/分布告诉它如何回到目标轨迹。训练时测试不是在训练中真正执行完整 rejection sampler,而是构造与多步草稿拓展相同的 attention 与输入结构,使 draft 对自己的历史鲁棒。

评价 draft 时要同时看接受长度和任务吞吐。1 个 draft 可能通过过度偏向高频 token 获得高单步命中,却在代码标识符处立即失败;平均接受长度按所有验证轮次计算,会暴露长前缀截断。逐域接受长度还能解释 SpecBundle 在数学代码的增益是否来自数据,而不是 inference 参数。

混合并行的通信账户

Target 张量并行每层需要 collective,但权重在服务启动后固定切分,不像 ZeRO-3 每次 forward 为参数 gather。专家并行只把 token 发给被路由专家,SGLang 已优化通信与内核。Draft 的 ZeRO-2 只在 backward reduce-scatter gradient、更新时处理 optimizer shard;它没有每层参数 gather。2 种通信发生在不同阶段,可以分别优化。

Monolithic ZeRO-3 的额外问题是 target 和 draft 被同一个 process group 与切分策略约束。Target 为容纳权重需要高切分度,draft 很小却跟着切,forward 中产生许多小 collective;或为了 draft 数据并行复制 target,造成 OOM。解耦后 target TP 可按模型容量选 1、4、8,draft 始终数据并行。表 1 中各模型 TP 不同,正是按规模选择,而非框架隐藏的统一默认。

若采用 disaggregate,通信账户会新增 target feature 传输。Feature 维度可能远大于 token 标识,且 TTT 每批都需要;可用流水重叠 target 下一批 prefill 与 draft 当前批训练,但要处理 backpressure 和 batch 对齐。复现时应把 colocate 与跨集群结果分开报告,并为跨集群路径记录 feature 字节数、网络时间和流水空洞。

为何 SGLang backend 影响如此大

训练循环中的 target 每个 batch 对整段序列做 1 次无梯度 prefill,本质就是高吞吐离线推理。SGLang 的 chunked prefill 控制峰值,FlashAttention/FlashInfer 减少中间读写,编译与 CUDA Graph 合并非 attention kernel,模型 runner 还包含混合专家的定制算子。Hugging Face 通用 forward 以模型正确性和易修改为主,未必有同等图捕获与内存管理。

后端最多 6.8× 并非所有模型相同。小 dense model 的通用 kernel 已较成熟,差距小;大 MoE 的路由、专家 GEMM 与通信更依赖专门实现,差距扩大。Hugging Face backend 在论文某些大 MoE 配置还遇到运行错误,这属于可用性而不只是速度。SpecForge 把 backend 设成可替换接口,使模型支持可随 SGLang 演进,而不是在训练仓库重复维护。

完整部署决策

团队应先确认目标模型在生产引擎中支持 EAGLE-3 feature 导出;若已有高质量 draft,直接评测,无需重新训练。没有 draft 时,收集与生产域匹配的 prompt,用 target 重生成回答;小规模训练验证 acceptance,再扩大到完整数据。训练吞吐优化只有在 draft 质量检查通过后才有意义。

上线前按真实并发、输入输出长度和采样温度搜索 speculative 配置,设置低接受率回退。监控每轮接受长度、draft 时间、target verification 时间和额外显存;模型版本、聊天模板或 tokenizer 更新后重新校准。SpecBundle 的公开权重降低启动成本,但不能替代目标流量验证。

许可证、权重来源和重生成数据治理也是框架外边界。Target 生成的 1.4M 回答可能含错误与敏感内容,训练前仍需过滤;draft 虽只用于加速,接受 token 最终会进入用户输出。标准验证保持目标分布,不会比 target 更安全,也不应绕过 target 的安全配置。

最终验收还要区分训练可扩展性与训练效率。8 卡结果说明单节点内 target 越大时解耦越有价值;扩到多节点后,ZeRO-2 gradient 通信、target feature 分发和数据读取都会增长,应逐节点报告吞吐缩放效率。长上下文超过 4096 时,稀疏树 attention 理论收益更大,位置编码、BlockMask 编译时间与缓存容量也要重新测。

与最终服务集成时,应确保训练与推理的 EAGLE-3 实现版本一致。树展开顺序、feature 层选择、归一化、草稿头结构或聊天模板任一变化,都可能让 checkpoint 能加载却接受率异常。发布 SpecBundle 需要同时保存 target 精确标识、训练提交、配置、tokenizer 和推荐参数。框架解决了制造工具,模型包的版本契约决定它能否被别人可靠复用。

因此完整复现报告至少同时给出:训练吞吐、峰值显存、最终训练 token 数、验证接受长度、各任务推理吞吐与任务得分。这组指标同时覆盖训练效率、草稿模型可用性和目标行为保持。

SpecForge 没有改变推测验证算法。冻结 target 使用成熟推理路径,可训练 draft 使用成熟训练路径;TTT 的树 attention 和大词表梯度再做专门的显存优化。SpecBundle 表明这条训练链能产出可部署的 draft。训练最高 9.99× 与推理最高 4.48× 是 2 类指标:前者降低训练 draft 的成本,后者衡量服务阶段的收益。

最终还应在相同硬件上重复多次训练吞吐与推理吞吐,报告波动和失败配置;这样 9.99× 与 4.48× 才有可审计的基线,而不是 1 次最佳运行。