FlexGen

FlexGen: 使用单个GPU对大型语言模型进行高吞吐量生成推理 [论文]

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

动机

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

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

FlexGen的目标是创建一个高吞吐量系统，以便在低成本硬件（如单个普通GPU而非昂贵的系统）上为面向吞吐量的任务启用新的、令人兴奋的基础模型应用。

查看FlexGen在单个普通GPU上运行的示例，包括基准测试和数据整理。

❌ 局限性。作为一个基于卸载、运行于较弱GPU上的系统，FlexGen也有其局限性。与拥有足够强大的GPU以容纳整个模型的情况相比，FlexGen可能显著较慢，尤其是在小批量的情况下。FlexGen主要针对面向吞吐量的批处理设置（例如批量分类或从大量文档中提取信息）进行了优化，适用于单个GPU。

这个项目得益于以下合作的支持：

安装

要求:

PyTorch >= 1.12 (帮助)

方法1: 使用pip

pip install flexgen

方法2: 从源代码安装

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

使用与示例

使用单个GPU入门

OPT-1.3B

要开始使用，可以先尝试一个小模型如OPT-1.3B。它适合在单个GPU上运行，因此无需卸载。FlexGen会自动从Hugging Face下载权重。

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

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

OPT-30B

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

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

OPT-175B

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

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

使用FlexGen运行HELM基准

FlexGen可以集成到HELM中，作为其执行后台。您可以使用以下命令在单个T4（16GB）GPU和200GB的DRAM上运行一个大规模多任务语言理解（MMLU）场景。

pip install crfm-helm
python3 -m flexgen.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场景经过测试。查看更多测试场景这里。

使用FlexGen运行数据整理任务

您可以按照这里的说明，运行本文中的示例，'Can Foundation Models Wrangle Your Data?'。

扩展至分布式GPU

如果你有多台带有GPU的机器，FlexGen可以结合卸载与流水线并行来实现扩展。例如，如果你有2个GPU，但总的GPU内存小于模型大小，你仍然需要卸载。FlexGen允许你使用这2个GPU进行流水线并行，以加速生成。但要实现扩展的性能，你需要在分布式机器上配置GPU。参见此处的示例。

API示例

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

生成API

FlexGen有一个遵循Hugging Face的transformers风格的生成API。

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 flexgen.apps.completion --model facebook/opt-6.7b

# 使用OPT-30B进行补全。你需要大约90GB的CPU内存。
python3 -m flexgen.apps.completion --model facebook/opt-30b --percent 0 100 100 0 100 0

# 使用经过指令微调的OPT-IML-MAX-30B进行补全。你需要大约90GB的CPU内存。
python3 -m flexgen.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上)
FlexGen	25.26 (2个在GPU上)	7.32 (144个在CPU上)	0.69 (256个在磁盘上)
FlexGen with Compression	29.12 (72个在GPU上)	8.38 (512个在CPU上)	1.12 (144个在CPU上)