打通投机采样最后一公里!SGLang联合美团技术团队开源投机采样训练框架

机智流 2025-07-28 22:21

资讯配图

编辑:机智流编辑部

SGLang团队联合美团搜推平台、Cloudsway.AI开源 SpecForge — 首个支持超大模型投机采样训练并开箱即用的框架。框架基于 Eagle3 的投机采样训练框架,并与 SGLang 推理引擎深度集成。一键打通投机采样训练推理全流程!

新开源训练框架 — SpecForge

对于超大尺寸的模型,除了进一步优化算子之外。还有像投机采样这样的技术能加速它们的推理。投机采样(Speculative Sampling)通过引入一个轻量级的辅助模型(draft model)来预测多个候选 token(词汇单元),然后利用主模型(target model)并行验证这些候选 token,从而减少推理过程中的计算开销,同时确保结果的质量和正确性。

投机采样已成为大语言模型(LLM)推理加速的共识方案,但其端到端训练工具的缺失仍是明显短板。目前开源社区中,尚无能够支持超大尺寸模型训练且与 SGLang 深度结合的框架,而这些工具层面的不足,直接导致此类模型的部署面临巨大挑战。

资讯配图

为此,项目开发了 SpecForge,一个专为投机采样训练而生、并与 SGLang 原生集成的生态系统,其核心功能包括:

  • 原生支持最新开源架构: SpecForge 支持主流模型,包括复杂的 MoE 层和 Transformer 变体。

  • 可扩展的分布式训练: SpecForge 集成了 FSDP (Fully Sharded Data Parallel) 和 TP (Tensor Parallelism) 等并行策略,可在 GPU 集群上实现高效扩展。

  • 内存高效训练优化: 显著降低了大规模模型训练时的内存开销;即使是万亿参数的基础模型,训练Eagle3也同样高效。

项目代码和模型已开源,地址如下:

  • 💻 GitHub 仓库: 训练框架的完整源代码,包括 TTT 和数据处理的实现细节。https://github.com/sgl-project/SpecForge

  • 🤗 Hugging Face 模型: 下载 LLaMA 4 Scout 和 Maverick Eagle3 heads(不含完整模型)用于您的项目。

SpecForge 核心特性

目前性能强劲的投机采样技术分别有MTP和Eagle3,但MTP其需要在预训练阶段与基础模型一起训练,限制了MTP在业界的广泛应用。而Eagle3作为一种训练后而集成的技术,很适合在已开源的超大尺寸模型上继续训练,而进一步提升推理效率。

Eagle3 集成

Eagle 通过训练一个专门的轻量级草稿模型来准确预测较大目标模型的 token 分布,从而实现高接受率和显著的性能提升。下图展示了Eagle3的端到端训练流程:


训练时测试(TTT)支持

EAGLE3的高性能提升主要来源于其创新的训练时测试 (Training-Time Test) 架构,该架构通过模拟多步生成来增强草稿模型的健壮性。尽管 TTT 性能强大,但其实现却极具有挑战性,因为它依赖于复杂的专用注意力掩码(specialized attention masks)和递归式数据循环(recursive data loops)。SpecForge 将这一复杂过程完全封装,提供了内置且经过验证的 TTT 支持。我们的实现严格参照了官方 Eagle3 的核心逻辑,以确保其计算的正确性与性能,从而免除底层实现负担。

双重训练模式:在线与离线

SpecForge 通过提供两种训练模式:在线 (Online) 和离线 (Offline)来简化隐藏状态的收集,隐藏层的收集是Eagle类模型的特点。它通过主模型的隐藏层训练草稿模型,让草稿模型整体分布和主模型对齐。

我们框架的双模式设计能让用户能找到高效的训练模式,关于在线和离线训练的优缺点如下文所示。


选择在线或离线模式,可以根据您的具体需求和资源调整训练过程。

  • 在线模式:可实现最大速度和灵活性。它非常适合快速实验和存储有限的场景,因为它能动态生成数据,无需大量磁盘空间。

  • 离线模式:适用于可复现性和数据复用至关重要的场景。通过预先计算和存储隐藏状态,此模式可保证实验之间的一致性,在存储空间充足时效率很高。

扩展性优先

SpecForge在设计时高度重视可扩展性,以满足工程生产需求。该框架使用模块化接口实现了新草稿模型和主模型的直接实现和注册。

为了实现可扩展性,我们实现了多种训练时并行策略。包括FSDP (Fully Sharded Data Parallel) 和TP并行实现,确保超大型模型的高效训练。

后续开发计划

SpecForge的Roadmap如下:

  • 支持更多模型架构,包括 Kimi K2 和 Qwen-3 MoE。

  • 将视觉-语言模型 (VLM) 集成到 SpecForge 中。

  • 通过更好的并行策略和kernel优化来支持更高效的训练。

实验

为验证 SpecForge 的有效性,我们利用它在包含 320K 样本的 ShareGPT 和 UltraChat 数据集上,为 LLaMA 4 训练了 Scout 和 Maverick 草稿模型。

这些模型在 MT-Bench 等行业标准基准上表现出色,充分证明了其模型质量以及与 Eagle3 架构的兼容性。特别值得一提的是,我们为 Llama 4 Maverick 训练的草稿模型在 MT-Bench 上实现了 2.18 倍的推理加速。详细的实验结果与性能指标总结如下。

在下图所示的所有测试中,x 轴代表投机采样步长,对应于 SGLang 中的 speculative-num-steps。同时,我们将 SGLang 的 speculative-eagle-topk 固定为 8,将 speculative-num-draft-tokens 固定为 10,这样可以使用 tree attention达到更高的接受率。为了找到最优的投机采样参数,我们可以使用 SGLang 代码库中的 bench_speculative 脚本。该脚本会在不同配置下运行吞吐量基准测试,帮助我们针对硬件调优出最佳性能。



致谢

衷心感谢以下团队和合作者:

  • SGLang核心团队 — Shenggui Li、Shuai shi、Fan Yin、Yikai Zhu、Yi Zhang、Yingyi Huang、Yineng Zhang 及其他成员。

  • 美团搜推平台 — Chao Wang

  • SafeAILab团队 — Yuhui Li、Hongyang Zhang及其成员


-- 完 --


机智流推荐阅读

1. Trae SOLO打造ACL 25专区, 书生Qwen3等生成3000+论文解读,SOLO模式太爽了!!!!

2. Trae Solo Code即将大放量,聊聊Kiro Spec和Trae Solo的区别,分别适合什么人使用

3. AWS AI IDE「Kiro」深度体验:免费解锁最强编程模型

4. 聊聊大模型推理系统之 MIRAGE:当KV Cache不够时,把模型参数“变”成缓存



关注机智流并加入 AI 技术交流群,不仅能和来自大厂名校的 AI 开发者、爱好者一起进行技术交流,同时还有
HuggingFace每日精选论文顶会论文解读Talk分享通俗易懂的Agent知识与项目前沿AI科技资讯大模型实战教学活动等。
在「机智流」公众号后台回复下方标红内容即可加入对应群聊:
  • cc | 大模型技术交流群
  • hf | HuggingFace 高赞论文分享群
  • 具身 | 具身智能交流群
  • 硬件 | AI 硬件交流群
  • 智能体 | Agent 技术交流群

声明:内容取材于网络,仅代表作者观点,如有内容违规问题,请联系处理。 
开源
more
实测爆火的阶跃星辰Step 3,性能SOTA,开源多模态推理之王
被 OpenAI 点名后,智谱刚刚发布新开源模型:国产最全能,一句话造出搜索引擎
扣子开源全家桶,Apache 2.0加持,AI Agent又一次卷到起飞
阿里又开源了!电影级视频模型,业界首用MoE架构,通义App已上线
Github呼吁:开源需要资金
资讯 | 荣耀 MagicGUI 大模型发布并开源;全球首款三涵道穿戴式载人飞行器亮相
全能高手+科学明星,全球领先开源科学多模态大模型『书生』Intern-S1发布
特朗普AI计划全文曝光:拆掉监管、力挺开源,中美AI竞赛彻底点燃,算力、能源、数据中心统统加急!
RISC-V生态: 开源架构能否成为国产车芯的核心竞争力? | 盖世汽车2025第五届汽车芯片产业大会
抢跑GPT-5,智谱开源新SOTA模型,一句话搞出能看视频、发弹幕的B站!
Copyright © 2025 成都科技区角科技有限公司
蜀ICP备2025143415号-1
  
川公网安备51015602001305号