SpecForge: Accelerating Speculative Decoding Training for SGLang

发表时间: 2025-07 · Blog post by The SGLang Team (lmsys.org)

原文: https://lmsys.org/blog/2025-07-25-spec-forge

作者/机构：The SGLang Team
发布日期：2025年7月25日

速读

一句话结论 针对投机解码中草稿模型训练工具匮乏且难以与推理引擎对齐的问题，本文开源了专为 SGLang 设计的 Eagle3 草稿模型训练框架 SpecForge，在 Llama 4 模型上实现了最高 2.18 倍的推理加速，并做到了从训练到部署的无缝衔接。

要解决什么问题 投机解码（Speculative decoding）是加速大语言模型推理的重要技术，但其核心环节——草稿模型（draft model）的训练，一直缺乏健壮的开源工具。现有的基于 Eagle3 算法的开源项目普遍存在维护不善、功能受限的问题，最致命的卡点是它们无法与 SGLang 等主流推理引擎兼容。这种训练与部署环境的割裂，导致模型训练完成后需要繁琐的适配才能上线，成为了实际部署的重大障碍。此外，当前的头部大模型（如 Llama 4、DeepSeek 等混合专家模型 MoE）结构日益复杂，传统的草稿模型训练代码在显存管理和分布式扩展性上无法支撑这种规模。缺乏能够处理这些复杂性并实现高效扩展的基础设施，极大地阻碍了前沿投机解码技术在工业界的实际落地，使得研究与实际部署之间存在巨大的鸿沟。

怎么做的 为了打通从研究到部署的壁垒，作者开发了 SpecForge 框架。其核心思路是构建一个原生集成 SGLang 的训练生态，使得草稿模型训练完成后能开箱即用地进行推理，无需任何额外适配。该框架的关键设计包含三个部分。首先是深度集成了 Eagle3 算法。Eagle3 通过训练一个轻量级的草稿模型来精准预测大目标模型的 token 分布，以此获得高接受率；其高性能依赖于训练时测试（Training-Time Test, TTT）架构，该架构通过模拟多步生成来增强草稿模型的鲁棒性。由于 TTT 需要处理特殊的注意力掩码（attention masks）和递归的数据循环，工程实现极具挑战，SpecForge 参照官方实现内置了 TTT 支持，封装了复杂逻辑以确保正确性。其次，框架设计了在线（Online）和离线（Offline）两种隐藏状态（hidden state）收集模式以适应不同硬件条件：在线模式在训练时实时生成目标模型的隐藏状态，适合快速迭代，不需要占用大量磁盘，但如果目标模型很大则需要更多 GPU；离线模式则提前计算并保存隐藏状态，训练时只需加载草稿模型，最低 1 张 GPU 即可运行，适合注重可复现性和数据复用的场景。最后，在底层工程实现上，框架提供了模块化接口以便于注册新的草稿和目标模型，原生支持复杂的 MoE 层和 Transformer 变体，并接入了 PyTorch 的完全分片数据并行（FSDP）和张量并行（TP），通过优化显存管理确保了超大底座模型草稿训练的可扩展性。

效果如何 实验基于 320K 条来自 ShareGPT 和 UltraChat 的样本数据集，使用 SpecForge 训练了 Llama 4 的 Scout 和 Maverick 两个版本的草稿模型。在 MT-Bench 基准测试中评估其推理加速效果。在具体设置上，测试固定了 SGLang 的 speculative-eagle-topk 参数为 8，speculative-num-draft-tokens 参数为 10，以此来确保树状注意力（tree attention）机制能够正常启用，并使用 SGLang 仓库中的 bench_speculative 脚本在不同配置下运行吞吐量基准测试以寻找最佳硬件性能。量化结果显示，Llama 4 Maverick 草稿模型在 MT-Bench 上实现了 2.18 倍的推理加速，而 Scout 版本也获得了 2.0 倍的加速比，证明了该框架训练出的模型能够直接胜任 Eagle3 的推理加速任务。作者也指出了该方案在离线模式下的存储代价，即需要极大的磁盘空间来预存隐藏状态（例如处理上述数据集需约 12TB 存储）。在未来的演进路线上，作者计划进一步支持 Kimi K2 和 Qwen-3 MoE 等更多模型架构，集成视觉语言模型（VLM），并开发更高效的并行策略和内核级优化。

1. 主要贡献

投机解码（Speculative decoding）是加速大语言模型（LLM）推理的强大技术。然而，当前生态中缺乏用于训练草稿模型（draft models）的稳健开源工具，这严重阻碍了该技术的实际应用与普及。为了解决这一核心痛点，SGLang团队开源了SpecForge——这是一个专为基于Eagle3的投机解码设计的全新训练框架。

SpecForge旨在提供极高的易用性，并与SGLang推理引擎进行了紧密的原生集成，实现了从模型训练到实际部署的无缝过渡。其核心创新点在于填补了学术研究与实际部署之间的空白，为复杂的前沿大模型（如Llama 4、DeepSeek及混合专家模型MoE）提供了全面支持、可扩展的分布式训练能力以及内存高效的训练方案。

2. 背景与设计原则

现有框架的局限性：尽管投机解码已成为加速大语言模型推理的突破性技术，但作为该过程关键组件的“草稿模型训练”却缺乏可靠的开源工具。现有的许多基于Eagle3的项目普遍面临维护不善、功能有限，或者与SGLang等主流推理框架不兼容的问题。这些固有的限制已经成为该技术在实际部署和广泛采用时的重大障碍。

SpecForge的设计初衷：为了弥合学术研究与实际部署之间的鸿沟，研究团队开发了SpecForge。这是一个专门构建的生态系统，用于训练能够与SGLang原生集成的草稿模型。其核心优势在于，一旦模型训练完成，即可开箱即用地直接进行推理，完全不需要进行任何额外的适配工作。同时，为当今的前沿大模型（如Llama 4、DeepSeek以及其他混合专家MoE模型）训练有效的草稿模型，需要能够应对其复杂性和庞大规模的底层基础设施。SpecForge从底层开始专门构建，正是为了满足这些严苛的需求。

SpecForge的核心能力：

• 对先进架构的原生支持：SpecForge全面支持最前沿的模型架构，其中包括复杂的混合专家（MoE）层以及各类Transformer变体。
• 可扩展的分布式训练：该框架集成了现代大规模训练策略，例如完全分片数据并行（Fully Sharded Data Parallel, FSDP）和张量并行（Tensor Parallelism, TP）。这使得SpecForge能够在GPU集群上实现高效的规模扩展。
• 内存高效的训练：通过采用优化的内存管理技术，SpecForge使得为超大规模的基础模型（base models）训练草稿模型成为现实。

3. 方法细节

Eagle3集成：Eagle是一种用于投机解码的最先进（state-of-the-art）方法，其核心目标是加速大语言模型的推理过程。该方法的工作原理是通过训练一个专门的、轻量级的草稿模型（draft model），使其能够准确预测更大规模目标模型（target model）的Token分布。这种机制能够带来极高的接受率（acceptance rates），从而实现显著的性能提升。

训练时测试（Training-time Test）支持：Eagle之所以能够实现如此高的性能，在很大程度上归功于其新颖的训练时测试（Training-Time Test, TTT）架构。TTT架构通过模拟多步生成（multi-step generation）过程，使得草稿模型具备极强的鲁棒性。然而，尽管TTT功能极其强大，但由于其需要使用专门的注意力掩码（attention masks）以及复杂的递归数据循环机制，导致其实现难度极高。为了解决这一痛点，SpecForge通过提供内置的TTT支持极大地简化了这种复杂性。在具体实现中，SpecForge严格参考了官方的Eagle3实现代码，以此确保算法的正确性并达到最优的性能表现。

两种训练模式（在线与离线）：为了简化隐藏状态（hidden state）的收集过程，SpecForge提供了两种功能多样的训练模式：在线模式（Online）和离线模式（Offline）。这种双模式设计确保了在不同的工作流中都能保持高度的灵活性，无论用户面临的是何种模型规模或硬件限制。

以下是两种训练模式的具体运行机制与对比：

• 在线模式（Online）：在该模式下，目标模型在训练期间被实时使用。其对磁盘空间的需求较低，但如果目标模型体积庞大，则需要调用更多的GPU资源。其核心工作原理是动态（on the fly）生成隐藏状态。在线模式非常适合需要敏捷性和最小化磁盘占用的场景，是进行快速迭代的理想选择。
• 离线模式（Offline）：在该模式下，目标模型仅在数据准备阶段被使用。其对磁盘空间的需求极高（例如，处理UltraChat和ShareGPT数据集大约需要占用12TB的存储空间），但对GPU的要求极低，最低仅需1张GPU即可运行（因为在训练阶段只需加载轻量级的草稿模型）。其核心工作原理是一次性预计算所有隐藏状态，并在后续训练中高效地重复使用这些数据。在存储空间充足的前提下，如果可重复性和数据复用是核心优先级，离线模式则是最佳选择。

对可扩展性与系统扩展性的优先级考量：SpecForge框架在设计之初就高度强调可扩展性（extensibility）和系统扩展性（scalability），以满足工程生产环境的严苛要求。通过提供模块化的接口，SpecForge允许开发者极其简便地实现并注册全新的草稿模型与目标模型。此外，为了支持超大规模的模型训练，SpecForge深度利用了PyTorch的FSDP（完全分片数据并行）技术，并无缝集成了张量并行（TP）机制，从而确保在多GPU集群环境中依然能够保持高效的训练吞吐量。

4. 实验环境

• 数据集规模及用途：使用了包含32万个样本（320K-sample）的数据集用于草稿模型的训练，数据来源于ShareGPT和UltraChat。
• 模型架构配置：基于Llama 4架构，训练了Scout和Maverick两个版本的Eagle3草稿模型（参数规模为17B）。
• 硬件配置：依托官方基础设施合作伙伴Voltage Park提供的计算基础，支持多GPU集群训练。
• 软件配置：深度集成SGLang推理引擎，底层训练依赖PyTorch框架，并启用了FSDP和TP分布式训练策略。测试与调优工具使用了SGLang仓库中的bench_speculative.py脚本。

5. 实验结果

研究团队使用SpecForge训练的模型在MT-Bench等基准测试中展现出了强大的性能，证明了其有效性以及对Eagle3推理的准备就绪状态。

• 性能加速比测试：团队在MT-Bench上评估了训练出的草稿模型。结果显示，Llama 4 Maverick草稿模型实现了2.18倍的速度提升；而Llama 4 Scout变体也实现了2.0倍的加速。这充分证明了SpecForge在不同模型变体上都能带来显著的推理性能增益。
• 草稿Token长度与参数调优评估：团队对Scout和Maverick在不同草稿Token长度下的表现进行了详细评估。在测试中，SGLang的speculative-num-steps参数（对应图表中的x轴步骤数）被作为变量进行测试。同时，为了确保能够启用树状注意力（tree attention）机制，测试严格固定了两个关键参数：SGLang的speculative-eagle-topk设置为8，speculative-num-draft-tokens设置为10。通过使用bench_speculative.py脚本运行不同配置下的吞吐量基准测试，能够有效帮助用户针对特定硬件调优出最佳性能。

6. 补充细节

代码与模型可用性：

• GitHub仓库：完整的训练框架源代码已开源，包含了训练时测试（TTT）和数据处理的具体实现细节。
• Hugging Face模型：预训练的Llama 4 Scout 17B（16E-Instruct-v1）和Maverick 17B（128E-Instruct-v1）的Eagle3草稿头（draft heads，仅包含草稿头，不包含完整模型）已开放下载，供开发者在实际项目中使用。

未来路线图（Roadmap）：
在不久的将来，SpecForge计划扩展以下核心支持：

• 支持更多模型架构：包括Kimi K2和Qwen-3 MoE。目前团队正与LinkedIn基础设施团队积极合作，训练额外的Qwen-3 MoE草稿模型，这些模型后续将由SpecForge提供原生支持。
• 集成视觉-语言模型（VLM）：计划将VLM的投机解码训练引入SpecForge框架。
• 提升训练效率：通过引入更先进的并行策略和底层内核优化（kernel optimization），进一步提升大规模训练的效率。

致谢与合作：
项目感谢了SGLang团队及社区成员、SafeAILab团队（因其在Eagle3算法上的开创性工作）以及美团（Meituan）的早期支持。特别强调了与官方基础设施合作伙伴Voltage Park的合作。Voltage Park为SpecForge提供了关键的算力基础，使得团队能够高效、可靠地训练和评估大规模投机解码模型。Voltage Park的首席产品与技术官Saurabh Giri表示，支持像SpecForge这样的高质量开源项目符合公共利益，对于推动高性能AI基础设施的民主化和技术进步至关重要。

7. 结论

SpecForge作为一个专为SGLang设计的Eagle3投机解码训练框架，成功填补了前沿大模型研究与实际部署之间的工具生态空白。通过提供内置的TTT（训练时测试）支持、灵活的在线/离线双训练模式，以及高度可扩展的分布式训练架构（FSDP与TP），SpecForge不仅大幅降低了训练复杂草稿模型（如MoE架构）的工程门槛，还在实际基准测试中实现了超过2倍的推理加速。未来，随着对更多模型架构（如Qwen-3 MoE、VLM）的支持以及底层内核优化的持续推进，SpecForge有望进一步推动开源大语言模型的高效推理与创新应用。