菜单
主题
vLLM Logo

Speculators v0.5.0：DFlash 支持与在线训练

6分钟阅读
Fynn Schmitt-Ulms, Helen Zhao, Rahul Tuli 和 Dipika Sikka (Red Hat AI 模型优化团队)
#推测解码#生态系统

v0.5.0 版本为推测解码模型训练带来了显著的架构改进，引入了 DFlash 算法支持、完全统一的在线训练能力，并完成了向 vLLM 原生隐藏状态提取系统的全面迁移。此版本标志着推测解码工作流在训练灵活性和生产就绪度方面迈出了重要一步。

主要功能包括：

  • DFlash 算法支持 - 通过块扩散（block diffusion）实现单次传递草稿 token 生成
  • Gemma 4 DFlash 的效果
  • vLLM 原生在线及离线训练，支持统一的隐藏状态提取
  • 更新了文档和示例，概述了关键工作流

DFlash 算法支持

v0.5.0 引入了对 DFlash 推测解码算法的训练支持。与自回归 Eagle 3 模型相比，这是一种从根本上不同的草稿 token 生成方法。Eagle 3 通过多次前向传递以自回归方式生成草稿 token，而 DFlash 则利用块扩散技术在单次前向传递中生成所有草稿 token。

DFlash 的单次传递特性可以显著降低推测解码的开销，对于较长的草稿序列尤为明显。推测器为每个前缀生成一个长度为 B 的 token 块。这种块结构完全通过注意力掩码（attention mask）实现。与 Eagle 3 的另一个关键区别在于，DFlash 使用非因果注意力模式，即块内的查询（query）可以关注同一块内的所有其他 token。

在训练期间，多个预测块会并行进行训练。一种简单的方法是在序列中的每个可能点之后启动一个预测块。然而，对于长序列，这会导致注意力掩码变得极其庞大，使得训练在内存使用和计算成本上变得不切实际。为避免这种情况，我们不会在所有地方启动块。相反，我们从实际对训练损失有贡献的位置中随机选择一小部分“锚点”（anchor）位置。预测块仅连接到这些锚点。无论序列长度如何，这都保持了预测块的数量固定，从而使训练能够扩展到更长的上下文，同时保持注意力掩码的可控性。

训练 DFlash 推测器

训练 DFlash 模型遵循与 Eagle 3 类似的在线工作流。完整的教程可点击此处查看。

与 Eagle 3 的关键区别在于训练命令中特定的推测器参数，如下所示：

torchrun --standalone --nproc_per_node 2 scripts/train.py \
    --verifier-name-or-path "Qwen/Qwen3-8B" \
    --vllm-endpoint "https://:8000/v1" \
    --speculator-type dflash \
    --draft-vocab-size 8192 \
    --block-size 8 \
    --max-anchors 3072 \
    --num-layers 5 \
    --target-layer-ids "2 18 33" \
    --epochs 5 --lr 1e-4

DFlash 的特定参数包括：

--block-size # Number of tokens generated per diffusion block
--max-anchors # Maximum anchor points for speculation during training
--speculator-type # Must specify dflash

Gemma 4 DFlash 推测器

利用 DFlash 算法支持，我们训练了一个 Gemma 4 31B DFlash 推测器，并评估了其在多种任务类型中的接受率。结果显示，它在推理和代码生成任务上表现尤为强劲。
Figure 1: Gemma 4 DFlash acceptance rates across diverse task types.
图 1：Gemma 4 DFlash 在不同任务类型中的接受率。

Gemma 4 DFlash 实现了比 Eagle 3 和独立的 FP8 量化验证器更好的 token 间延迟。将 DFlash 与 FP8 量化验证器结合使用，可以获得更大的收益，如下所示：

Figure 2: Gemma 4 DFlash inter-token latency comparison.
图 2：Gemma 4 DFlash token 间延迟对比。

在 vLLM 中部署 DFlash 模型

DFlash 模型可与 vLLM 的推测解码基础设施无缝集成（自 PR #38300 起，该 PR 已包含在 vllm>=0.20.0 中）。

与 Eagle 3 模型类似，DFlash 模型在 config.json 中包含一个 speculators_config，其中包含有关目标模型、推测 token、推测算法名称等的详细信息。有了此配置，模型可以使用基本的 vllm serve 命令进行部署，如下所示。

vllm serve -tp 2 RedHatAI/gemma-4-31B-it-speculator.dflash

统一的在线与离线训练支持

v0.5.0 通过 vLLM 的隐藏状态提取系统（在 vLLM v0.18.0 中引入）为在线和离线训练模式添加了原生支持。Speculators 的先前版本使用 vLLM 的底层工具提取隐藏状态，要求将 vLLM 作为直接的 Python 依赖项。这种方法将训练流水线与 vLLM 的内部 API 紧密耦合，而这些 API 在 vLLM 版本更新之间经常变化，并需要手动与上游更改保持同步。此次集成移除了以往的自定义数据生成流水线，并消除了对 vLLM 作为直接 Python 依赖的需求。

现在，两种训练模式都使用相同的基于 vLLM 的提取路径：

  • 在线训练：在训练过程中动态提取隐藏状态
  • 离线训练：预先生成并缓存隐藏状态到磁盘，然后进行训练

通过利用 vLLM 原生的隐藏状态提取功能，Speculators 继承了 vLLM 的所有推理优化，包括高效的内存管理、批处理策略和硬件加速支持。现在的训练过程通过标准 REST API 与运行中的 vLLM 服务器通信，将训练基础设施与 vLLM 的内部实现细节解耦。这种架构转变提供了更好的版本稳定性，使团队能够更轻松地独立于 Speculators 训练框架更新 vLLM。

在线训练期间会发生什么：

  1. vLLM 服务器使用基础模型（以及一些特殊配置）初始化
  2. 训练提示（prompt）被发送到 vLLM 进行推理
  3. 提取隐藏状态并临时写入磁盘（或 RAM 磁盘）
  4. 训练进程加载提取的隐藏状态并删除文件
  5. 推测器模型基于提取的状态进行训练

本教程提供了有关在线训练工作流的更多信息。

离线数据生成也已更新，以使用与在线训练相同的隐藏状态提取系统和数据格式。我们开发了新的脚本，用请求填充正在运行的 vLLM 服务器并将其写入磁盘。这两种方法耦合得非常紧密，甚至可以将它们结合起来使用。例如，您可以部分离线生成隐藏状态，然后运行训练并加载现有的隐藏状态，同时生成任何缺失的状态。您还可以运行在线训练作业，在生成后不清除文件，从而可以在第一个 epoch 生成一次，然后在后续 epoch 加载这些文件。

本教程更详细地介绍了离线训练工作流。

添加了全面的文档

另一个值得强调的功能是更新后的文档站点。我们为 Speculators 支持的推测解码算法添加了简洁的介绍，并提供了训练推测器模型的详细分步教程。对于开发人员，我们还引入了一份指南，涵盖了如何向 Speculators 库添加新的推测解码算法，以及全面的 API 参考。