跑ChatGPT体量模型，从此只需一块GPU：加速百倍的方法来了

计算成本是人们打造 ChatGPT 等大模型面临的重大挑战之一。

据统计，从 GPT 进化到 GPT-3 的过程也是模型体量增长的过程 ——

参数

量从 1.17 亿增加到了 1750 亿，预训练数据量从

5G

B 增加到 45TB，其中 GPT-3 训练一次的费用是 460 万美元，总训练成本达 1200 万美元。

除了训练，推理也很花钱。有人估算，现在 OpenAI 运行 ChatGPT 的算力费用每天就有 10 万美元。

在发展技术，让大模型掌握更多能力的同时，也有人在尝试降低 AI 所需的算力资源。最近，一种名为 FlexGen 的技术因为「一块 RTX 3090 跑 ChatGPT 体量模型」而获得了人们的关注。

虽然 FlexGen 加速后的大模型看起来仍然很慢 —— 跑 1750 亿

参数

的

语言模型

时每秒 1 个 token，但令人印象深刻的是，它已经把不可能变成了可能。

传统上，大

语言模型

（LLM）推理的高计算和内存要求使人们必须使用多个高端 AI 加速器进行训练。本研究探索了如何将 LLM 推理的要求降低到一个消费级 GPU 并实现实用性能。

近日，来自斯坦福大学、UC Berkeley、苏黎世联邦理工学院、Yandex、莫斯科国立高等经济学院、Meta、卡耐基梅隆大学等机构的新研究提出了 FlexGen，这是一种用于运行有限 GPU 内存的 LLM 的高吞吐量生成引擎。

通过聚合来自 GPU、CPU 和磁盘的内存和计算，FlexGen 可以在各种硬件资源限制下灵活配置。通过

线性规划


优化器

，它搜索存储和访问

张量

的最佳模式，包括

权重

、激活和注意力键 / 值（KV）缓存。FlexGen 将

权重

和 KV 缓存进一步压缩到 4 位，精度损失低到可以忽略不计。与最先进的 offloading 系统相比，FlexGen 在单个 16GB GPU 上运行 OPT-175B 的速度提高了 100 倍，并首次实现了 1 token/s 的实际生成吞吐量。如果提供了更多的分布式 GPU，FlexGen 还带有流水线并行 runtime，以允许在解码时进行超线性扩展。

目前，该技术已经放出代码，获得了几千 Star 量：https://github.com/FMInference/FlexGen

简介

近年来，大

语言模型

在广泛的任务中表现出卓越的性能。LLM 在展现出前所未有的通用智能的同时，也让人们在构建时面临着前所未有的挑战。这些模型可能有数十亿甚至数万亿个

参数

，这导致运行它们需要极高的计算和内存要求。例如，GPT-175B（GPT-3）仅用于存储模型

权重

就需要 32

5G

B 的内存。要让此模型进行推理，至少需要五块英伟达 A100（80GB）和复杂的并行策略。

降低 LLM 推理资源需求的方法是最近人们经常讨论的内容。这些努力分为三个方向：

（1）模型压缩以减少总内存占用量；

（2）协同推理，通过去中心化分摊成本；

（3）Offloading 以利用 CPU 和磁盘的内存。

这些技术显着降低了使用 LLM 的计算资源需求。然而，人们通常假设模型适合 GPU 内存，而现有的基于 offloading 的系统仍然难以使用单块 GPU 以可接受的吞吐量运行 1750 亿

参数

规模的模型。

在新研究中，作者专注于高吞吐量生成推理的有效 offloading 策略。当 GPU 显存不够用时，我们需要将其卸载到二级存储，通过部分加载的方式，逐段进行计算。在典型的机器上，内存层次结构分为三级，如下图所示。高级内存速度快但稀缺，低级内存速度慢但充裕。

在 FlexGen 中，作者不追求低延迟，而是瞄准面向吞吐量的场景，这些场景在

基准

测试、信息提取、数据整理等应用中很受欢迎。实现低延迟对于 offloading 来说本质上是一个挑战，但是对于面向吞吐量的场景，可以大大提高 offloading 的效率。图 1 说明了三个具有 offloading 的推理系统的延迟吞吐量权衡。通过仔细的

调度

，I/O 成本可以通过大量输入分摊并与计算重叠。在研究中，作者展示了就单位算力成本而言，单块消费级 GPU 吞吐量优化的 T4 GPU 效率要比云上延迟优化的 8 块 A100 GPU 的效率高 4 倍。

图 1. OPT-175B（左）和 OPT-30B（右）上三个基于 offloading 的系统的延迟和吞吐量权衡。FlexGen 实现了新的帕累托最优边界，OPT-175B 的最大吞吐量提高了 100 倍。由于内存不足，其他系统无法进一步提高吞吐量。

尽管已有研究在训练的背景下讨论了 offloading 的延迟 – 吞吐量权衡，但尚未有人将其用于生成 LLM 推理，这是一个截然不同的过程。由于 LLM 的自回归性质，生成推理提出了独特的挑战。除了存储所有

参数

外，它还需要顺序解码并维护一个大的注意力键 / 值缓存（KV 缓存）。现有的 offload 系统都无法应对这些挑战，因此它们执行过多的 I/O，只能实现远低于硬件能力的吞吐量。

为生成推理设计良好的 offloading 策略具有一定挑战性。首先，这个过程中存在三种

张量

：

权重

、激活和 KV 缓存。该策略应指定在三级层次结构上的卸载内容、位置以及卸载时机。其次，逐个 batch、逐个 token 和逐个 layer 计算的结构形成了一个复杂的依赖图，可以通过多种方式进行计算。该策略应该选择一个可以最小化执行时间的时间表。这些选择共同构成了一个复杂的设计空间。

为此，在新方法 FlexGen 上，人们提出了一种用于 LLM 推理的 offloading 框架。FlexGen 聚合来自 GPU、CPU 和磁盘的内存，并能有效地

调度

I/O 操作，作者也讨论了可能的压缩方法和分布式管道并行性。

该研究的主要贡献如下：

1、作者正式定义了可能的 offloading 策略的搜索空间，并使用成本模型和

线性规划

求解器搜索最佳策略。值得关注的是，研究人员证明了搜索空间捕获了一个几乎 I/O 最优的计算顺序，其 I/O 复杂度在最优计算顺序的 2 倍以内。搜索算法可以针对各种硬件规格和延迟 / 吞吐量限制进行配置，从而提供一种平滑导航权衡空间的方法。与现有策略相比，FlexGen 解决方案统一了

权重

、激活和 KV 缓存的放置，从而实现了更大的 batch size。

2、研究表明，可以将 OPT-175B 等 LLM 的

权重

和 KV 缓存压缩到 4 位，而无需重新训练 / 校准，精度损失可忽略不计。这是通过细粒度分组

量化

实现的，可以显著降低 I/O 成本。

3、通过在英伟达 T4 GPU (16GB) 上运行 OPT-175B 来展示 FlexGen 的效率。在单块 GPU 上，给定相同的延迟要求，与 DeepSpeed Zero-Inference (Aminabadi et al., 2022) 和 Hugging Face Accelerate (HuggingFace, 2022) 相比，不压缩的 FlexGen 可以实现高出 65 倍的吞吐量，后者是目前业内最先进的基于 offloading 的推理系统。如果允许更高的延迟和压缩，FlexGen 可以进一步提高吞吐量并达到 100 倍的改进。FlexGen 是第一个可以使用单块 T4 GPU 为 OPT-175B 实现 1 token/s 速度吞吐量的系统。如果给定多块分布式 GPU，具有流水线并行性的 FlexGen 可在解码时实现超线性扩展。

在研究中，作者还将 FlexGen 和 Petals 作为 offloading 和去中心化集合推理方法的代表进行了比较。结果表明，具有单块 T4 GPU 的 FlexGen 在吞吐量方面胜过具有 12 块 T4 GPU 的分散式 Petal 集群，并且在某些情况下甚至可以实现更低的延迟。

运行机制

通过聚合来自 GPU、CPU 和磁盘的内存和计算，FlexGen 可以在各种硬件资源限制下灵活配置。通过

线性规划


优化器

，它搜索存储和访问

张量

的最佳模式，包括

权重

、激活和注意力键 / 值 (KV) 缓存。FlexGen 将

权重

和 KV 缓存进一步压缩到 4 位，精度损失可以忽略不计。

FlexGen 的一个关键思想是进行延迟 – 吞吐量权衡。实现低延迟对于卸载方法来说本来就具有挑战性，但对于面向吞吐量的场景，可以极大地提升卸载效率（见下图）。FlexGen 利用块

调度

来重用

权重

并将 I/O 与计算重叠，如下图 (b) 所示，而其他基线系统使用低效的逐行

调度

，如下图 (a) 所示。

目前，该研究作者的下一步计划包括对苹果 M1、M2 芯片的支持和 Colab 部署的支持。

FlexGen 自发布后在 GitHub 上的 Star 量很快上千，在社交网络上热度也很高。人们纷纷表示这个项目很有前途，似乎运行高性能大型

语言模型

的障碍正在被逐渐克服，希望在今年之内，单机就能搞定 ChatGPT。

有人用这种方法训练了一个

语言模型

，结果如下：

虽然没有经过大量数据的投喂，AI 不知道具体知识，但回答问题的

逻辑

似乎比较清晰，或许未来的游戏中，我们能看见这样的 NPC？

参考内容：


https://news.ycombinator.com/item?id=34869960