FlexLLMGen：单 GPU 实现大语言模型的高吞吐量生成推理 [论文]

FlexLLMGen 是一个高吞吐量生成引擎，用于在有限的 GPU 内存下运行大语言模型。FlexLLMGen 通过高效的 IO 卸载、压缩和大有效批处理大小来实现高吞吐量生成。

动机

近年来，大语言模型在各种任务中展现出卓越的性能。LLM 不仅越来越多地应用于交互式应用（如聊天），也应用于许多“后台”任务。这些任务包括基准测试、信息提取、数据整理和表单处理。

这些应用的一个关键特征是面向吞吐量：它们需要以批处理方式对数百万个 token 运行 LLM 推理，例如，公司语料库中的所有私有文档，或 HELM 基准测试中的所有任务。这些工作负载对延迟不敏感——用户启动一个作业并让其运行一整夜——但提高吞吐量对于降低成本至关重要。吞吐量是衡量整个作业运行期间（可能长达数小时）每秒处理的 token 数量的指标。面向吞吐量的工作负载提供了用延迟换取更高吞吐量的机会，这使得利用低成本商用 GPU 变得更加容易。

FlexLLMGen 的目标是创建一个高吞吐量系统，使基础模型能够在低成本硬件（例如单个商用 GPU，而非昂贵的系统）上应用于面向吞吐量的任务，从而激发新的、令人兴奋的应用。

查看 示例，了解您可以在单个商用 GPU 上使用 FlexLLMGen 运行的任务，包括基准测试和数据整理。

❌ 局限性。作为一个在弱 GPU 上运行的基于卸载的系统，FlexLLMGen 也有其局限性。当您拥有足够强大的 GPU 来容纳整个模型时，FlexLLMGen 可能会明显更慢，尤其是在小批量情况下。FlexLLMGen 主要针对单 GPU 上面向吞吐量的批处理场景（例如，批量分类或从许多文档中提取信息）进行了优化。

本项目得益于与以下机构的合作：

目录

• 安装
• 使用与示例
• 单 GPU 入门
• 使用 FlexLLMGen 运行 HELM 基准测试
• 使用 FlexLLMGen 运行数据整理任务
• 扩展到分布式 GPU
• API 示例
• 常见问题
• 性能结果
• 工作原理
• 路线图

安装

要求：
- PyTorch >= 1.12 (帮助)

方法 1：使用 pip

pip install flexllmgen

方法 2：从源码安装

git clone https://github.com/FMInference/FlexLLMGen.git
cd FlexLLMGen
pip install -e .

使用与示例

单 GPU 入门

OPT-1.3B

要开始使用，您可以先尝试像 OPT-1.3B 这样的小模型。它适合单个 GPU，因此不需要卸载。FlexLLMGen 会自动从 Hugging Face 下载权重。

python3 -m flexllmgen.flex_opt --model facebook/opt-1.3b

您应该会看到一些由 OPT-1.3B 生成的文本以及基准测试结果。

OPT-30B

要运行像 OPT-30B 这样的大模型，您需要使用 CPU 卸载。您可以尝试下面的命令。
--percent 参数分别指定了参数、注意力缓存和隐藏状态的卸载策略。该参数的确切含义可以在这里找到。

python3 -m flexllmgen.flex_opt --model facebook/opt-30b --percent 0 100 100 0 100 0

OPT-175B

要运行 OPT-175B，您需要从 metaseq 下载权重，并将权重转换为 Alpa 格式。
然后，您可以尝试将所有权重卸载到磁盘：

python3 -m flexllmgen.flex_opt --model facebook/opt-175b --percent 0 0 100 0 100 0 --offload-dir YOUR_SSD_FOLDER

使用 FlexLLMGen 运行 HELM 基准测试

FlexLLMGen 可以集成到语言模型基准测试框架 HELM 中，作为其执行后端。
您可以使用以下命令，在单个 T4 (16GB) GPU 和 200GB DRAM 上运行大规模多任务语言理解 (MMLU) 场景。

pip install crfm-helm
python3 -m flexllmgen.apps.helm_run --description mmlu:model=text,subject=abstract_algebra,data_augmentation=canonical --pad-to-seq-len 512 --model facebook/opt-30b --percent 20 80 0 100 0 100 --gpu-batch-size 48 --num-gpu-batches 3 --max-eval-instance 100

请注意，仅测试了 HELM 场景的一个子集。查看更多已测试的场景这里。

使用 FlexLLMGen 运行数据整理任务

您可以按照此处的说明，运行论文《Can Foundation Models Wrangle Your Data?》中的示例。

扩展到分布式 GPU

如果您拥有多台带 GPU 的机器，FlexLLMGen 可以将卸载与流水线并行相结合以实现扩展。
例如，如果您有 2 个 GPU，但聚合的 GPU 内存小于模型大小，您仍然需要卸载。FlexLLMGen 允许您在这 2 个 GPU 上进行流水线并行以加速生成。
但要获得扩展性能，您应该在分布式机器上拥有 GPU。
查看示例这里。

API 示例

我们在 completion.py 中演示了 FlexLLMGen API 的用法。
此示例展示了如何为两个句子运行生成。要获得 FlexLLMGen 的最佳吞吐量，通常需要批处理更多句子。

生成 API

FlexLLMGen 的生成 API 遵循 Hugging Face transformers 的风格。

output_ids = model.generate(
    input_ids,
    do_sample=True,
    temperature=0.7,
    max_new_tokens=32,
    stop=stop)

示例命令

您可以使用下面的示例命令。
如果您没有足够的 GPU/CPU 内存，请参阅处理内存不足部分。

# 使用 OPT-6.7B 完成。您至少需要 15GB GPU 内存。
python3 -m flexllmgen.apps.completion --model facebook/opt-6.7b

# 使用 OPT-30B 完成。您需要大约 90GB CPU 内存。
python3 -m flexllmgen.apps.completion --model facebook/opt-30b --percent 0 100 100 0 100 0

# 使用指令调优的 OPT-IML-MAX-30B 完成。您需要大约 90GB CPU 内存。
python3 -m flexllmgen.apps.completion --model facebook/opt-iml-max-30b --percent 0 100 100 0 100 0

常见问题

如何设置卸载策略和 --percent？

我们稍后将发布一个自动策略优化器，但目前您必须手动尝试几种策略。
高吞吐量生成的思路是尽可能多地将参数和注意力缓存卸载到 CPU，必要时卸载到磁盘。
您可以在我们的基准测试这里查看参考策略。
为避免内存不足，您可以调整 --percent 以将更多张量卸载到 CPU 和磁盘。

如何处理内存不足？

如果您没有足够的 GPU/CPU 内存，可以尝试以下几种方法。它们会节省更多内存，但运行速度较慢。

• 不要固定权重，添加 --pin-weight 0。这可以将 CPU 上的权重内存使用量减少约 20% 或更多。
• 启用权重压缩，添加 --compress-weight。这可以将权重内存使用量减少约 70%。
• 将所有权重卸载到磁盘，使用 --percent 0 0 100 0 100 0。这需要非常少的 CPU 和 GPU 内存。

性能结果

生成吞吐量 (token/s)

括号内为对应的有效批处理大小和最低卸载设备。更多详情请见此处。
| 系统 | OPT-6.7B | OPT-30B | OPT-175B |
| ------ | -------- | ------- | -------- |
| Hugging Face Accelerate | 25.12 (2 在 GPU 上) | 0.62 (8 在 CPU 上) | 0.01 (2 在磁盘上) |
| DeepSpeed ZeRO-Inference | 9.28 (16 在 CPU 上) | 0.60 (4 在 CPU 上) | 0.01 (1 在磁盘上) |
| Petals | 8.25 (2 在 GPU 上) | 2.84 (2 在 GPU 上) | 0.08 (2 在 GPU 上) |
| FlexLLMGen | 25.26 (2 在 GPU 上) | 7.32 (144 在 CPU 上) | 0.69 (256 在磁盘上) |
| FlexLLMGen 带压缩 | 29.12 (72 在 GPU 上) | 8.38 (512 在 CPU 上) | 1.12 (144 在 CPU 上) |

• 硬件：GCP 上的一个 NVIDIA T4 (16GB) 实例，配备 208GB DRAM 和 1.5TB SSD。
• 工作负载：输入序列长度 = 512，输出序列长度 = 32。批处理大小调整为一个较大值，以最大化每个系统的生成吞吐量。
• 指标：生成吞吐量 (token/s) = 生成的 token 数量 / (处理提示的时间 + 生成的时间)。

如何复现。

延迟-吞吐量权衡

下图显示了三个基于卸载的系统在 OPT-175B（左）和 OPT-30B（右）上的延迟和吞吐量权衡。
FlexLLMGen 实现了一个新的帕累托最优前沿，两个模型的最大吞吐量都显著提高。其他系统由于内存不足无法进一步提高吞吐量。
"FlexLLMGen(c)" 表示带压缩的 FlexLLMGen。

工作原理

FlexLLMGen 可以通过聚合 GPU、CPU 和磁盘的内存和计算资源，在各种硬件资源限制下灵活配置。通过线性规划优化器，它搜索存储和访问张量（包括权重、激活和注意力键/值 (KV) 缓存）的最佳模式。FlexLLMGen 进一步将权重和 KV 缓存压缩到 4 位，且精度损失可忽略不计。

FlexLLMGen 的一个关键思想是利用延迟-吞吐量的权衡。对于卸载方法来说，实现低延迟本身具有挑战性，但对于面向吞吐量的场景，卸载的 I/O 效率可以大大提高（见上图）。
FlexLLMGen 利用块调度来重用权重并重叠 I/O 与计算，如下面图 (b) 所示，而其他基线系统使用低效的逐行调度，如下面图 (a) 所示。

更多技术细节请参阅我们的论文。

路线图

我们计划开发以下功能。

• [ ] 优化同一台机器上多个 GPU 的性能
• [ ] 支持更多模型 (BLOOM, CodeGen, GLM)
• [X] 发布成本模型和策略优化器
• [ ] Macbook 支持 (M1 和 M2)
• [ ] AMD 支持