投机采样如何将 vLLM 性能提升高达 2.8 倍

2024年10月17日阅读时间 10 分钟

vLLM 团队

vLLM 中的投机采样是一种强大的技术，它通过结合小型模型和大型模型来加速 token 生成。在这篇博客中，我们将剖析 vLLM 中的投机采样、其工作原理以及它带来的性能提升。

本文内容基于我们的双周 vLLM 线上交流会（vLLM Office Hours），我们会在会上讨论优化 vLLM 性能的技术和更新。您可以在此处查看会议幻灯片。如果您更喜欢观看视频，可以在 YouTube 上观看完整录像。我们非常期待您参加未来的会议——请务必注册！

投机采样简介

投机采样（Leviathan 等人，2023）是降低大型语言模型（LLM）token 生成延迟的一项关键技术。该方法利用较小的模型处理简单的 token 预测，同时利用较大的模型来验证或调整这些预测。通过这种方式，投机采样无需牺牲精度即可加速生成，是一种无损且高效的 LLM 性能优化方法。

为什么投机采样可以降低延迟？传统上，LLM 以自回归方式一次生成一个 token。例如，给定一个提示词，模型生成三个 token T1、T2、T3，每个都需要单独的前向传递。投机采样改变了这一过程，允许在一次前向传递中提出并验证多个 token。

其工作流程如下

草稿模型（Draft Model）：一个更小、更高效的模型逐个提出候选 token。
目标模型验证（Target Model Verification）：更大的模型在一次前向传递中验证这些 token。它确认正确的 token 并纠正错误的 token。
一次传递生成多个 token：该方法不是每次传递生成一个 token，而是同时处理多个 token，从而降低了延迟。

通过采用这种方法，投机采样加速了 token 生成，使其成为小型和大型语言模型部署的有效方法。

投机采样在 vLLM 中是如何运作的

在 vLLM 中，投机采样与系统的连续批处理（continuous batching）架构相结合，其中不同的请求在单个批次中一起处理，从而实现了更高的吞吐量。vLLM 使用两个关键组件来实现这一点：

草稿运行器（Draft Runner）：该运行器负责执行较小的模型以提出候选 token。
目标运行器（Target Runner）：目标运行器通过运行较大的模型来验证这些 token。

vLLM 的系统经过优化，能够高效处理此过程，使投机采样能够与连续批处理无缝协作，从而提升整体系统性能。

为了在 vLLM 中实现投机采样，必须修改两个关键组件：

调度器（Scheduler）：调度器经过调整，可以在单次前向传递中处理多个 token 插槽，从而实现多个 token 的同步生成和验证。
内存管理器（Memory Manager）：内存管理器现在负责处理草稿模型和目标模型的 KV 缓存，确保投机采样过程中的平滑处理。

vLLM 支持的投机采样类型

vLLM 支持三种类型的投机采样，每种类型都针对不同的工作负载和性能需求进行了定制：

基于草稿模型的投机采样

这是最常用的投机采样形式，较小的模型预测下一个 token，较大的模型对其进行验证。一个常见的例子是使用 Llama 68M 模型为 Llama 2 70B 模型预测 token。此方法需要仔细选择草稿模型，以平衡精度和开销。

选择正确的草稿模型对于最大限度地提高投机采样的效率至关重要。草稿模型需要足够小，以避免产生明显的额外开销，但又要足够准确，以提供有意义的性能提升。

然而，选择合适的草稿模型可能具有挑战性。例如，在 Llama 3 等模型中，由于词汇表大小的差异，很难找到合适的草稿模型。投机采样要求草稿模型和目标模型共享相同的词汇表，在某些情况下，这会限制投机采样的使用。因此，在接下来的章节中，我们将介绍几种无需草稿模型的投机采样方法。

提示词查找采样（Prompt Lookup Decoding）

也称为 n-gram 匹配，这种方法对于摘要和问答等用例非常有效，因为这些场景中提示词和答案之间存在大量重叠。系统不是使用小模型来提出 token，而是根据提示词中已有的信息进行推测。当大模型在答案中重复提示词的部分内容时，这种方法效果特别好。

Medusa/Eagle/MLPSpeculator

在这种方法中，额外的层（或头）被添加到大模型本身，允许它在一次前向传递中预测多个 token。这减少了对单独草稿模型的需要，转而利用大模型自身的并行生成能力。尽管尚处于初步阶段，但随着更优化的内核（kernels）开发，该方法在提升效率方面展现出了潜力。

投机采样性能洞察：加速比与权衡

投机采样在低 QPS（每秒查询数）环境中提供了显著的性能优势。例如，在 ShareGPT 数据集的测试中，使用基于草稿模型的投机采样时，vLLM 的 token 生成速度提升了高达 1.5 倍。同样，在应用于 CNN/DailyMail 等摘要数据集时，提示词查找采样表现出了高达 2.8 倍的加速。

然而，在高 QPS 环境中，投机采样可能会引入性能权衡。提出和验证 token 所需的额外计算有时会减慢系统速度（当系统已处于计算瓶颈时），例如当每秒请求数增加时。在这种情况下，投机采样的开销可能会超过其收益，导致性能下降。

路线图：实现更高性能的动态调整

为了克服投机采样在高 QPS 设置中的局限性，vLLM 正在致力于实现动态投机采样。欢迎查阅相关论文了解详情。这也是 vLLM 目前的研究方向之一！此功能将允许 vLLM 根据系统负载和草稿模型的准确性来调整投机 token 的数量。简而言之，动态投机采样会在系统负载较高时缩短提议长度。然而，如下图所示，当平均 token 接受率较高时，缩短的影响较小。

未来，系统将能够自动修改每一步的投机程度，确保无论工作负载如何，投机采样始终是有益的。这将允许用户在启用投机采样时，无需担心它是否会拖慢系统。

如何在 vLLM 中使用投机采样

在 vLLM 中设置投机采样非常简单。启动 vLLM 服务器时，您只需添加必要的标志来指定投机模型、token 数量和张量并行大小即可。

以下代码配置 vLLM 以离线模式使用草稿模型进行投机采样，每次投机 5 个 token：

from vllm import LLM
 
llm = LLM(
    model="facebook/opt-6.7b",
    speculative_model="facebook/opt-125m",
    num_speculative_tokens=5,
)
outputs = llm.generate("The future of AI is")
 
for output in outputs:
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

以下代码配置 vLLM 使用投机采样，其中的提议是通过匹配提示词中的 n-gram 生成的：

from vllm import LLM
 
llm = LLM(
    model="facebook/opt-6.7b",
    speculative_model="[ngram]",
    num_speculative_tokens=5,
    ngram_prompt_lookup_max=4,
    ngram_prompt_lookup_min=1,
)
outputs = llm.generate("The future of AI is")
 
for output in outputs:
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

有时，您可能希望草稿模型以与目标模型不同的张量并行大小运行以提高效率。这允许草稿模型使用更少的资源并减少通信开销，将资源密集型计算留给目标模型。在 vLLM 中，您可以配置草稿模型使用张量并行大小为 1，而目标模型使用大小为 4，如下例所示。

from vllm import LLM
 
llm = LLM(
    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,
    speculative_model="ibm-fms/llama3-70b-accelerator",
    speculative_draft_tensor_parallel_size=1,
)
outputs = llm.generate("The future of AI is")
 
for output in outputs:
    print(f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}")

未来的更新（论文, RFC）将允许 vLLM 自动选择投机 token 的数量，从而无需手动配置，进一步简化流程。

请关注我们关于 vLLM 投机采样的文档以开始使用。参加我们的双周交流会来提出问题并反馈意见。

结论：vLLM 中投机采样的未来

vLLM 中的投机采样带来了巨大的性能提升，尤其是在低 QPS 环境中。随着动态调整机制的引入，它即使在高 QPS 设置下也将成为一种非常有效的工具，使其成为减少延迟和提高 LLM 推理效率的多功能必备特性。