宣布 vLLM 支持 Llama 3.1

今天，vLLM 团队很高兴能与 Meta 合作，宣布支持 Llama 3.1 模型系列。Llama 3.1 带来了令人振奋的新特性，包括更长的上下文长度（最高支持 128K token）、更大的模型规模（最高 405B 参数）以及更先进的模型能力。vLLM 社区增加了许多增强功能，以确保更长、更大的 Llama 模型能在 vLLM 上顺畅运行，这些功能包括分块预填充（chunked prefill）、FP8 量化和流水线并行。我们将在本文中介绍这些新的增强功能。

简介

vLLM 是一个快速、易于使用且开源的大语言模型服务引擎。vLLM 目前支持超过 40 种开源大语言模型，兼容多种硬件平台（Nvidia GPU、AMD GPU、AWS Inferentia、Google TPU、Intel CPU/GPU、Gaudi 等），并提供全方位的推理优化。点击这里了解更多关于 vLLM 的信息。

对于新的 Llama 3.1 系列，vLLM 能够以完整的 128K 上下文窗口运行这些模型。为了支持如此大的上下文窗口，vLLM 会自动启用分块预填充。分块预填充不仅能有效控制内存使用，还能减少因长提示词处理而导致对现有请求的打断。您可以通过运行以下命令或使用我们的官方 Docker 镜像（vllm/vllm-openai）来安装 vLLM。

pip install -U vllm

针对大型 Llama 405B 模型，vLLM 提供了多种支持方案：

• FP8： vLLM 可以在 8xA100 或 8xH100 上原生运行官方的 FP8 量化模型。
• 流水线并行： vLLM 通过将模型的不同层放置在不同节点上，从而在多节点上运行官方的 BF16 版本。
• 张量并行： vLLM 也可以通过将模型分片到多个节点以及节点内的多个 GPU 上来运行。
• AMD MI300x 或 NVIDIA H200： vLLM 可在单个 8xMI300x 或 8xH200 机器上运行该模型，这些 GPU 分别拥有 192GB 和 141GB 显存。
• CPU 卸载（Offloading）： 作为最后的手段，vLLM 可以在执行前向传播时将部分权重卸载到 CPU，从而允许您在有限的 GPU 内存上运行大型模型，并保持完全精度。

请注意，虽然 vLLM 支持所有这些方法，但性能表现仍处于初步阶段。vLLM 社区正在积极进行优化工作，我们欢迎大家的贡献。例如，我们正在积极探索更多的模型量化方法，以及提高流水线并行的吞吐量。本文后面列出的性能数据仅供早期参考；我们预计在接下来的几周内性能将会有显著提升。

在所有方法中，我们建议单节点使用 FP8，多节点使用流水线并行。让我们更详细地探讨一下这些方案。

FP8

FP8 使用 8 位来表示浮点数。当前这一代的 GPU（H100, MI300x）通过专用张量核心对 FP8 提供原生支持。目前，vLLM 可以在 KV 缓存、注意力机制和 MLP 层运行 FP8 量化模型。这减少了内存占用、提高了吞吐量、降低了延迟，且精度损失极小。

目前，vLLM 通过利用 MLP 层中的通道量化（per-channel quantization），支持经 FBGEMM 量化的官方 Meta Llama 3.1 405B FP8 模型。特别是，上投影（up）、门控（gate）和下投影（down）的每个通道都被量化，并乘以一个静态缩放因子。结合跳过第一层和最后一层的量化以及静态上限，这种方法对模型精度的影响微乎其微。您可以使用最新的 vLLM，通过以下命令在单台 8xH100 或 8xA100 上运行该模型：

$ vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct-FP8 --tensor-parallel-size 8

在使用 FP8 量化模型处理平均输入长度为 1024 个 token、平均输出长度为 128 个 token 的请求时，服务器每秒可维持 2.82 个请求。相应的服务吞吐量分别为每秒 2884.86 个输入 token 和每秒 291.53 个输出 token。

我们还独立证实了 FP8 检查点的精度损失极小。例如，使用 lm-eval-harness 对 GSM8K 基准进行测试（8-shot 和思维链），我们观察到准确匹配分数为 95.38%（标准差 ±0.56%），与官方 BF16 分数 96.8% 相比，下降幅度极小。

流水线并行

如果您想在不量化的情况下运行 Llama 3.1 405B 模型该怎么办？您可以使用 vLLM 的流水线并行，在 16xH100 或 16xA100 GPU 上实现！

流水线并行将模型拆分为较小的层组，以流水线方式在两个或多个节点上并行执行。与需要昂贵的 all-reduce 操作的张量并行不同，流水线并行是在层边界处对模型进行分区，仅需要低成本的点对点通信。当您拥有多个不一定通过 InfiniBand 等高速互连设备连接的节点时，这一点尤为有用。

vLLM 支持将流水线并行与张量并行结合使用。例如，通过 2 节点的 16 个 GPU，您可以使用 2 路流水线并行和 8 路张量并行来优化硬件使用。该配置将模型的一半映射到每个节点，利用 NVLink 进行 all-reduce 操作，将每层划分到 8 个 GPU 上。您可以使用以下命令运行 Llama 3.1 405B 模型：

$ vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --tensor-parallel-size 8 --pipeline-parallel-size 2

如果您拥有像 InfiniBand 这样的高速互连设备，则可以使用 16 路张量并行。

$ vllm serve meta-llama/Meta-Llama-3.1-405B-Instruct --tensor-parallel-size 16

在 16xH100 GPU 上，使用合成数据集（平均输入长度 1024，平均输出长度 128）的服务吞吐量。

我们观察到，当节点之间没有连接 InfiniBand 时，流水线并行是必不可少的。与 16 路张量并行相比，将 2 路流水线并行与 8 路张量并行相结合，性能提升了 6.6 倍。另一方面，在具备 InfiniBand 的情况下，两种配置的性能表现相似。

要了解有关使用 vLLM 进行分布式推理的更多信息，请参考此文档。关于 CPU 卸载，请参考此示例。

致谢

我们要感谢 Meta 提供的预发布合作，让我们得以测试该模型。在发布之外，我们还要感谢以下为本博客文章中提到的功能做出贡献的 vLLM 贡献者：Neural Magic（FP8 量化）；CentML 和 Snowflake AI Research（流水线并行）；Anyscale（分块预填充功能）。评估是在具备 InfiniBand 的 Lambda 一键集群上运行的，我们感谢 Lambda 提供的资源以及流畅的集群设置体验。