深入探究 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 验证器(Verifier)的单层及多层草稿模型的训练能力。

大规模推理

在过去十年中,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(State-of-the-Art)是 Eagle3 (Zhang et al., 2025)

Eagle3 草稿模型将验证器模型三个层的隐藏状态作为输入,以捕捉验证器的潜在特征。结合 Token ID,这些隐藏状态通过更小的草稿模型,自回归地生成草稿 Token。

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

  1. 验证器模型的隐藏状态(来自三个中间层)
  2. Token ID
  3. 损失掩码(Loss mask,用于仅在模型响应上进行训练,忽略用户输入提示)
  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 的插件系统,通过自定义工作进程扩展来实现。它通过修补模型的前向传递过程,在预填充(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。

任何训练实现中的另一个重要功能是批处理(Batching)。由于序列长度通常不同,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:用于兼容性检查的预期验证器架构