FMInference/FlexLLMGen

# FlexLLMGen：使用单 GPU 实现大型语言模型的高吞吐量生成推理 [[论文](https://arxiv.org/abs/2303.06865)] FlexLLMGen 是一个高吞吐量生成引擎，旨在显存受限的情况下运行大型语言模型。FlexLLMGen 通过高 IO 效率的卸载、压缩以及**更大的有效 batch size**，实现了**高吞吐量**的生成。 ## 动机 近年来，大型语言模型（LLM）在各种任务中展现出了出色的表现。LLM 不仅越来越多地被应用于交互式应用（如聊天），还应用于许多“后台”任务中。 这些任务包括基准测试、信息提取、数据整理和表单处理。 这些应用的一个关键特征是它们是**吞吐量导向的**：它们需要批量处理数百万个 token 的 LLM 推理，例如，公司语料库中的所有私有文档，或者 [HELM](https://crfm.stanford.edu/helm/latest/) 基准测试中的所有任务。 这些工作负载对延迟不太敏感——用户启动一个任务并让它通宵运行——但提高吞吐量对于降低成本至关重要。 吞吐量是指在任务整个运行期间（可能长达数小时）每秒处理的 token 数量的衡量指标。 吞吐量导向的工作负载提供了以延迟换取更高吞吐量的机会，这使得更容易利用低成本的商品级 GPU。 FlexLLMGen 的目标是创建一个高吞吐量系统，在低成本硬件（例如单个商品级 GPU，而不是昂贵的系统）上，为基础模型应用于吞吐量导向的任务开启全新且令人兴奋的应用。 查看可以使用 FlexLLMGen 在单个商品级 GPU 上运行的[示例](#examples)，包括基准测试和数据整理。 ❌ **局限性**。作为一个在性能较弱的 GPU 上运行的基于卸载的系统，FlexLLMGen 也有其自身的局限性。 与你拥有足够强大的 GPU 来容纳整个模型的情况相比，FlexLLMGen 的速度可能会明显慢一些，尤其是在小 batch 的情况下。 FlexLLMGen 主要针对单 GPU 上吞吐量导向的批处理设置（例如，批量对许多文档进行分类或提取信息）进行了优化。 ## 目录 - [安装](#installation) - [用法与示例](#usage-and-examples) - [单 GPU 快速入门](#get-started-with-a-single-gpu) - [使用 FlexLLMGen 运行 HELM 基准测试](#run-helm-benchmark-with-flexllmgen) - [使用 FlexLLMGen 运行数据整理任务](#run-data-wrangling-tasks-with-flexllmgen) - [扩展至分布式 GPU](#scaling-to-distributed-gpus) - [API 示例](#api-example) - [常见问题解答](#frequently-asked-questions) - [性能结果](#performance-results) - [工作原理](#how-it-works) - [路线图](#roadmap) ## 安装说明 环境要求： - PyTorch >= 1.12 [(帮助)](https://pytorch.org/get-started/locally/) ### 方法 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` 参数分别指定了参数、注意力缓存和隐藏状态的卸载策略。 该参数的确切含义可以在[这里](https://github.com/FMInference/FlexLLMGen/blob/9d092d848f106cd9eaf305c12ef3590f7bcb0277/flexllmgen/flex_opt.py#L1271-L1279)找到。 ``` python3 -m flexllmgen.flex_opt --model facebook/opt-30b --percent 0 100 100 0 100 0 ``` #### OPT-175B 要运行 OPT-175B，你需要从 [metaseq](https://github.com/facebookresearch/metaseq/tree/main/projects/OPT) 下载权重，并将权重转换为 Alpa [格式](https://alpa.ai/tutorials/opt_serving.html#convert-opt-175b-weights-into-alpa-formats)。 然后你可以尝试通过以下命令将所有权重卸载到磁盘上 ``` 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](https://crfm.stanford.edu/helm)（一个语言模型基准测试框架）中。 你可以使用以下命令在单个 T4 (16GB) GPU 和 200GB DRAM 上运行大规模多任务语言理解（MMLU）[场景](https://crfm.stanford.edu/helm/latest/?group=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/apps/helm_passed_30b.sh)查看更多已测试的场景。 ### 使用 FlexLLMGen 运行数据整理任务 你可以按照[此处](flexllmgen/apps/data_wrangle)的说明，运行论文[《基础模型能整理你的数据吗？》](https://arxiv.org/abs/2205.09911)中的示例。 ### 扩展至分布式 GPU 如果你拥有多台带 GPU 的机器，FlexLLMGen 可以将卸载与流水线并行结合起来，从而实现扩展。 例如，如果你有 2 个 GPU，但 GPU 总内存小于模型大小，你仍然需要卸载。FlexLLMGen 允许你在这 2 个 GPU 上进行流水线并行，以加速生成过程。 但是要获得扩展性能，你应该使用分布式机器上的 GPU。 在[此处](https://github.com/FMInference/FlexLLMGen/tree/main/benchmark/flexllmgen#distributed-gpus)查看示例。 ### API 示例 我们在 [completion.py](flexllmgen/apps/completion.py) 中演示了 FlexLLMGen API 的用法。 此示例展示了如何对两个句子进行生成。 为了从 FlexLLMGen 中获得最佳吞吐量，你通常需要批量处理更多的句子。 #### 生成 API FlexLLMGen 提供了一个遵循 Hugging Face transformers 风格的生成 API。 ``` output_ids = model.generate( input_ids, do_sample=True, temperature=0.7, max_new_tokens=32, stop=stop) ``` #### 示例命令 你可以使用以下示例命令。 如果你没有足够的 GPU/CPU 内存，请参阅[处理内存溢出](#handle-out-of-memory)部分。 ``` # 使用 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 ``` ``` # 使用经过 instruction-tuned 的 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，必要时也可以卸载到磁盘。 你可以在我们的基准测试中查看参考策略，请点击[这里](https://github.com/FMInference/FlexLLMGen/blob/9d092d848f106cd9eaf305c12ef3590f7bcb0277/benchmark/flexllmgen/bench_suite.py#L39-L79)。 为了避免内存溢出，你可以调整 `--percent`，将更多的 tensor 卸载到 CPU 和磁盘上。 #### 如何处理内存溢出？ 如果你没有足够的 GPU/CPU 内存，可以尝试以下几种方法。 它们能节省更多内存，但运行速度会变慢。 - 添加 `--pin-weight 0` 以不固定权重。这可以减少 CPU 上约 20% 或更多的权重内存使用。 - 添加 `--compress-weight` 启用权重压缩。这可以减少约 70% 的权重内存使用。 - 使用 `--percent 0 0 100 0 100 0` 将所有权重卸载到磁盘。这只需要很少的 CPU 和 GPU 内存。 ## 性能结果 ### 生成吞吐量 (token/s) 括号内为相应的有效 batch size 和最低卸载设备。更多详情请参见[此处](benchmark/batch_size_table.md)。 | 系统 | 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。针对每个系统，将 batch size 调整为能使生成吞吐量最大化的**一个大值**。 - 指标：生成吞吐量 (token/s) = 生成的 token 数量 /（处理 prompt 的时间 + 生成时间）。 如何[复现](benchmark/flexllmgen)。 ### 延迟-吞吐量权衡 下图展示了三种基于卸载的系统在 OPT-175B（左）和 OPT-30B（右）上的延迟和吞吐量权衡。 FlexLLMGen 实现了新的帕累托最优边界，并且在两个模型上都获得了显著更高的最大吞吐量。 由于内存溢出，其他系统无法进一步提高吞吐量。 "FlexLLMGen(c)" 是开启了压缩的 FlexLLMGen。 ## 工作原理 FlexLLMGen 可以通过聚合来自 GPU、CPU 和磁盘的内存与计算资源，在各种硬件资源限制下进行灵活配置。通过线性规划优化器，它会搜索最佳的存储和访问 tensor 的模式，包括权重、激活值和注意力键值（KV）缓存。FlexLLMGen 进一步将权重和 KV 缓存压缩至 4 bit，且准确率损失可以忽略不计。 FlexLLMGen 的一个关键理念是利用延迟-吞吐量的权衡。对于卸载方法来说，实现低延迟本身就具有挑战性，但在吞吐量导向的场景中，卸载的 I/O 效率可以得到极大的提升（见上图）。 FlexLLMGen 利用块调度来重用权重并将 I/O 与计算重叠，如下面的图所示，而其他基线系统则使用低效的逐行调度，如下面的图所示。 更多技术细节请参阅我们的[论文](https://arxiv.org/abs/2303.06865)。 ## 路线图 我们计划开发以下功能。 - [ ] 优化同一台机器上多个 GPU 的性能 - [ ] 支持更多模型 (BLOOM, CodeGen, GLM) - [X] 发布成本模型和策略优化器 - [ ] Macbook 支持 (M1 和 M2) - [ ] AMD 支持

标签：DLL 劫持, GPU卸载, 凭据扫描, 单GPU推理, 大语言模型, 推理引擎, 逆向工具, 高吞吐量