深入探究 vLLM 的投机采样训练支持:Speculators v0.3.0

阅读时间 11 分钟
Fynn Schmitt-Ulms, Helen Zhao, Rahul Tuli 和 Dipika Sikka (Red Hat AI 模型优化团队)

核心亮点

  • 推测解码是一种提升推理性能的优化技术。然而,为每个大语言模型(LLM)训练专属的草稿模型往往困难且耗时,同时目前缺乏用于生成生产级 vLLM 模型的可直接使用的训练工具。
  • Speculators v0.3.0 提供了端到端的 Eagle3 草稿模型训练支持,可与 vLLM 无缝衔接。
  • 训练支持包括使用 vLLM 进行离线数据生成,以及针对 MoE 和非 MoE 验证模型的单层和多层草稿模型训练能力。

大规模推理

在过去十年中,LLM 的规模和能力飞速扩张,同时也带来了对推理性能的严峻考验。由于 LLM 采用顺序方式生成 Token,每个 Token 都需要完整的数十亿参数前向传播,生成成本迅速攀升。随着模型规模的不断增长,这种顺序计算成为了主要的瓶颈,使得当今的 LLM 功能强大但运行缓慢。

缓解这一挑战的一个可行优化方案是推测解码。它允许较小的草稿模型提出 Token 预测,并由较大的模型进行快速验证,从而加速生成过程。

本文将深入探讨推测解码作为一种优化技术,介绍 Speculators 库,并深入剖析其最近发布的 v0.3.0 版本。Speculators 为研究人员、工程师和机器学习从业者提供了端到端生成推测解码模型并与 vLLM 无缝集成的工具。

什么是推测解码?

推测解码允许 LLM 在单次前向传播中生成多个 Token。其工作原理是利用一个小型“草稿”模型配合全尺寸的“验证”模型(即你试图部署的原始 LLM)。草稿模型成本低廉、运行迅速(通常只需一个 Transformer 块),它负责承担繁重的工作并自回归地预测多个 Token。验证模型并行处理这些 Token。对于每个 Token,验证模型决定是否接受草稿的预测。如果验证模型拒绝某个 Token,序列的其余部分将被丢弃;否则,这些 Token 将被包含在验证模型的响应中。

该方法具有以下优势:

  1. 最终响应的分布与单独使用验证模型完全一致,确保了使用推测解码时模型性能不会下降。
  2. 验证模型能够并行生成多个 Token。
  3. 由于草稿模型规模较小,运行它的额外开销通常极小。

总体而言,这可以使模型延迟降低 1.5 到 3 倍,从而显著加快生成速度。

在 vLLM 中使用推测解码模型

vLLM 和 Speculators 使得运行推测解码模型就像使用 vllm serve 服务任何其他模型一样简单。特别地,推测解码在低吞吐量场景下表现最佳,此时 GPU 未完全饱和,能够利用验证模型的并行 Token 生成能力。此外,草稿模型与验证模型的高度对齐至关重要,这也是为什么我们要为每个验证模型训练特定的草稿模型。然而,训练 LLM 专属的草稿模型可能很困难且耗时。幸运的是,Speculators 库简化了这一训练过程,使用户能够轻松生成并将其无缝集成到 vLLM 中。

创建新的草稿模型

目前推测解码算法的 SOTA(最先进)方案是 Eagle3 (Zhang et al., 2025)

Eagle3 草稿模型以验证模型的三个层级的隐藏状态作为输入,从而捕获验证模型的潜在特征。这些隐藏状态与 Token ID 结合,输入到较小的草稿模型中,进行自回归 Token 生成。

这意味着训练 Eagle3 草稿模型需要一个包含以下组件的样本序列数据集:

  1. 验证模型的隐藏状态(来自三个中间层)
  2. Token ID
  3. 损失掩码(用于仅针对模型响应进行训练,忽略用户提示词)
  4. 验证模型的输出概率(作为草稿模型的训练目标)

数据生成

直接从 vLLM 中提取这些值并非易事。幸运的是,Speculators v0.3.0 支持通过隐藏状态生成器进行离线训练数据生成,该生成器可以从标准 LLM 文本数据集中生成隐藏状态张量。这些张量随后被保存到磁盘,以便在训练过程中使用。

数据生成主要分为三个部分:预处理、隐藏状态生成和保存。

data_generation_overview
data_generation_overview

预处理阶段接收原始数据集:

  1. 重新格式化并规范化对话轮次
  2. 应用模型的对话模板
  3. 对对话进行 Token 化
  4. 基于助手响应范围计算损失掩码
  5. 将结果与 Token ID 一起保存到磁盘
  6. 收集 Token 频率统计数据,并保存到磁盘以备后用

损失掩码确保训练仅关注机器生成的 Token。对于通常只在最后响应中插入思考 Token 的推理模型,Speculators 提供了一个额外的标记来随机丢弃对话轮次,以确保模型在各种对话长度上进行训练。

隐藏状态生成器利用 vLLM 的插件系统,通过自定义的工作扩展(worker extension)运行。它通过给模型的前向传递过程打补丁,从而在预填充(prefill)阶段拦截并捕获中间隐藏状态。生成器使用 vLLM 的多进程执行器进行高效的批处理推理,并支持大模型所需的张量并行。此过程如下图所示。

hidden_state_generator
hidden_state_generator

在保存阶段,每个处理后的样本作为一个独立的 .pt 文件保存到磁盘,包含:

  • input_ids:Token 化的输入序列
  • hidden_states:每个被捕获层的张量列表
  • loss_mask:指示可训练 Token 的二进制掩码

生成器使用带有 ThreadPoolExecutor 的异步 I/O 来并行化磁盘写入操作,同时继续生成隐藏状态,从而最大化吞吐量。

除了数据文件外,还会额外保存两个文件:

  • data_config.json:包含有关数据生成的元数据
  • token_freq.pt:包含关于 Token 频率的信息

存储在 token_freq.pt 中的频率数据用于构建额外的目标转草稿(t2d)和草稿转目标(d2t)文件。这些文件充当验证模型完整词汇表与草稿模型较小词汇表之间的映射。这种精简的“草稿”词汇表仅包含最频繁出现的 Token,从而提高了草稿模型的效率。

以下脚本可用于开启离线数据生成:

训练

Speculators v0.3.0 支持训练 Eagle3 草稿模型。训练步骤接收来自上述步骤生成的样本和词汇表映射文件,连同模型配置信息,并初始化一个新的 Eagle3DraftModel 实例。该模型使用 Eagle3 作者引入的“训练时测试”(train-time-testing)技术进行训练。训练时测试在训练过程中模拟多步草稿采样过程,以确保模型不仅学会预测第一个 Token,还能学会后续的 Token。

flex_attention
flex_attention

来自 Eagle3 (Zhang et al., 2025) 论文的图示。

上图展示了训练时测试过程以及每一步的注意力掩码。对于每个前缀,草稿模型生成下一个 Token(蓝色)。然后,对于每个前缀加上第一步生成的步骤,模型生成第二个 Token(黄色),依此类推。

训练时测试的实现颇具挑战性,因为注意力掩码是稀疏的,而典型的注意力实现难以在计算和内存效率上处理这种情况。因此,Speculators 使用 FlexAttention (He et al., 2024) 进行注意力计算。FlexAttention 将注意力掩码分割为块,仅在非空区域计算注意力。结合 torch.compile,这不仅加速了计算,还极大地减少了反向传播所需的激活 VRAM。

批处理是任何训练实现中的另一个重要功能。由于序列长度通常不同,LLM 训练的样本批处理变得更加复杂。有两种解决方法:第一种是通过截断和填充使序列长度相同。这在统一长度的数据集上效果良好,但在需要大量填充的数据集上会导致计算浪费。相比之下,Speculators v0.3.0 使用第二种方法,即将序列沿“序列”维度连接起来,然后配置注意力掩码将其视为独立序列。这种方法与 FlexAttention 实现配合良好,尤其是在结合智能批采样算法以高效将样本打包至接近最大序列长度时,性能更佳。

综上所述,这些组件使得 Speculators Eagle3 模型训练既快速又节省内存,所有这些都可以通过单个 train.py 脚本来应用。

在 vLLM 中运行 Speculators 模型

训练完成后,该库会生成一个完整的模型制品,并附带一个扩展的 config.json 文件,其中包括 speculators_config。随后可以使用简单的 vllm serve 命令在 vLLM 中无缝运行模型:

vllm serve RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3

执行此命令时,vLLM 将读取存储在 speculators_config 中的推测解码设置(例如验证模型的名称)。该信息用于将草稿模型和验证模型同时加载到同一服务器中并设置推测解码。speculators_config 提供了标准化的配置格式,使得模型能够自包含并知晓其运行方式,从而使推测解码模型的部署像运行任何其他 LLM 一样简单。有关 speculators_config 的更多详情,请参阅下方的示例

虽然简化的单命令部署非常适合入门,但 vLLM 也提供了长格式语法,以便在需要更多控制时使用。这在以下场景下非常有用:

  • 使用与配置文件中不同的验证模型
  • 调整推测解码参数,例如推测 Token 的数量

长格式命令服务于基础(验证)模型,并通过 --speculative-config 标志指定推测器。这种灵活性对于实验和优化至关重要。例如,你可能希望换用验证模型的量化版本以进一步提升性能。

vllm serve RedHatAI/Qwen3-8B-FP8-dynamic \
  --tensor-parallel-size 1 \
  --gpu-memory-utilization 0.9 \
  --speculative-config '{"model": "RedHatAI/Qwen3-8B-speculator.eagle3", "num_speculative_tokens": 5, "method": "eagle3"}'

在此示例中,我们使用 FP8 量化的 Qwen3-8B 作为验证模型(取代 speculators_config 中引用的默认 BF16 版本),并将推测 Token 的数量从默认的 3 个增加到 5 个,以获得更高的潜在吞吐量。

vLLM 集成:生产就绪的推测解码

Speculators 与 vLLM 的紧密集成将推测解码从一种研究技术转变为生产级功能。vLLM 对 Eagle3 的支持使得在各种模型架构和配置下进行无缝部署成为可能。

vLLM 服务与 Speculators 训练支持::

  • Llama (3.1, 3.2, 3.3):8B 至 70B 参数
  • Qwen3:8B, 14B, 32B 参数
  • Qwen3 MoE:235B-A22B 参数(混合专家模型)
  • GPT-OSS:20B, 120B 参数

仅支持 vLLM 服务::

  • 多模态:Llama 4 视觉-语言模型

未来规划

Speculators 将重点关注以下下一阶段功能:

  • 在线数据生成(在训练时生成隐藏状态,无需中间磁盘缓存)
  • 对视觉语言模型的数据生成支持
  • 重新生成验证模型响应(用验证模型生成的响应替换数据集中的“助手”响应,以实现更优的对齐训练数据)

参与其中!

有兴趣深入了解推测解码吗?查看 Speculators 仓库,并通过浏览 Good First Issues 为仓库做出贡献!

如需更多资源、文档和 Slack 频道,请查看:

附录

Eagle3 算法

Eagle3 Algorithm
Eagle3 算法

speculators_config:

{
  "architectures": ["Eagle3Speculator"],
  "auto_map": {"": "eagle3.Eagle3SpeculatorConfig"},
  "Speculators_model_type": "eagle3",
  "Speculators_version": "0.3.0",
 
  "draft_vocab_size": 10000,
  "transformer_layer_config": {
    "num_hidden_layers": 1,
    "hidden_size": 4096,
    ...
  },
 
  "Speculators_config": {
    "algorithm": "eagle3",
    "proposal_methods": [{
      "proposal_type": "greedy",
      "speculative_tokens": 3,
      ...
    }],
    "verifier": {
      "name_or_path": "meta-llama/Llama-3.1-8B-Instruct",
      "architectures": ["LlamaForCausalLM"]
    }
  }
}

此配置将推测器定义为一个完整的模型,包含:

  • 模型标识
    • architectures:推测器的模型类(例如 Eagle3Speculator)
    • auto_map:用于 Hugging Face 兼容性的自定义模型加载
    • Speculators_model_type:具体的推测器实现
  • 草稿模型架构
    • transformer_layer_config:草稿模型 Transformer 层的完整规范
    • draft_vocab_size:用于高效草稿生成的简化词汇表大小(通常为 10k-32k 个 Token)
    • 模型特定的配置选项
  • 推测解码配置
    • algorithm:推测解码算法(EAGLE3)
    • proposal_methods:带有参数的 Token 生成策略
      • speculative_tokens:每步生成的草稿 Token 数量
      • verifier_accept_k:验证期间要考虑的前 k 个预测数量
      • accept_tolerance:接受草稿 Token 的概率阈值
    • verifier:要使用并进行验证的目标验证模型
      • name_or_path:HuggingFace 模型 ID 或本地路径
      • architectures:兼容性检查所需的验证模型架构