认识FlexGen：一种用于在有限的GPU内存上运行大型语言模型（LLMs）的高吞吐量生成引擎
FlexGen A high-throughput generation engine for running large language models (LLMs) on limited GPU memory.
大型语言模型（LLMs）最近在各种任务上展现出令人印象深刻的性能。生成性LLM推理具有前所未有的能力，但也面临特定的困难。这些模型可以包含数十亿或数万亿个参数，这意味着运行它们需要巨大的内存和计算能力。例如，GPT-175B只需要325GB的GPU RAM来加载其模型权重。至少需要五个A100（80GB）GPU和复杂的并行技术才能将此模型适配到GPU上。因此，最近对减少LLM推理所需的资源产生了很大的兴趣。
LLMs用于各种“后台”操作，包括基准测试、信息提取、数据整理、表单处理以及聊天机器人等交互式用例。在这项研究中，他们专注于所谓的以吞吐量为导向的生成性推理情况。这些任务经常需要在大量标记上批量进行LLM推理，例如公司语料库中的所有论文，对标记生成的延迟不太敏感，这是这些工作的一个重要特征。因此，在某些工作负载中，通过以延迟为代价获得更好的吞吐量，可以降低资源需求。
已经使用了三种方法来减少LLM推理所需的资源：模型压缩以减少总体内存占用，协作推理通过分散推理成本进行分散，以及卸载以更好地利用CPU和磁盘上的内存。尽管存在明确的限制，但这些策略已经大大减少了雇佣LLMs所需的资源。在前两种方法的研究中，通常需要帮助在单个普通GPU上运行175B级模型，因为它假设模型适配GPU内存。另一方面，由于低效的I/O调度和张量放置，第三类基于卸载的尖端系统无法在单个GPU上达到可接受的吞吐量。
在单个普通GPU上，他们的主要目标是构建用于高吞吐量生成性推理的有效卸载机制。他们可以部分加载LLM，并将计算分批次地卸载到二级存储器中，以在受限的GPU内存中操作LLM。在典型系统中，内存层次结构分为三个层次。较低的级别速度较慢，但更充足，而较高的级别速度更快，但更稀缺。小批量大小可能导致这些系统出现瓶颈。通过使用较高的批量大小并将昂贵的I/O操作分布在与处理重叠的大批量输入的多个内存层次中，他们可以在以吞吐量为导向的场景中牺牲延迟。
即使他们可以牺牲延迟，实现具有受限GPU内存的高吞吐量生成性推理仍然很困难。第一个困难是制定一个成功的卸载计划。该计划应该明确说明哪些张量应该被卸载，它们应该在三级内存结构中的哪个位置卸载，并在推理过程中的何时卸载。生成性推理中使用了三种类型的张量：权重、激活和键值（KV）缓存。
由于算法的逐批次、逐标记和逐层结构，有几种计算方式。这些选项共同构成了一个复杂的设计空间。目前使用的基于卸载的推理系统继承了基于训练的方法，这些方法进行了过多的I/O操作，并且实现的吞吐量远低于理论硬件限制，使它们成为一些推理的不良领域。创建高效的压缩算法是第二个问题。早期的出版物已经展示了LLMs的权重和激活的有希望的压缩结果。然而，在高吞吐量生成性推理中，当压缩和卸载结合在一起时，额外的压缩策略受到I/O成本和权重和KV缓存的内存减少的驱动。
来自UCB、斯坦福大学、CMU、Meta、Yandex、ETH和HSE的研究人员共同介绍了FlexGen，这是一个用于高吞吐量LLM推理的卸载框架，用于克服这些问题。FlexGen通过将来自GPU、CPU和磁盘的内存结合起来，有效地调度I/O活动、潜在的压缩技术和分布式管道并行性。以下是他们的贡献：
- 他们明确描述了潜在卸载选项的搜索空间，考虑了计算调度、张量放置和计算委派。他们证明了他们的搜索空间捕捉到了一个I/O复杂度在最优性的2个范围内的计算顺序。接下来，他们创建了一个基于线性规划的搜索算法，在搜索空间内最大化吞吐量。
- 他们展示了在不重新训练或校准的情况下，可以将类似OPT-175B的LLMs的权重和KV缓存减少到4位，几乎没有精度损失。适用于在卸载期间降低I/O成本和内存使用的细粒度组内量化实现了这一点。
- 他们通过在NVIDIA T4（16GB）GPU上运行OPT-175B来展示了FlexGen的效率。与两个尖端的基于卸载的推理算法DeepSpeed Zero-Inference和Hugging Face Accelerate相比，FlexGen通常允许更大的批量大小。因此，FlexGen可以实现更大的吞吐量。
请查看Paper和Github。本研究所有功劳归于该项目的研究人员。此外，请不要忘记加入我们的16k+ ML SubReddit、Discord Channel和Email Newsletter，在那里我们分享最新的AI研究新闻、酷炫的AI项目等等。
使用Tensorleap的可解释性平台揭示深度学习的秘密
本文介绍了FlexGen: 一种用于在有限GPU内存下运行大型语言模型（LLMs）的高吞吐量生成引擎。
本文最初发布于MarkTechPost。
