使用 Dataflow 和 vLLM 对 Gemma 进行推理

十一月 13, 2024

Danny McCormick Senior Software Engineer Dataflow

像 Gemma 这样的大型语言模型 (LLM) 功能强大且用途广泛。它们可以翻译语言，撰写不同类型的文本内容，并为您的问题给出信息丰富的答案。然而，将这些 LLM 部署到生产环境（特别是用于流式传输用例）仍可能会是一个重大挑战。

本文将探讨如何使用两种最先进的工具 vLLM 和 Dataflow，以最少的代码高效地大规模部署 LLM。首先，我们将介绍 vLLM 如何使用连续批处理来更有效地为 LLM 服务。其次，我们将介绍 Dataflow 的模型管理器如何简化 vLLM 和其他大型模型框架的部署。

什么是 vLLM？

vLLM 是一个专为高吞吐量、低延迟 LLM 推理而设计的开源库。它通过采用几种专业技术（包括连续批处理）来优化 LLM 的服务。

要理解连续批处理的工作原理，让我们先看看模型传统上如何批处理输入。GPU 擅长并行处理，同时执行多个计算。GPU 可以通过批处理利用其所有可用内核一次性处理整批数据，而不是单独处理每个输入。这大大加快了推理过程；通常，同时对 8 个输入记录进行推理使用的资源与对单个记录进行推理差不多。

您可以将批处理想象成餐厅厨房：厨师可以将类似的订单分组并一起烹饪，而不是单独准备每道菜，从而节省时间和资源。如果您有 8 个炉灶，那么制作 1 个煎蛋卷或 8 个煎蛋卷所需的时间和精力都差不多。

然而，传统的批处理也有一些缺点。最重要的是在这种情况下，当执行推理需要不同的时间时，批处理效果不佳。大多数框架无法访问且不了解执行推理的底层机制，因此它们通常只是等待所有请求完成，然后再开始新的批处理。这意味着即使批次中的其他记录已经完成，单个慢速记录也可能消耗 GPU 的所有处理能力，从而导致作业更慢且成本更高。

batches of varying lengths lead to wasted compute

使用大型语言模型 (LLM) 运行推理时，等待整个批次完成可能会非常缓慢且昂贵。这是由于为一条记录生成推理所需的时间与记录长度的关系是 1:1。例如，假设我们将以下 2 个请求批处理到 LLM：

墨西哥的首都是哪里？

2. 墨西哥和美国之间有哪些文化差异和相似之处？

我们预期问题 (1) 的答案较短，而问题 (2) 的答案较长。但是，由于回答问题 (2) 需要更长的时间，因此我们必须等待该问题完成，然后才能返回批处理的任何结果，而它在这段时间里独占了我们的 GPU。

prompts of varying length lead to wasted compute

连续批处理允许 vLLM 在请求仍在运行时更新批次。它通过利用 LLM 执行推理的方式来做到这点：即通过重复生成响应中下一个令牌的过程来循环执行。因此，在生成句子“墨西哥的首都是墨西哥城”时，我们其实运行了 7 次推理（每个输出的字一次）。vLLM 的连续批处理技术让它可以在 LLM 每次为批次生成一组令牌时重新计算批次，而不是一次性批处理输入。这样，它就能在批次中即时添加请求，并在批次中的一个记录完全完成时返回早期结果。

continuous batching fully utilizes the GPU

在某些情况下，vLLM 的动态批处理和其他优化已证明可将主流 LLM 的推理吞吐量提高 2-4 倍，因此 vLLM 已成为模型服务的非常有用的工具。有关 vLLM 的更多信息，请查看此白皮书。

在 Dataflow 中使用 vLLM

在流式传输流水线中部署 vLLM 实例可能很复杂。如果使用传统方式，您需要：

在您的流水线中启动一个 vLLM 服务器。这通常比较麻烦，因为大多数流式传输系统都会启动多个工作进程，因此您需要选出一个主要进程来启动专用的 vLLM 流程。

确保所有工作进程都可以与该单例服务器通信。

监控并将服务投入生产，使其能够承受 vLLM 服务器故障。

这涉及大量的多处理，可能很耗时、容易出错，并且需要专业知识。此外，还需要对底层拓扑结构有深入的了解。如果您想尝试不同的机器配置（例如比较 8 核机器和 16 核机器的性能），拓扑结构往往会发生变化。

幸运的是，Dataflow 通过其模型管理器简化了这一过程。此功能抽象化了在流水线中管理和部署模型的复杂性。默认情况下，Dataflow 在其工作机器上为每个可用核心配置一个工作进程。这些进程负责处理进出工作器的 I/O 以及对数据执行的任何转换，而且它们完全独立运行。对于大多数流水线（包括用于 ML 用例的数据准备流水线），这是最佳拓扑结构。

Beam worker processes operate independently in parallel

然而，对于需要服务大型模型（例如 Gemma 模型）的流水线而言，这种拓扑结构就不适用了。将大型模型的副本加载到每个进程中既不经济，也得不到高性能，因为流水线很可能会遇到内存不足的问题。大多数此类流水线的理想拓扑结构是仅加载大型模型的单个副本。

Dataflow 的模型管理器旨在让用户可以控制在其流水线中部署的模型的具体副本数，无论网络拓扑如何。应用 RunInference 转换时，Dataflow 能够理解您的意图，以便为流水线创建理想的拓扑并部署最佳数量的模型；您只需提供一些配置参数即可。

the model manager allows Beam workers to share a single model

使用 vLLM 时，Dataflow 的模型管理器不会加载模型，而是在专用的推理进程中启动单个 vLLM 实例。然后，工作进程可以高效地将记录发送到此实例以进行推理。

the model manager allows Beam to spin up a dedicated vLLM process

Dataflow 可通过此模型管理器充分利用 vLLM 的连续批处理；当工作器接收到传入的请求时，它会异步将其附加到 vLLM 的请求队列并等待响应，从而允许 vLLM 动态批处理尽可能多的请求。

利用 Dataflow 的模型管理器和 RunInference 转换，可以轻松地将 vLLM 整合到流水线中。您只需要指定一些配置详细信息和几行代码。由于 Dataflow 能够理解您的底层意图，它可以为您配置流水线拓扑的其余部分。只需 5 行代码，您就可以编写一个完整的端到端流水线来读取数据，通过 vLLM 运行数据，然后将其输出到接收器。

model_handler = VLLMCompletionsModelHandler('google/gemma-2b')
with beam.Pipeline() as p:
  _ = (p | beam.ReadFromSource(<config>)
         | RunInference(model_handler) #将提示发送到 vLLM 并获取响应。
         | beam.WriteToSink(<config>))

您可以在此处找到完整的流水线并自己运行：https://cloud.google.com/dataflow/docs/notebooks/run_inference_vllm

性能

vLLM 显著提升了 Dataflow 流水线中 LLM 推理的性能。为了将 vLLM 的性能与使用固定大小批处理的原始流水线进行比较，我们使用具有 T4 GPU 的单个工作器运行了 2 个流水线。每个流水线都从 P3 数据集中读取提示，根据 google/gemma-2b 模型运行，并记录结果。

使用原始（默认）批处理策略时，处理 10,000 个提示需要 59.137 个 vCPU 小时。使用 vLLM 进行连续批处理时，仅需 2.481 个 vCPU 小时即可处理相同的 10,000 条提示。提高幅度超过 23 倍！

这里有一些注意事项：具体来说，两个流水线都没有进行调整，如果调整为使用更大或更均匀的批次，原始流水线可能会表现得更好。话虽如此，但这也是 vLLM 魔力的一部分；用不到 20 行代码且无需调整工作，我们就能得到高性能的 LLM 服务流水线！如果想比较另一个模型，只需更改模型处理程序中的一个字符串，就能以高性能的方式进行比较！

后续工作

通过将 vLLM 和 Dataflow 的强大功能相结合，您可以轻松为您的流式传输应用程序高效地部署和扩展 LLM。如需了解更多详细信息，请浏览此示例笔记本：https://cloud.google.com/dataflow/docs/notebooks/run_inference_vllm。

如需详细了解 Gemma 模型及其应用，请查看 Gemma 文档：https://ai.google.dev/gemma/docs

如需详细了解 vLLM 及其用于优化服务的其他机制，请访问 vLLM 文档：https://docs.vllm.ai/en/latest/