FlexGen: 使用单个GPU对大型语言模型进行高吞吐量生成推理 [论文]
FlexGen是一个用于在有限的GPU内存下运行大型语言模型的高吞吐量生成引擎。通过IO高效卸载、压缩和大有效批量,FlexGen实现了高吞吐量的生成。
近年来,大型语言模型(LLM)在广泛的任务中表现出色。LLM不仅被应用于交互式应用(如聊天),还被广泛应用于许多“后台”任务。这些任务包括基准测试、信息提取、数据整理和表单处理。
这些应用的一个关键特征是它们是面向吞吐量的:它们需要对数百万个标记进行批量的LLM推理,例如公司语料库中的所有私密文档,或HELM基准中的所有任务。这些工作负载对延迟不太敏感——用户启动一个作业并让它运行一整夜——但增加吞吐量对于降低成本至关重要。吞吐量是整个作业运行时(可能是数小时)内每秒处理的标记数的度量。面向吞吐量的工作负载提供了用更高吞吐量来权衡延迟的机会,这使得利用低成本的普通GPU变得更加容易。
FlexGen的目标是创建一个高吞吐量系统,以便在低成本硬件(如单个普通GPU而非昂贵的系统)上为面向吞吐量的任务启用新的、令人兴奋的基础模型应用。
查看FlexGen在单个普通GPU上运行的示例,包括基准测试和数据整理。
❌ 局限性。作为一个基于卸载、运行于较弱GPU上的系统,FlexGen也有其局限性。与拥有足够强大的GPU以容纳整个模型的情况相比,FlexGen可能显著较慢,尤其是在小批量的情况下。FlexGen主要针对面向吞吐量的批处理设置(例如批量分类或从大量文档中提取信息)进行了优化,适用于单个GPU。
这个项目得益于以下合作的支持:
<a href="https://cs.stanford.edu/"><img src="https://yellow-cdn.veclightyear.com/2b54e442/7bd48afe-bab9-4f04-beba-a4a39126ba95.png" height="20"></a> <a href="https://sky.cs.berkeley.edu/"><img src="https://yellow-cdn.veclightyear.com/2b54e442/00aaad63-de12-46fa-8bcf-46c14b613695.png" height="22"></a> <a href="https://www.andrew.cmu.edu/user/beidic/"><img src="https://upload.wikimedia.org/wikipedia/commons/9/9b/Carnegie_Mellon_wordmark.svg" height="20"></a> <a href="https://www.together.xyz/"><img src="https://yellow-cdn.veclightyear.com/2b54e442/2df596ca-ba67-4061-9e27-2a21f6948f35.png" height="20"></a> <a href="https://research.yandex.com/"><img src="https://yellow-cdn.veclightyear.com/2b54e442/9b8a84f0-a24c-4797-8d74-32b9e6d4f555.png" height="20"></a> <a href="https://ds3lab.inf.ethz.ch/"><img src="https://yellow-cdn.veclightyear.com/2b54e442/e62c2abd-2c47-4e4f-88cc-a7e226f00f06.png" height="20"></a>
要求:
pip install flexgen
git clone https://github.com/FMInference/FlexGen.git
cd FlexGen
pip install -e .
要开始使用,可以先尝试一个小模型如OPT-1.3B。它适合在单个GPU上运行,因此无需卸载。FlexGen会自动从Hugging Face下载权重。
python3 -m flexgen.flex_opt --model facebook/opt-1.3b
您应该会看到由OPT-1.3B生成的文本以及基准测试结果。
要运行像OPT-30B这样的大型模型,您需要使用CPU卸载。您可以尝试以下命令。--percent
参数分别指定参数、注意力缓存和隐藏状态的卸载策略。该参数的确切含义可以在这里找到。
python3 -m flexgen.flex_opt --model facebook/opt-30b --percent 0 100 100 0 100 0
要运行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中,作为其执行后台。您可以使用以下命令在单个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场景经过测试。查看更多测试场景这里。
您可以按照这里的说明,运行本文中的示例,'Can Foundation Models Wrangle Your Data?'。
如果你有多台带有GPU的机器,FlexGen可以结合卸载与流水线并行来实现扩展。例如,如果你有2个GPU,但总的GPU内存小于模型大小,你仍然需要卸载。FlexGen允许你使用这2个GPU进行流水线并行,以加速生成。但要实现扩展的性能,你需要在分布式机器上配置GPU。参见此处的示例。
我们在completion.py中展示了FlexGen API的使用示例。此示例展示了如何对两个句子进行生成。为了从FlexGen中获得最佳吞吐量,你通常需要批处理更多的句子。
FlexGen有一个遵循Hugging Face的transformers风格的生成API。
output_ids = model.generate( input_ids, do_sample=True, temperature=0.7, max_new_tokens=32, stop=stop)