FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU

发表时间: 2023-07 · ICML 2023

原文: https://proceedings.mlr.press/v202/sheng23a.html

作者/机构：Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Beidi Chen, Percy Liang, Christopher Re, Ion Stoica, Ce Zhang

速读

一句话结论 本文提出了一个名为 FlexGen 的高吞吐量大语言模型生成引擎，通过统筹调度 GPU、CPU 和磁盘的内存与计算，并结合张量压缩技术，首次实现在单张 16GB 显存的 GPU 上以 1 token/s 的吞吐量运行 OPT-175B 模型。

要解决什么问题 大语言模型推理面临极高的显存墙瓶颈，例如 175B 参数的模型仅加载权重就需要 325GB 显存，通常必须依赖多张高端 GPU。在离线批处理等对延迟不敏感但对吞吐量要求极高的场景下，工程师希望用单张平价 GPU 跑通大模型。现有的卸载（Offloading）系统（如 DeepSpeed Zero-Inference 和 Hugging Face Accelerate）直接沿用了训练阶段的策略，采用逐层计算的调度方式。这种方式在生成任务中会导致极其低效的 I/O 调度，因为每次生成一个 token 都要把整个模型的权重在设备间搬运一次。为了防止显存溢出，这些系统只能将批处理大小限制在 1 或 2。更严重的是，在长序列和大批处理大小下，KV cache（键值缓存）会急剧膨胀，例如 OPT-175B 在批处理大小为 512 时，KV cache 高达 1.2TB，是权重的 3.8 倍。现有系统无法有效处理这种规模的 KV cache，导致 I/O 成本无法被摊薄，吞吐量远低于硬件理论上限。

怎么做的 FlexGen 的核心思路是将生成推理过程建模为一个在 GPU、CPU 和磁盘三级存储架构下的图遍历问题，通过寻找最优的张量存放位置和计算顺序来最大化吞吐量。首先，它设计了锯齿状块调度（Zig-zag block schedule）机制。传统方法逐层计算导致每层权重被反复加载；FlexGen 则将一批请求组合成一个块，加载某层权重后，让块内所有请求都在该层完成计算，从而摊薄权重 I/O 开销。其次，它构建了一个包含计算调度、张量放置和计算委托的搜索空间。通过建立解析代价模型，预估预填充和解码阶段的延迟，FlexGen 将寻找最优卸载策略转化为一个线性规划问题。其优化目标是最小化每个 token 的平均生成延迟，即：

$$ \min \frac{T}{bls} $$

其中 $T$ 是计算一个块的总延迟，$bls$ 是块大小。约束条件为各级设备内存不超载，决策变量则是权重、激活值和 KV cache 在三级存储上的存放比例。在这个框架下，当 KV cache 存放在 CPU 时，FlexGen 会将注意力分数的计算委托给 CPU 执行，避免将庞大的 KV cache 搬运回 GPU。最后，为了进一步压缩 I/O 成本，FlexGen 引入了细粒度的分组非对称量化技术，将权重和 KV cache 均压缩至 4 bit。它以 64 为组大小，对权重沿输出通道维度分组，对 KV cache 沿隐藏层维度分组，在不重新训练的情况下减少了内存占用，从而允许系统使用更大的批处理大小。

效果如何 实验硬件为单张 NVIDIA T4 (16GB) GPU、208GB CPU 内存和 1.5TB 固态硬盘，测试模型为 OPT 系列。对比的基线方法包括代表单机卸载路线的 DeepSpeed Zero-Inference 和 Hugging Face Accelerate，以及代表分布式协同推理路线的 Petals。在输入 512、输出 32 的设置下运行 OPT-175B，若限制总延迟 5000 秒，FlexGen 可使用 64 的有效批处理大小，吞吐量比 DeepSpeed（批处理大小为 1）高 40 倍以上，而 Hugging Face Accelerate 直接内存溢出。如果放宽延迟限制到 12000 秒，FlexGen 能将有效批处理大小扩大至 256，最大吞吐量比基线高出 69 倍。在开启 4 bit 压缩后，FlexGen 能够将所有权重和缓存塞进 CPU 内存而完全摆脱缓慢的磁盘 I/O，以 144 的有效批处理大小实现了 100 倍的吞吐量提升，达到了 1.12 token/s。在真实任务测试中，FlexGen 仅用 21 小时就在单张 16GB GPU 上完成了 30B 模型在 HELM 基准测试 7 个子场景下的评估。对比 Petals，FlexGen 的单卡吞吐量在各种网络条件下均超越了后者的单卡平均吞吐量。不过，这种极致的吞吐量提升是以牺牲单次响应延迟为代价的，仅适用于离线处理场景。此外作者也承认，由于 4 bit 压缩在 CPU 上的解压开销过大，开启量化时必须关闭 CPU 计算委托功能。

主要贡献

大型语言模型（LLM）推理具有极高的计算和内存需求，通常只有使用多个高端加速器才可行。例如，加载 GPT-175B 的模型权重就需要 325GB 的 GPU 内存，这至少需要 5 张 A100 (80GB) GPU 并配合复杂的并行策略。由于当前对支持批量处理且对延迟不敏感的任务（如基准测试、信息提取、数据整理等）的需求不断涌现，本文旨在研究如何使用有限的资源（例如单张商用 GPU）实现高吞吐量的 LLM 推理。

本文提出了 FlexGen，这是一个在有限 GPU 内存下运行 LLM 的高吞吐量生成引擎。FlexGen 的核心创新点包括：
1. 统一的离线调度策略：通过聚合 GPU、CPU 和磁盘的内存与计算资源，FlexGen 可以在各种硬件资源限制下灵活配置。它通过定义包含计算调度、张量放置和计算委托的搜索空间，并求解线性规划问题，来寻找存储和访问张量的高效模式。该搜索空间捕获的计算顺序其 I/O 复杂度在最优解的 $2 \times$ 以内。
2. 高效的压缩技术：FlexGen 能够将权重和注意力缓存（KV cache）压缩到 4 bits，而精度损失可忽略不计，且无需重新训练或校准。这进一步减少了 I/O 成本和内存占用。
3. 显著的吞吐量提升：这些技术使 FlexGen 拥有更大的批处理大小（Batch Size）选择空间，从而显著提高最大吞吐量。在单张 16GB GPU 上运行 OPT-175B 时，FlexGen 相比最先进的卸载（offloading）系统实现了显著更高的吞吐量，首次在有效批处理大小为 144 的情况下达到了 1 token/s 的生成吞吐量。在 HELM 基准测试中，FlexGen 可以在 21 小时内使用 16GB GPU 对 30B 模型完成 7 个代表性子场景的基准测试。

图1. OPT-175B（左）和OPT-30B（右）在单张NVIDIA T4 (16 GB) GPU配备208 GB CPU DRAM和1.5 TB SSD上的块总延迟与吞吐量权衡。FlexGen实现了新的帕累托最优边界，OPT-175B的最大吞吐量提高了100倍。其他系统由于内存不足问题无法进一步提高吞吐量。"(c)"表示压缩。
图1. OPT-175B（左）和OPT-30B（右）在单张NVIDIA T4 (16 GB) GPU配备208 GB CPU DRAM和1.5 TB SSD上的块总延迟与吞吐量权衡。FlexGen实现了新的帕累托最优边界，OPT-175B的最大吞吐量提高了100倍。其他系统由于内存不足问题无法进一步提高吞吐量。"(c)"表示压缩。

背景知识

生成式推理过程：典型的 LLM 生成式推理任务包含两个阶段：i) 预填充（prefill）阶段，该阶段接收一个提示（prompt）序列并为 LLM 的每个 Transformer 层生成键值缓存（KV cache）；ii) 解码（decoding）阶段，该阶段利用并更新 KV cache 来逐步生成 token，其中当前 token 的生成依赖于先前生成的 tokens。

推理计算公式：对于特定的推理计算，设批处理大小为 $b$，输入序列长度为 $s$，输出序列长度为 $n$，Transformer 的隐藏维度为 $h_1$，第二个 MLP 层的隐藏维度为 $h_2$，Transformer 的总层数为 $l$。给定一个 Transformer 层的权重矩阵 $\mathbf{w}_K^i, \mathbf{w}_Q^i, \mathbf{w}_V^i, \mathbf{w}_O^i, \mathbf{w}_1^i, \mathbf{w}_2^i$，其中 $\mathbf{w}_K^i, \mathbf{w}_Q^i, \mathbf{w}_V^i, \mathbf{w}_O^i \in \mathcal{R}^{h_1 \times h_1}$，$\mathbf{w}_1^i \in \mathcal{R}^{h_1 \times h_2}$，且 $\mathbf{w}_2^i \in \mathcal{R}^{h_2 \times h_1}$。
在 prefill 阶段，第 $i$ 层的输入由 $\mathbf{x}^i$ 指定，注意力层的 key、value、query 和 output 分别由 $\mathbf{x}_K^i, \mathbf{x}_V^i, \mathbf{x}_Q^i, \mathbf{x}_{\mathrm{Out}}^i$ 指定，其中这些变量均属于 $\mathcal{R}^{b \times s \times h_1}$。缓存的 key 和 value 计算如下：

$$ \mathbf{x}_K^i = \mathbf{x}^i \cdot \mathbf{w}_K^i ; \quad \mathbf{x}_V^i = \mathbf{x}^i \cdot \mathbf{w}_V^i $$


第 $i$ 层的其余计算为：

$$ \mathbf{x}_Q^i = \mathbf{x}^i \cdot \mathbf{w}_Q^i $$
$$ \mathbf{x}_{\mathrm{Out}}^i = f_{\mathsf{Softmax}} \left( \frac{\mathbf{x}_Q^i (\mathbf{x}_K^i)^T}{\sqrt{h_1}} \right) \cdot \mathbf{x}_V^i \cdot \mathbf{w}_O^i + \mathbf{x}^i $$
$$ \mathbf{x}^{i+1} = f_{\mathsf{relu}} \left( \mathbf{x}_{\mathrm{Out}}^i \cdot \mathbf{w}_1^i \right) \cdot \mathbf{w}_2^i + \mathbf{x}_{\mathrm{Out}}^i $$
在 decode 阶段，给定 $\mathbf{t}^i \in \mathcal{R}^{b \times 1 \times h_1}$ 作为第 $i$ 层当前生成 token 的嵌入，推理计算需要 i) 更新 KV cache：
$$ \mathbf{x}_K^i \leftarrow \mathsf{Concat}(\mathbf{x}_K^i, \mathbf{t}^i \cdot \mathbf{w}_K^i) $$
$$ \mathbf{x}_V^i \leftarrow \mathsf{Concat}(\mathbf{x}_V^i, \mathbf{t}^i \cdot \mathbf{w}_V^i) $$
并 ii) 计算当前层的输出：
$$ \mathbf{t}_Q^i = \mathbf{t}^i \cdot \mathbf{w}_Q^i $$
$$ \mathbf{t}_{\mathrm{Out}}^i = f_{\mathsf{Softmax}} \left( \frac{\mathbf{t}_Q^i (\mathbf{x}_K^i)^T}{\sqrt{h_1}} \right) \cdot \mathbf{x}_V^i \cdot \mathbf{w}_O^i + \mathbf{t}^i $$
$$ \mathbf{t}^{i+1} = f_{\mathsf{relu}} \left( \mathbf{t}_{\mathrm{Out}}^i \cdot \mathbf{w}_1^i \right) \cdot \mathbf{w}_2^i + \mathbf{t}_{\mathrm{Out}}^i $$

内存瓶颈分析：LLM 推理的内存占用主要来自模型权重和 KV cache。对于 FP16 精度的 OPT-175B 模型，存储参数的总字节数大约为 $l(8h_1^2 + 4h_1h_2)$。存储 KV cache 的峰值总字节数为 $4 \times blh_1(s + n)$。在真实场景中，OPT-175B（$l = 96, h_1 = 12288, h_2 = 49152$）占用 325 GB 内存。当批处理大小 $b = 512$，输入序列长度 $s = 512$，输出序列长度 $n = 32$ 时，存储 KV cache 所需的总内存为 1.2 TB，这是模型权重的 $3.8 \times$，使得 KV cache 成为大批量高吞吐量推理的新瓶颈。在 FlexGen 中，为了使 OPT-175B 达到 0.69 token/s 的吞吐量，有效批处理大小被扩大到了 256。

吞吐量与延迟定义：考虑有效批处理大小 $b$、输入序列长度 $s$ 和输出序列长度 $n$，延迟 $t$ 定义为处理提示并生成所有 $bn$ 个 tokens 所花费的总秒数。生成吞吐量定义为 $bn/t$。

图2. LLM推理的计算图。
图2. LLM推理的计算图。

方法细节

4. 卸载策略 (Offloading Strategy)

问题形式化定义：考虑一台配备 GPU、CPU 和磁盘三个设备的机器。其中 GPU 和 CPU 可以执行计算，而磁盘不能。这三个设备构成了一个三级内存层次结构，GPU 的内存最小但速度最快，而磁盘的内存最大但速度最慢。当一个 LLM 无法完全装入 GPU 时，我们需要将其卸载到二级存储中，并通过部分加载 LLM 来逐部分执行计算。
三级内存层次结构

计算图遍历约束：我们将带有卸载的生成式推理形式化为一个图遍历问题。如计算图（Fig 2）所示，模型有 4 层，每个 prompt 生成 3 个 token。由于目标是面向吞吐量的场景，假设给定数据集中有无限数量的 prompt 需要处理。图中一个方块代表对一个层的 GPU 批次进行计算，相同颜色的方块共享相同的层权重。有效路径定义为遍历（即计算）所有方块的路径，同时受以下条件约束：
* 只有当同一行中位于某方块左侧的所有方块都已计算完毕时，该方块才能被计算。
* 为了在某个设备上计算一个方块，其所有输入（权重、激活值、缓存）必须加载到同一个设备上。
* 计算完成后，一个方块产生两个输出：激活值和 KV cache。激活值应被存储，直到其右侧相邻的方块被计算完毕。KV cache 应被存储，直到同一行中最右侧的方块被计算完毕。
* 在任何时候，存储在一个设备上的张量总大小不能超过其内存容量。
目标是寻找一条使总执行时间（包括计算成本和在设备之间移动张量的 I/O 成本）最小化的有效路径。

计算调度设计：直观来看，遍历计算图有两种顺序：逐行（row-by-row）和逐列（column-by-column）。所有现有系统【1，Deepspeed-inference: Enabling efficient inference of transformer models at unprecedented scale，2022，SC22】、【12，Hugging face accelerate，2022，https://huggingface.co/docs/accelerate/index】均采用逐行遍历， 如 Fig 3(a) 所示。这种做法是合理的，因为它是完成一个批次生成的最快方式，且 KV cache 可以在一行结束后立即释放。然而，由于每两个相邻的方块不共享权重，这种调度必须重复加载权重，导致巨大的 I/O 成本。

降低权重 I/O 的列遍历与块调度：为了减少权重的 I/O 成本，我们可以逐列遍历计算图。一列中的所有方块共享权重，因此我们可以让权重保留在 GPU 上以供重用，并仅加载/卸载激活值和 KV cache。但是，我们不能一直遍历到底，因为激活值和 KV cache 仍然需要被存储。因此，当它们填满 CPU 和磁盘内存时，我们必须停止。综合考虑这些因素，我们收敛到一种“之字形块调度”（zig-zag block schedule），如 Fig 3(b) 所示。此外，我们提出了另一种更先进且 I/O 最优的调度策略，但由于最优策略在实际实现中的困难，我们仅实现了较简单的块调度。不过，我们在附录 A.2 中证明了块调度的性能最多比最优调度差两倍。

调度重叠与算法实现：另一种典型的优化是重叠（overlapping）。我们可以将下一层的权重加载、下一个批次的缓存/激活加载、上一个批次的缓存/激活存储以及当前批次的计算进行重叠。将重叠添加到块调度中就得到了 Algorithm 1。最内层循环的前六个函数可以看作是由六个逻辑线程并行启动的，因为它们之间没有依赖关系。最后一个函数随后同步这六个逻辑线程。我们依赖操作系统和 CUDA 驱动程序来解决底层硬件资源的调度问题。结论是，该算法在我们的搜索空间中引入了两个参数：GPU 批处理大小（GPU batch size）和一个块中的 GPU 批次数量。GPU 批处理大小与 GPU 批次数量的乘积被称为块大小（block size，或有效批处理大小 effective batch size）。

# Algorithm 1 Block Schedule with Overlapping
for i = 1 to generation_length do 
    for j = 1 to num_layers do 
        // Compute a block with multiple GPU batches 
        for k = 1 to num_GPU_batches do 
            // Load the weight of the next layer 
            load_weight(i, j + 1, k) 
            // Store the cache and activation of the prev batch 
            store_activation(i, j, k - 1) 
            store_cache(i, j, k - 1) 
            // Load the cache and activation of the next batch 
            load_cache(i, j, k + 1) 
            load_activation(i, j, k + 1) 
            // Compute this batch 
            compute(i, j, k) 
            // Synchronize all devices 
            synchronize() 
        end for 
    end for 
end for
图3. 两种不同的调度策略。红色箭头表示计算顺序。
图3. 两种不同的调度策略。红色箭头表示计算顺序。

张量放置策略：除了计算调度外，策略还应指定如何在内存层次结构中存储这些张量。我们使用三个变量 $wg, wc, wd$ 分别定义存储在 GPU、CPU 和磁盘上的权重百分比。同样，我们使用 $hg, hc, hd$ 定义激活值的百分比，使用 $cg, cc, cd$ 定义 KV cache 的百分比。给定百分比后，仍然有多种划分张量的方法。以权重张量为例，从粗粒度到细粒度，我们可以按模型粒度划分权重（例如，将模型中 50% 的层分配给 GPU）、按层粒度划分（例如，将一层中 50% 的张量分配给 GPU），或者按张量粒度划分（例如，将一个张量中 50% 的元素分配给 GPU）。较粗的粒度会导致较低的运行时开销，但灵活性较差且成本难以分析。综合考虑运行时开销和所需的灵活性，我们对权重使用层粒度，对激活值和 KV cache 使用张量粒度。

计算委托（CPU 计算）机制：虽然 CPU 比 GPU 慢得多，但我们发现在某些情况下使用 CPU 计算仍然是有益的。这是因为解码期间注意力分数的计算是 I/O 密集型的。考虑 KV cache 存储在 CPU 上的情况。在 GPU 上计算注意力分数需要将整个 KV cache 移动到 GPU，由于 KV cache 非常庞大，这会产生大量的 I/O 成本。相比之下，在 CPU 上计算注意力分数不需要移动 KV cache，它只需要将激活值从 GPU 移动到 CPU。定量地看，设 $b$ 为 GPU 批处理大小，$s$ 为序列长度，$h_1$ 为隐藏层大小。移动的 KV cache 大小为 $b \times s \times h_1 \times 4$ 字节，而移动的激活值大小为 $b \times h_1 \times 4$ 字节，因此在 CPU 上计算注意力分数可将 I/O 减少 $s \times$。对于长序列（例如 $s \geq 512$），如果相关的 KV cache 未存储在 GPU 上，最好在 CPU 上计算注意力分数。

成本模型构建：上述调度和放置策略构建了一个包含多个参数的搜索空间。现在我们开发一个分析性成本模型，以在给定这些算法参数和硬件规格的情况下估算执行时间。成本模型预测一个块中预填充期间一层的延迟（记为 $T_{pre}$）以及解码期间一层的平均延迟（记为 $T_{gen}$）。然后，计算一个块的总延迟可估算为 $T = T_{pre} \cdot l + T_{gen} \cdot (n - 1) \cdot l$，其中 $l$ 是层数，$n$ 是要生成的 token 数。假设完美重叠，$T_{pre}$ 可估算为 $T_{pre} = \max(ctog^p, gtoc^p, dtoc^p, ctod^p, comp^p)$，其中 $ctog^p, gtoc^p, dtoc^p, ctod^p, comp^p$ 分别表示预填充期间一层的 CPU 到 GPU 读取、GPU 到 CPU 写入、磁盘到 CPU 读取、CPU 到磁盘写入以及计算的延迟。同理，$T_{gen}$ 可估算为 $T_{gen} = \max(ctog^g, gtoc^g, dtoc^g, ctod^g, comp^g)$，各变量含义类似，针对解码阶段。

I/O 与计算项估算：对于像 $dtoc^g$ 这样的 I/O 项，它是通过累加 I/O 事件来估算的，这些事件包含权重、激活值和缓存的读取。一个 Transformer 层的 FP16 权重大小为 $8h_1^2 + 4h_1 \cdot h_2$ 字节，其中 $h_1$ 表示隐藏大小，$h_2$ 表示第二个 MLP 层的隐藏大小。设 $bls$ 为块大小，$s$ 为提示长度；则一层激活值的大小为 $2 \cdot bls \cdot h_1$。一层 KV cache 的平均大小为 $4 \cdot bls \cdot (s + \frac{n}{2}) \cdot h_1$。我们必须从磁盘分别加载 $wd, hd, cd$ 比例的权重、激活值和 KV cache，因此磁盘读取的总延迟为：

$$ dtoc^g = \frac{1}{\mathrm{disk\_to\_cpu\_bandwidth}} ((8h_1^2 + 4h_1 \cdot h_2) \cdot wd + 4 \cdot bls \cdot (s + \frac{n}{2}) \cdot h_1 \cdot cd + 2 \cdot bls \cdot h_1 \cdot hd) $$


对于计算项，我们同样累加所有计算事件，包括 CPU 和 GPU 上的矩阵乘法和批量矩阵乘法。除了延迟估算，我们还估算 GPU、CPU 和磁盘的峰值内存使用量，并随后添加内存约束（详细模型见附录 A.3）。

策略搜索过程：一个策略包含 11 个变量：块大小 $bls$、GPU 批处理大小 $gbs$、权重放置 $wg, wc, wd$、激活放置 $hg, hc, hd$ 以及 KV cache 放置 $cg, cc, cd$。在实践中，百分比不能是 0 到 1 之间的任意实数，因为张量不能任意拆分。然而，由于百分比是渐变的，我们在成本模型中放宽这些变量为 0 到 1 之间的任意实数。我们将该问题作为一个双层优化问题来解决。首先，我们枚举几个 $(bls, gbs)$ 元组的组合。通常 $gbs$ 是 4 的倍数，而 $bls$ 小于 20，因此组合并不多。然后，在固定 $bls, gbs$ 的情况下，寻找最佳放置 $\boldsymbol{p} = (wg, wc, wd, cg, cc, cd, hg, hc, hd)$ 就变成了一个线性规划问题（如下式）。由于只有 9 个变量，该线性规划问题可以非常快地求解。这种公式化方法还可以灵活扩展以包含延迟约束和模型近似方法（如压缩）。

$$ \min_{\boldsymbol{p}} T / bls $$

$$ \text{s.t. } \text{gpu peak memory} < \text{gpu mem capacity} $$
$$ \text{cpu peak memory} < \text{cpu mem capacity} $$
$$ \text{disk peak memory} < \text{disk mem capacity} $$
$$ wg + wc + wd = 1 $$
$$ cg + cc + cd = 1 $$
$$ hg + hc + hd = 1 $$
为了使用成本模型，我们在硬件上运行分析工具以采样一些数据点并拟合硬件参数。然后调用优化器获取卸载策略。由于我们的松弛处理以及精确建模峰值内存使用量（例如内存碎片）的困难，有时策略搜索得出的策略会导致内存不足。在这种情况下，我们会手动对策略进行微调。成本模型通常能返回一个良好的策略，但通过手动调整通常能获得更好的策略。

多 GPU 扩展：我们讨论如何在有多张 GPU 的情况下扩展 FlexGen 的卸载策略。虽然我们能为单张 GPU 找到近乎最优的策略，但该策略仍严重受限于 I/O 且 GPU 利用率较低。如果提供更多的 GPU 和 CPU，可以利用模型并行化来减轻每张 GPU 的内存压力，这有可能在解码吞吐量上实现超线性扩展。模型并行化分为张量并行和流水线并行两种。张量并行可以减少单查询延迟，而流水线并行由于通信成本低，可以在吞吐量上实现良好的扩展。因为我们以吞吐量为目标，FlexGen 实现了流水线并行。我们通过在 $m$ 张 GPU 上等分一个 $l$ 层的 LLM 来使用流水线并行，所有 GPU 的执行遵循相同的模式。问题被简化为在一张 GPU 上运行一个 $n/m$ 层的 Transformer。我们可以直接复用为单张 GPU 开发的策略搜索。为了实现微批次流水线（micro-batch pipelining），在 Algorithm 1 中添加了一个新的 for 循环，以将迭代级流水线并行执行调度与我们的单设备卸载运行时结合起来。

5. 近似方法 (Approximate Methods)

组级量化（Group-wise Quantization）：上一节重点关注精确计算。然而，通过允许某些近似计算，推理吞吐量可以得到极大提升，且精度损失可忽略不计，因为 LLM 通常对谨慎的近似具有鲁棒性。我们展示了 OPT-175B 的权重和 KV cache 都可以直接量化为 4-bit 整数，而无需任何重新训练或校准，同时保持相似的精度。与一些主要尝试使用整数矩阵乘法来加速计算的相关工作相比，我们量化的目标主要是为了压缩和降低 I/O 成本。因此，我们可以选择有利于高压缩比的细粒度量化格式，并在计算前将张量反量化回 FP16。我们使用了一种细粒度的组级非对称量化方法【28，Q-bert: Hessian based ultra low precision quantization of bert，2020，AAAI】。给定一个张量，我们沿某个特定维度选择 $g$ 个连续元素作为一个组。对于每个组，我们计算组内元素的最小值和最大值，并将每个元素 $x$ 量化为 $b$-bit 整数：

$$ x_{quant} = round\Big(\frac{x - min}{max - min} \times (2^b - 1)\Big) $$


张量以量化格式存储，并在计算前转换回 FP16。由于权重和 KV cache 都消耗大量内存，我们将两者都压缩为 4 bits，组大小设为 64。有多种方法可以选择沿哪个维度进行分组。我们发现沿输出通道维度对权重进行分组，沿隐藏维度对 KV cache 进行分组，在实践中既能保持精度又具有运行时效率。需要提及的一点是，FlexGen 中这种细粒度的组级量化在压缩和解压缩时会产生一些开销。如果运行在 CPU 上，这种开销可能会非常显著，使得 CPU 委托变得毫无用处，因此我们在启用量化时会关闭 CPU 委托。

稀疏注意力（Sparse Attention）：我们证明了可以通过仅加载 OPT-175B 中前 10% 的注意力 value 缓存来利用自注意力的稀疏性，同时保持模型质量。我们提出了一种简单的 Top-K 稀疏近似方法。在计算出注意力矩阵后，对于每个 query，我们计算其在 K cache 中 Top-K 个 token 的索引。然后，我们直接丢弃其他 token，并仅根据这些索引加载 V cache 的一个子集。这些近似方法的应用非常直接，FlexGen 是一个通用框架，可以无缝插入许多近似方法。

实验环境

实验结果

最大吞吐量基准测试：

延迟-吞吐量权衡：
- 实验内容：配置系统在不同延迟约束下实现最大吞吐量，绘制权衡曲线（Fig 1）。
- 实验结果：FlexGen 设定了新的帕累托最优边界。在 5000 秒的延迟要求下，无压缩的 FlexGen 相比 DeepSpeed 和 Accelerate 实现了 $40 \times$ 的吞吐量提升。若允许更高延迟和压缩，FlexGen 通过将所有内容装入 CPU 内存，可达到 $100 \times$ 的改进。

运行时分解与消融实验：
- 实验内容：分析 OPT-175B 在 FlexGen 上的运行时分解，并逐一禁用优化技术以评估其贡献（Table 4）。
- 实验结果：GPU 计算利用率在 prefill 阶段为 82%，在 decoding 阶段为 13%。消融实验表明，缺乏良好的策略搜索（"No policy search"）会导致 OPT-175B 吞吐量从 0.69 降至 0.27；禁用重叠（No overlapping）或禁用 CPU 计算（No CPU compute）均带来显著的性能下降。若直接将 DeepSpeed 的策略移植到 FlexGen 运行时，OPT-175B 吞吐量仅为 0.01 token/s，证明了现有策略的次优性。

HELM 与数据整理任务：
- 实验内容：使用 FlexGen 评估未包含在官方 HELM 发布中的新模型 OPT-IML-30B，以及执行数据整理任务。
- 实验结果：FlexGen 在包含所有系统开销的情况下，于 21 小时内完成了 7 个代表性子场景的基准测试。

近似方法验证：
- 实验内容：在 Lambada 和 WikiText 数据集上测试 4-bit 压缩和 4-bit + 10% 稀疏注意力（4-bit-S）对精度的影响（Table 5）。
- 实验结果：两种近似方法在 OPT-30B 和 OPT-175B 上均表现出与 FP16 几乎相同的准确率（Acc）和困惑度（PPL），证明了 LLM 对这些近似的鲁棒性。

卸载与协作推理对比：
- 实验内容：在不同网络延迟和带宽下，对比单 GPU 的 FlexGen 与 4 节点（每节点 1 GPU）的 Petals 去中心化集群在 OPT-30B 上的性能（Fig 4）。
- 实验结果：FlexGen 的单 GPU 吞吐量在所有测试网络条件下均优于 Petals 集群的单 GPU 平均吞吐量。Petals 未利用卸载，无法使用大批处理大小。在慢速网络和短生成任务中，由于网络带宽成为激活传输的瓶颈，FlexGen 甚至实现了比 Petals 更低的总延迟。

结论

本文介绍了 FlexGen，一个用于 LLM 推理的高吞吐量生成引擎，专门针对资源受限场景下对延迟不敏感的批量处理任务。FlexGen 通过统一的张量放置策略、之字形块调度、计算委托以及 4-bit 压缩技术，在单张商用 GPU 上显著突破了现有的吞吐量极限，为在有限硬件上部署和基准测试超大规模语言模型提供了一种高效可行的解决方案。

附录

A.1 符号说明

详细符号含义见原始论文 Table 6，包括模型层数 $l$、序列长度 $s$ 和 $n$、块大小 $bls$、隐藏维度 $h_1, h_2$ 等。

A.2 计算调度的最优性分析

本节讨论仅当模型无法装入单张 GPU 且不使用 CPU 计算时的图遍历问题。为了生成一个 prompt 的 $n$ 个 token，我们需要将整个模型权重交换进出 $n$ 次。因此，最好将加载的权重重用于一批 prompt，以分摊权重 I/O 时间。
之字形块调度（Zig-zag block schedule）：计算第一列的 $bls$ 个样本，保存缓存和激活，然后计算第二列，直到最后一列。在一次块计算中生成 $n \cdot bls$ 个 token，需要加载 $n$ 次模型权重。峰值内存约束为 $\text{peak\_mem} = w + 2h_1 \cdot bls + 4h_1 \cdot bls \cdot l \cdot (s + n) \leq cmem$。由此推导出块大小上限 $bls_1$。
对角线块调度（Diagonal block schedule）：如 Fig 5 所示，系统在预热后，每次迭代计算一个包含多个子对角线的对角线。每次计算一个对角线，每层的权重被加载一次，生成的 token 数为 $bls$。因此，在相同的 $bls$ 下，每个 token 的 I/O 与之字形块调度相同。其峰值内存估算为 $\text{peak\_mem} = w + 2h_1 \cdot bls + \frac{4h_1 \cdot bls \cdot l(2s+n)(n-1)}{2n}$。由此推导出的块大小上限 $bls_2$ 相比 $bls_1$ 有大约 $2 \times$ 的提升（当 $n \gg s$ 时）。
Theorem 4.1 证明：通过定义工作状态（Definition A.1）和状态转移，证明了在有限内存容量下，对角线块调度在渐进意义上是 I/O 最优的（Theorem A.3）。而之字形块调度的峰值内存消耗在加载权重时不断变化，其平均内存消耗至少为容量的一半，因此其 I/O 复杂度最多是最优解的 $2 \times$。尽管对角线调度有优势，但由于 KV cache 缓冲区的动态更新导致实现复杂，FlexGen 实际实现了之字形块调度。

图5. 对角线块调度
图5. 对角线块调度

A.3 成本模型

成本模型的完整约束条件如下，目标是最小化 $T / bls$：
延迟计算：

$$ T = T_{pre} \cdot l + T_{gen} \cdot (n - 1) \cdot l $$

$$ T_{pre} = \max(ctog^p, gtoc^p, dtoc^p, ctod^p, comp^p) $$
$$ ctog^p = \frac{1}{ctog\_bdw} ((wc + wd)(8h_1^2 + 4h_1 \cdot h_2) + 2(hc + hd)s \cdot h_1 \cdot bls) $$
$$ gtoc^p = \frac{1}{gtoc\_bdw} (4(cc + cd)(s + 1)h_1 \cdot bls + 2(hc + hd)s \cdot h_1 \cdot bls) $$
$$ dtoc^p = \frac{1}{dtoc\_bdw} (wd(8h_1^2 + 4h_1 \cdot h_2) + 2hd \cdot s \cdot h_1 \cdot bls) $$
$$ ctod^p = \frac{1}{ctod\_bdw} (4cd \cdot bls \cdot (s + 1) \cdot h_1 + 2hd \cdot s \cdot h_1 \cdot bls) $$
$$ comp^p = \frac{bls(8s \cdot h_1^2 + 4s \cdot h_1 \cdot h_2)}{mm\_flops} + \frac{4bls \cdot s^2 \cdot h_1}{bmm\_flops} $$
解码阶段的 $T_{gen}$ 也有类似的基于 $ctog^g, gtoc^g, dtoc^g, ctod^g, comp^g$ 的公式，其中计算项 $comp^g = gpu\_comp^g + cpu\_comp^g$，CPU 计算项为 $cpu\_comp^g = \frac{4(cc + cd)bls \cdot (s + n/2) \cdot h_1}{cpu\_flops}$。

峰值内存约束：
- GPU 预填充期间峰值内存：$gpu\_peak^p = gpu\_home^p + gpu\_w^p < gmem$。其中 $gpu\_home^p$ 是保留在 GPU 上的权重、激活和缓存比例占用的静态内存；$gpu\_w^p$ 是工作内存，包含 QKV、Attention、Embed 和 MLP 层的中间结果的最大值。
- GPU 解码期间峰值内存：$gpu\_peak^g = gpu\_home^g + gpu\_w^g < gmem$。
- CPU 和 NVMe 峰值内存：同样分为静态保留内存和工作内存，必须分别小于 $cmem$ 和 $nmem$。

A.4 补充细节 (Tables and Additional Experimental Results)

引用文献汇总

在方法细节及后续章节中，本文引用了以下关键文献来支撑其设计与对比：
* [1, Deepspeed-inference: Enabling efficient inference of transformer models at unprecedented scale + 2022 + SC22]：在“计算调度设计”段落中被引用。原文指出所有现有系统（如 DeepSpeed）均采用逐行遍历（row-by-row）的计算图遍历顺序，这导致了重复加载权重的巨大 I/O 成本。
* [12, Hugging face accelerate + 2022 + https://huggingface.co/docs/accelerate/index]：在“计算调度设计”段落中与文献 [1] 一同被引用，作为采用次优的逐行遍历调度的现有系统的代表。
* [28, Q-bert: Hessian based ultra low precision quantization of bert + 2020 + AAAI]：在“组级量化”段落中被引用。原文描述 FlexGen 使用了该文献提出的一种细粒度的组级非对称量化方法，用于将权重和 KV cache 压缩到 4-bit 以降低 I/O 成本。
* [21, Pagecachemangagement + 2008 + https://code.google.com/archive/p/pagecache-mangagement/source/default/source]：在消融实验（A.4 补充细节）的描述中被引用。原文说明在进行策略消融研究时，开启了该工具以防止操作系统的自动磁盘缓存，从而使消融结果更加准确，尽管这带来了一些额外开销。