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. 为什么 speculative decoding 还没有普及
- 2. baseline 的系统错配
- 3. Target–draft 解耦与混合并行
- 4. 两个显存优化
- 5. 完整例子:Qwen3-235B-A22B 的一轮训练
- 6. 训练性能证据
- 7. SpecBundle:框架产出的 draft 是否真有用
- 8. 贡献、限制及与推测解码系统的关系
- 9. 推测解码的数学前置与接受长度
- 10. EAGLE-3 的特征、树与训练时测试
- 11. Monolithic DeepSpeed 为什么在大目标模型上失效
- 12. 一轮分布式训练的数据流与同步
- 13. 稀疏树注意的内存推导
- 14. 原地损失反传逐式解释
- 15. 端到端训练表如何拆解
- 16. SpecBundle 数据配方与推理评测
- 17. 消融、开销与适用边界
- 18. 复现清单与结论边界
1. 为什么 speculative decoding 还没有普及
推测解码让小 draft 一次提议 γ 个 token,大 target 一次并行验证;接受前缀越长,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:用同一个训练后端无法同时优化两种路径。
target prefill 是训练每个 batch 的固定大头,普通 training stack 缺少 chunked prefill、FlashAttention/FlashInfer、CUDA Graph 和成熟 MoE kernel。SpecForge 因此把“训练 speculative model”重写成一个 inference service 与一个 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。
- 一个 batch 的完整 token 先送 target,SGLang 产生 EAGLE-3 所需 feature/logit。
- feature 传给每卡 draft replica,draft 预测多步候选。
- TTT 在训练时模拟多步 speculative rollout,让 draft 见到自身预测造成的分布偏移。
- 损失只反传 draft;target 冻结且留在 inference-optimized backend。
4. 两个显存优化
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_causal OR M_suffix
M_causal: query 可看历史且 key 未 padding
M_suffix: 只允许同一 TTT 分支的对应位置相连
In-place log-softmax backward
forward 完成后 logits 不再需要。自定义 Triton kernel 直接把梯度覆盖回 logits buffer,不另申请同尺寸 gradient tensor:
s = Σ(p * g)
π = softmax(z)
z <- -(p*g - π*s) # z 原位变成 dL/dz
该优化按 context length 和词表大小再降低 30%–40% 显存。二者作用不同:BlockMask 解决 attention activation,原地 backward 解决大词表 loss buffer。
5. 完整例子:Qwen3-235B-A22B 的一轮训练
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 两个 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. 推测解码的数学前置与接受长度
自回归大模型每生成一个 token 都要读取大量权重,单请求解码常受显存带宽限制。小草稿模型先连续提出若干候选,大目标模型把这些 token 作为一个序列并行计算概率;从左到右接受与目标分布一致的前缀,在第一个不接受处由目标模型采样纠正。一次目标 forward 因而可能推进多个 token,同时标准校正规则保持目标模型分布。
设草稿每轮提议的 token 数为 γ,单个位置在目标下被接受的平均概率为 α,草稿一次 forward 成本相对目标为 c。论文给出的理论分析把期望接受长度与 α 的几何级数联系起来;α 越高,验证一次平均推进越长,但 γ 增大也会增加草稿生成和目标验证的计算。实际速度还受 batch、并行度、树宽和 kernel 影响,不能只用接受率预测。
Acceptance length 是每次目标验证平均确认的 token 数,是 draft 质量的直接指标。它比草稿单独 next-token accuracy 更贴近部署:前几个位置错误会截断后续前缀,即使后面 token 单独正确也无法接受。EAGLE 系列不直接训练一个完整小语言模型,而是利用目标模型中间 feature 加轻量自回归 head,减少草稿难度并提高连续接受。
10. EAGLE-3 的特征、树与训练时测试
EAGLE-3 从目标模型多个层级取得混合 feature,让 draft 同时看到较低层词法信息和较高层语义信息,再预测目标 token。推理时 draft 可以按树展开多个候选分支,目标模型一次验证整棵树;分支数越多,覆盖目标选择的机会越高,也增加 attention 与显存。
普通 teacher forcing 训练时,draft 每一步都条件在真实历史;推理时它条件在自己上一轮预测,早期小错误会改变后续输入,产生暴露偏差。Training-Time Test 在一个训练 forward 中模拟多步草稿生成,把 draft 自己产生的 token/feature 接回后续步骤,使训练路径接近推理。长度七意味着同一 batch 要展开七步,保存多轮 attention、隐藏状态和 logits,这正是 SpecForge 显存问题的来源。
树中不同分支不能互相看见,否则某分支会利用实际推理不存在的兄弟 token。注意 mask 既要允许每个节点看 causal prefix,又只允许它沿自己的祖先链读取 suffix。论文的 BlockMask 算法用 query、key 索引、prefix 长度和序列真实长度判断属于 causal 区还是同分支对角关系,把逻辑编码成可跳块的稀疏掩码。
11. Monolithic DeepSpeed 为什么在大目标模型上失效
原实现把冻结 target 与可训练 draft 包成同一个模块,再统一使用 ZeRO。二阶段切 optimizer state 与 gradient,但 target 没有 optimizer/gradient,最占空间的冻结参数仍在每卡复制;七十亿模型尚可,七百亿或混合专家模型很快超显存。三阶段连参数也切,可容纳大 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、计算损失并反传。两边不共享一个分布式模块,所以 target 可用 SGLang 的张量、专家、流水并行,draft 用数据并行与 ZeRO 二阶段。只要接口和模型层映射正确,新 target backend 不要求重写训练器。
12. 一轮分布式训练的数据流与同步
- 数据加载器形成长度四千零九十六的对话 batch,并按 target tokenizer 得到 token、标签和 padding 信息。
- SGLang target runner 在多卡上执行 prefill,提取 EAGLE-3 所需的多层 feature 与目标 logits;target 全程冻结,不保留 backward activation。
- 每张 draft replica 收到对应样本 feature。草稿参数复制,optimizer state 与 gradient 按 ZeRO 二阶段切分,避免小模型做频繁参数聚合。
- Draft 做多步训练时测试。每一步新 key/value 写入局部 cache,BlockMask 保证分支可见性,FlexAttention 流式计算。
- 根据 target 分布形成监督,计算 masked log-softmax loss;自定义 kernel 在 backward 把 logits buffer 原地覆盖成梯度。
- 各数据并行 replica reduce-scatter gradient,更新参数,进入下一 batch。Target 仍驻留 SGLang,不随优化器 step 改变。
论文实验 colocate 两部分,即同一组八张 H200 同时放 target shard 与 draft replica。这样 feature 不跨网络,但两者竞争显存。架构也允许 disaggregate:target 独占推理卡,draft trainer 在另一组卡;此时 feature 传输可能成为瓶颈,论文没有把它作为主结果,因此不能假设分离一定更快。
13. 稀疏树注意的内存推导
Naive attention 先物化分数矩阵 S=QK^T,再 softmax 并与 V 相乘。序列长度为 n 时,每个头每层保存平方级分数;训练时为了 backward 还要保留或重算。TTT 把多个 suffix 分支拼接,虽然每个节点只能看祖先,dense 实现仍为不可见位置分配矩阵。论文 profile 中这些 attention logits 占 activation memory 的八成。
FlexAttention 采用类似 FlashAttention 的分块流式计算:一个 query block 只读取允许的 key block,在片上计算局部最大值、归一化和输出,不把完整分数写回显存。BlockMask 预先标出全空、全可见和部分可见块,全空直接跳过,部分块才执行元素级规则。它同时减少平方级中间量和无效计算,而不改变允许边的结果。
论文在单张 H200、训练时测试长度七的 microbenchmark 上,序列四千零九十六时相对原生 SDPA 把执行时间降低六十二点一百分比、峰值内存降低九十三点五百分比。这里是 attention 子算子结果,不等于端到端九点九九倍;端到端还包含 target prefill、draft 前馈、loss 和通信。
14. 原地损失反传逐式解释
设 z 是词表 logits,p 是监督目标分布,g 是上游对逐 token loss 的梯度。Softmax 概率记作 π。标准实现会保留 z,再申请同尺寸缓冲存对 z 的梯度;大词表、长序列和多步 TTT 下,这个缓冲很大。Forward 结束后若后续算子不再需要原 logits,就可复用其存储。
s = sum(p * g)
pi = softmax(z)
z = -(p * g - pi * s)
s:目标分布与上游梯度逐元素乘积之和
pi:由原 logits 计算的概率
最后一行:把 dLoss/dz 写回 z 的原内存
自定义 Triton kernel 必须在覆盖前完成 softmax 所需的归约,否则前半块写梯度会破坏后半块读原值。实现按词表行做稳定最大值与指数和,得到 π 后再原地写。论文报告该优化按上下文与词表规模再节省三到四成内存。它与树 attention 互补:前者压 loss/logit 缓冲,后者压 attention activation。
15. 端到端训练表如何拆解
八张 H200、长度四千零九十六下,LLaMA 三点一八十亿的 baseline 每步一点零四秒、吞吐六万三千零一十五点四;SpecForge batch 可从十六增到六十四,step 两点零七秒但吞吐十二万六千六百三十九点六,整体二点零一倍。单步更慢不代表系统更差,因为一次处理样本更多,必须以 token 吞吐比较。
LLaMA 三点三七百亿的二阶段 baseline OOM,三阶段吞吐一万四千八百二十七点一;SpecForge 用 target TP 四、draft 二阶段达到二万零六百零八点八,一点三九倍。Qwen 三十亿激活三十亿的混合专家模型,最佳 baseline 二万九千二百五十七点一,SpecForge 十二万六千零三十点八,四点三一倍。
Qwen 二千三百五十亿激活二百二十亿时,二阶段 baseline OOM,三阶段每步十一点二秒、二千零二十五点七 token 每秒;SpecForge target TP 八,每步一点六二秒、二万零二百二十七点二,九点九九倍。模型越大,避免 ZeRO 三阶段 target 参数 all-gather 的收益越明显,所以最大倍数不能外推到小 target。
16. SpecBundle 数据配方与推理评测
框架快不代表 draft 好。作者用一百四十万对话的 Open-PerfectBlend,覆盖聊天、数学、代码和指令;不直接沿用原 assistant 文本,而让对应 target 以温度零点八重生成,使训练答案来自实际目标分布。Draft 从头训练两个 epoch,学习率十的负四次方并用余弦退火。数据组成决定接受长度:缺少数学代码样本的旧 checkpoint 在这些域更容易早拒绝。
推理在 SGLang 与 H200 上运行。小 LLaMA 八十亿并发八,其余并发十六;对 speculative step、树宽和 draft token 数尝试多组配置,报告每个模型任务的最高吞吐。这个选择反映部署会调参,也使结果是可达上界,不是单一默认参数的无条件性能。
LLaMA 八十亿在 Math500 从一百九十一到六百三十八 token 每秒,三点三四倍;Llama-4-Scout 的 FinanceQA 从二百八十八点九到一千一百八十九点六,四点一二倍;整套结果最高相对无 draft 四点四八倍,相对已有 draft 一点三五倍。某些通用域旧 draft 可能接近或略优,说明平衡数据存在领域取舍。
17. 消融、开销与适用边界
Backend 消融在相同模型与并行配置下比较 Hugging Face、团队自写和 SGLang target runner,SGLang 最多快六点八倍;这隔离了推理优化 prefill 的价值。Attention microbenchmark 隔离 BlockMask/FlexAttention;端到端表再叠加混合并行、batch 容量和 backend。三组结果不能互相相乘,因为收益重叠。
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. 复现清单与结论边界
复现需固定八张 H200、序列四千零九十六、同一数据与 EAGLE-3 超参数;分别测官方 baseline 的二、三阶段并取最佳可运行值。SpecForge 要记录 target TP、draft ZeRO 二阶段、最大 batch、step time、token 吞吐和峰值显存。Backend 实验保持并行度不变,kernel 实验固定 TTT 七和相同 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 不是草稿参数量比例。一个只有目标百分之三参数的 draft,若 kernel 小、batch 低、启动多,相对延迟可能远高于百分之三;混合专家 target 又只激活部分专家。复现应直接测每轮 draft 与 target 时间,再结合平均接受长度估计速度,并用端到端结果校验。增加 γ 让几何级数的期望接受前缀增长趋于饱和,而草稿成本继续上升,所以论文推理评测搜索多组步骤和 token 数而不固定最大值。
手工检查一个三步训练时测试树
设 prefix 有四个真实 token,草稿第一步产生节点甲,第二步从甲分成乙、丙,第三步乙产生丁、丙产生戊。甲可以看四个 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 对齐上相同。若错一位,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 八十亿从十六提高到六十四。若只报告 step time,二点零七秒看似慢于一点零四秒;除以每步处理 token 后才得到二倍吞吐。反过来,靠极大 batch 提高吞吐也可能改变优化动态,需要保持总 token、学习率缩放和有效 batch 可比,验证 draft acceptance 不因训练配方变化而下降。系统性能与模型质量要分别守恒。
大模型表还应注明哪个 baseline OOM。不能把 OOM 记为零吞吐再计算无限倍,而是对可运行的三阶段基线计算九点九九倍;二阶段 OOM 说明扩展性失败,是定性证据。七十亿与七百亿的加速差反映 target 参数通信占比,混合专家模型的 SGLang 优势还含专家 kernel 与路由实现。跨模型平均一个倍数会遮蔽这些原因。
数据生成和质量控制
用 target 重生成回答使训练文本与目标模型条件分布对齐,但温度零点八会引入多样性和错误。数据流水线应保存 prompt、目标回答、生成参数、停止原因和 tokenizer 版本,过滤截断、空回答与解析异常。聊天模板必须与部署一致;同一内容换 system token 或角色分隔符会改变 target feature,draft 接受率可能下降。
平衡一百四十万对话并非只按样本数平均。代码和数学回答通常更长,对 token 总量贡献更大;训练两个 epoch 时长样本被截断的比例也影响领域。复现应报告各域样本数、token 数、长度分布和过滤规则,并分别在通用、数学、代码、金融任务测接受长度。总体吞吐提高可能由某一域主导,逐域结果才能判断生产流量是否匹配。
推理表中的吞吐与质量
表二、三报告同一 benchmark 运行时的 token 每秒和相对无推测 target 的倍数。推测验证设计上保持 target 输出分布,因此 draft 不应改变 benchmark 正确率;但实现错误、浮点差异或不同采样参数仍可能改变输出。复现需要同时跑任务分数,确认速度提升没有通过缩短回答、不同停止条件或错误接受获得。
对已有 draft 的比较要使用相同 target、SGLang 版本、并发与参数搜索预算。SpecBundle 在数学代码上明显领先,与 PerfectBlend 覆盖有关;若给新 draft 搜索五种树配置、旧 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。第一步 draft 预测偏离后,第二步输入已经不是训练语料真实前缀;算法仍用目标 feature/分布告诉它如何回到目标轨迹。训练时测试不是在训练中真正执行完整 rejection sampler,而是构造与多步草稿拓展相同的 attention 与输入结构,使 draft 对自己的历史鲁棒。
评价 draft 时要同时看接受长度和任务吞吐。一个 draft 可能通过过度偏向高频 token 获得高单步命中,却在代码标识符处立即失败;平均接受长度按所有验证轮次计算,会暴露长前缀截断。逐域接受长度还能解释 SpecBundle 在数学代码的增益是否来自数据,而不是 inference 参数。
混合并行的通信账户
Target 张量并行每层需要 collective,但权重在服务启动后固定切分,不像 ZeRO 三阶段每次 forward 为参数 gather。专家并行只把 token 发给被路由专家,SGLang 已优化通信与内核。Draft 的二阶段只在 backward reduce-scatter gradient、更新时处理 optimizer shard;它没有每层参数 gather。两种通信发生在不同阶段,可以分别优化。
Monolithic 三阶段的额外问题是 target 和 draft 被同一个 process group 与切分策略约束。Target 为容纳权重需要高切分度,draft 很小却跟着切,forward 中产生许多小 collective;或为了 draft 数据并行复制 target,造成 OOM。解耦后 target TP 可按模型容量选一、四、八,draft 始终数据并行。表一中各模型 TP 不同,正是按规模选择,而非框架隐藏的统一默认。
若采用 disaggregate,通信账户会新增 target feature 传输。Feature 维度可能远大于 token 标识,且 TTT 每批都需要;可用流水重叠 target 下一批 prefill 与 draft 当前批训练,但要处理 backpressure 和 batch 对齐。论文架构允许这种部署,却未给出主表数字,因此复现不应把 colocate 结果标成跨集群结果。
为何 SGLang backend 影响如此大
训练循环中的 target 每个 batch 对整段序列做一次无梯度 prefill,本质就是高吞吐离线推理。SGLang 的 chunked prefill 控制峰值,FlashAttention/FlashInfer 减少中间读写,编译与 CUDA Graph 合并非 attention kernel,模型 runner 还包含混合专家的定制算子。Hugging Face 通用 forward 以模型正确性和易修改为主,未必有同等图捕获与内存管理。
后端最多六点八倍并非所有模型相同。小 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 生成的一百四十万回答可能含错误与敏感内容,训练前仍需过滤;draft 虽只用于加速,接受 token 最终会进入用户输出。标准验证保持目标分布,不会比 target 更安全,也不应绕过 target 的安全配置。
最终验收还要区分训练可扩展性与训练效率。八卡结果证明单节点内 target 越大时解耦越有价值,但没有给出多节点 draft 数据并行的网络扩展曲线。若扩到多节点,二阶段 gradient 通信、target feature 分发和数据读取都会增长;应逐节点报告吞吐缩放效率,而不是用单节点九点九九倍推断集群线性扩展。长上下文超过四千零九十六时,稀疏树 attention 理论收益更大,位置编码、BlockMask 编译时间与缓存容量也要重新测。
与最终服务集成时,应确保训练与推理的 EAGLE-3 实现版本一致。树展开顺序、feature 层选择、归一化、草稿头结构或聊天模板任一变化,都可能让 checkpoint 能加载却接受率异常。发布 SpecBundle 需要同时保存 target 精确标识、训练提交、配置、tokenizer 和推荐参数。框架解决了制造工具,模型包的版本契约决定它能否被别人可靠复用。
因此完整复现报告至少同时给出:训练吞吐、峰值显存、最终训练 token 数、验证接受长度、各任务推理吞吐与任务得分。只满足其中一项,无法证明系统既训练得快,又训练出了真正可用且保持目标行为的草稿模型。
SpecForge 没有改变推测验证算法。冻结 target 使用成熟推理路径,可训练 draft 使用成熟训练路径;TTT 的树 attention 和大词表梯度再做专门的显存优化。SpecBundle 表明这条训练链能产出可部署的 draft。训练最高九点九九倍与推理最高四点四八倍是两类指标:前者降低训练 draft 的成本,后者衡量服务阶段的收益。
最终还应在相同硬件上重复多次训练吞吐与推理吞吐,报告波动和失败配置;这样九点九九倍与四点四八倍才有可审计的基线,而不是一次最佳运行。