论文阅读：FlexGen
这篇论文讨论的是一个很现实的问题：如果我只有一张消费级/普通云 GPU，还能不能跑非常大的 LLM？更准确地说，是能不能在不追求交互式低延迟的前提下，用单卡完成大批量生成任务，并且把吞吐做得尽量高。
论文标题是 FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU，发表于 ICML 2023。它的定位不是“让 175B 模型在单张 T4 上像 ChatGPT 一样实时聊天”，而是面向另一类场景：benchmark、数据处理、信息抽取、批量文档处理等。这类任务可以接受几千秒级的单批延迟，但希望单位时间内处理尽可能多的 token。
一句话总结 FlexGen：
FlexGen 把 GPU、CPU 内存和磁盘统一看成一个分层存储系统，通过搜索 offloading 策略、重排计算顺序、压缩权重和 KV cache，让超大模型在单张小显存 GPU 上以更高吞吐完成批量生成。
1. 论文要解决什么问题？
LLM 推理的资源瓶颈主要来自两个地方：模型权重和 KV cache。
以 OPT-175B 这种规模为例，FP16 权重大约需要 325GB 显存。单张 16GB T4 显然放不下，甚至单张 80GB A100 也放不下，论文提到至少需要五张 A100（80GB）才能装下权重。传统做法是多卡并行，但这会提高硬件门槛，也引入复杂的并行策略。
如果只是把权重放到 CPU 或磁盘，需要时再搬到 GPU，看起来可以解决“放不下”的问题，但会带来另一个问题：I/O 太慢。GPU 算得快，PCIe、CPU 内存和 SSD 的带宽远低于显存带宽，如果调度不好，GPU 大部分时间都在等数据。
GPU、CPU 内存和磁盘构成一个三级存储层级：越往上带宽越高、容量越小，越往下容量越大、带宽越低。下图给出论文实验环境（单张 T4）的典型容量和带宽，FlexGen 要解决的核心问题就是如何在这个层级之间调度数据。
更麻烦的是，生成式推理和训练不太一样。一次生成通常分为两个阶段：
- Prefill：把 prompt 跑一遍，生成每层的 KV cache。
- Decode：每次生成一个新 token，并不断读取、更新 KV cache。
在小 batch、短文本场景里，模型权重可能是主要内存压力；但在大 batch、高吞吐场景里，KV cache 也会变成巨大的内存负担。论文里提到，FlexGen 为了提高吞吐会把 effective batch size 放大到很大，这时 KV cache 不再是可以忽略的中间状态，而是必须被认真调度和压缩的对象。
所以问题不是简单的“显存不够怎么办”，而是：
- 权重、activation、KV cache 分别应该放在 GPU、CPU 还是磁盘？
- 什么时机搬运这些 tensor，才能尽量减少重复 I/O？
- batch size 该设多大，才能在吞吐和延迟之间取得合理平衡？
- 如果允许近似计算，压缩哪些 tensor 最划算？
2. 现有 offloading 为什么不够好？
论文对 DeepSpeed ZeRO-Inference 和 Hugging Face Accelerate 这类 offloading 系统的一个批评是：它们很多策略继承自训练系统，没有充分利用生成式推理的结构。
直观理解，可以把生成过程画成一个网格：纵向是不断到来的多个 GPU batch，横向是模型的各层（沿 token 方向还会继续向右延伸），同一层用同一份权重。每个格子代表某个 batch 在某一层上的一次计算。
已有系统更接近“按行遍历”：先把一个 batch 推完所有层、所有 token，再处理下一个 batch。这种方式单个 batch 的延迟较低，KV cache 用完可以较快释放；但一行之内横向相邻的两步属于不同层、不共享权重，于是同一层权重会被反复从 CPU 或磁盘搬到 GPU。
对于单卡大模型 offloading 来说，重复加载权重是非常致命的。因为 OPT-175B 的权重本身已经远大于显存，只要调度方式导致频繁从慢存储读取权重，吞吐就会被 I/O 卡死。
FlexGen 的关键判断是：既然目标是 throughput-oriented inference，就应该愿意牺牲单批延迟，换取更大的 batch 和更少的重复 I/O。
3. FlexGen 的核心方法
3.1 Block schedule：先复用权重，再考虑延迟
FlexGen 没有采用传统的 row-by-row schedule，而是提出了一个 zig-zag block schedule。
简单说，它会让一组 GPU batch 组成一个 block。当某一层的权重被加载到 GPU 后，尽量让这个 block 里的多个 batch 都用上这份权重，然后再进入下一层。这样做的目标是摊薄权重加载成本。
把生成过程看成一个网格，纵向是同一个 block 里的多个 batch，横向是模型的各层（颜色相同表示同一层、共享同一份权重）。已有系统按行遍历：先把一个 batch 推完所有层，再做下一个 batch，横向相邻的两步属于不同层，于是每一步都要重新加载权重。FlexGen 改成按列遍历：一层权重加载后，让 block 里的多个 batch 都用完这份权重，再换下一层，权重只需加载一次。
按列推进的顺序是：用加载好的某层权重，自上而下扫过 block 内的每个 batch（b1 → b4），算完这一列再跳回顶端、换下一层（L1·b4 → L2·b1 …）。把这条计算顺序连起来，就是上图 (b) 中红色箭头那样的“之”字形，FlexGen 因此称它为 zig-zag block schedule。
这个策略会增加单个 batch 的完成时间，因为一个 batch 可能要等同一个 block 里的其他 batch 一起推进。但对于批量离线任务来说，这种延迟是可以接受的；换来的收益是权重 I/O 被更好地复用，整体吞吐上升。
论文还加入了 I/O 和计算重叠：加载下一层权重、读取下一个 batch 的 activation/KV cache、写回上一个 batch 的中间状态、当前 batch 的 GPU 计算，可以尽量并行发生。实现上，FlexGen 基于 PyTorch，管理多个 CUDA stream 和 CPU thread 来做 overlap。
3.2 统一决定权重、activation、KV cache 的位置
回到前面的三级存储层级，一个 offloading policy 需要决定三类 tensor 的放置比例：
- weights 放多少在 GPU / CPU / disk；
- activations 放多少在 GPU / CPU / disk；
- KV cache 放多少在 GPU / CPU / disk。
这比“权重不够就放 CPU”要细得多。FlexGen 会建立 cost model，估计不同策略下每层 prefill 和 decode 的计算时间、CPU/GPU/Disk I/O 时间以及峰值内存占用。然后它枚举 batch/block 相关参数，并把 tensor placement 部分转成线性规划问题来搜索较优策略。
这里我觉得有意思的一点是：FlexGen 不是给出一个固定规则，而是把硬件规格纳入策略搜索。CPU 内存更大、SSD 更快、GPU 显存更多时，搜索结果都可能变化。
3.3 CPU compute：有时少搬数据比用 GPU 算更重要
论文还提到一个反直觉的点：虽然 CPU 计算能力远弱于 GPU，但某些 decoding 阶段的 attention score 可以交给 CPU 算。
原因是，如果 KV cache 已经放在 CPU 上，把它搬到 GPU 可能比直接在 CPU 上算更亏。尤其序列很长时，KV cache 体积远大于当前 token 的 activation。此时让 CPU 就地处理一部分 attention，虽然算得慢，但可以避免大规模数据搬运。
这说明 FlexGen 优化的不是“每个算子都最快”，而是整个 memory-compute pipeline 的总时间。
3.4 4-bit 压缩权重和 KV cache
FlexGen 还用了近似方法来进一步降低内存和 I/O：
- 对 weights 做 4-bit group-wise quantization；
- 对 KV cache 也做 4-bit group-wise quantization；
- 可选地使用 sparse attention，只加载 top 10% 的 value cache。
这里的量化目标主要不是用 INT4 矩阵乘加速计算，而是减少存储和搬运成本。论文采用的是细粒度 group-wise asymmetric quantization，存储时保持低比特，计算前再反量化回 FP16。
一个重要结果是：在 OPT-30B 和 OPT-175B 上，4-bit 权重 + 4-bit KV cache 几乎没有带来明显精度损失。论文在 Lambada 和 WikiText 上的结果大致如下：
| 模型 | 配置 | Lambada acc | WikiText ppl |
|---|---|---|---|
| OPT-30B | FP16 | 0.725 | 12.72 |
| OPT-30B | 4-bit | 0.724 | 12.90 |
| OPT-175B | FP16 | 0.758 | 10.82 |
| OPT-175B | 4-bit | 0.756 | 10.94 |
不过论文也提到，3-bit 压缩不能保持精度；而且细粒度量化/反量化本身有开销，开启量化后 FlexGen 会关闭 CPU delegation，避免 CPU 上的量化开销抵消收益。
4. 效果如何？
论文实验主要在 Google Cloud 的 NVIDIA T4 16GB GPU 上完成，CPU 内存 208GB，SSD 1.5TB。测试模型是 OPT-6.7B、OPT-30B 和 OPT-175B，生成 32 个 token，prompt 长度主要测试 512 和 1024。
最核心的吞吐结果如下，单位是 generated token/s：
| 系统 | OPT-6.7B seq 512 | OPT-30B seq 512 | OPT-175B seq 512 | OPT-6.7B seq 1024 | OPT-30B seq 1024 | OPT-175B seq 1024 |
|---|---|---|---|---|---|---|
| Accelerate | 25.12 | 0.62 | 0.01 | 13.01 | 0.31 | 0.01 |
| DeepSpeed | 9.28 | 0.60 | 0.01 | 4.59 | 0.29 | OOM |
| Petals | 8.25 | 2.84 | 0.08 | 6.56 | 1.51 | 0.06 |
| FlexGen | 25.26 | 7.32 | 0.69 | 13.72 | 3.50 | 0.35 |
| FlexGen (4-bit) | 29.12 | 8.70 | 1.12 | 13.18 | 3.98 | 0.42 |
这个表里最显眼的是 OPT-175B：
- 在 seq length 512 下，Accelerate 和 DeepSpeed 只有 0.01 token/s；
- Petals 的 per-GPU throughput 是 0.08 token/s；
- FlexGen 不压缩达到 0.69 token/s；
- FlexGen 开启 4-bit 压缩后达到 1.12 token/s。
也就是说，在一张 16GB T4 上，FlexGen 第一次把 OPT-175B 的生成吞吐推到了约 1 token/s 的量级。这个数字对交互式聊天当然很慢，但对“放一批任务在后台慢慢跑”的场景已经有意义。
论文还做了 ablation，说明各个组件并不是装饰：
| 配置 | OPT-30B | OPT-175B |
|---|---|---|
| All optimizations | 7.32 | 0.69 |
| No policy search | 7.26 | 0.27 |
| No overlapping | 5.86 | 0.59 |
| No CPU compute | 4.03 | 0.62 |
| No disk | 7.32 | OOM |
| w/ DeepSpeed policy | 1.57 | 0.01 |
可以看到，对 OPT-175B 来说，policy search 非常关键；如果把 DeepSpeed/Accelerate 风格的策略移植到 FlexGen runtime 里，吞吐会掉到 0.01 token/s。对 OPT-30B 来说，CPU compute 和 overlap 的收益更明显。
5. 我对这篇论文的理解
FlexGen 的核心价值不在于某一个单点技巧，而在于它把 LLM inference 的问题重新表述成了一个系统调度问题。
很多时候我们会自然地把“模型推理优化”理解成：
- kernel 更快；
- 量化更低比特；
- batch 更大；
- 多卡并行更高效。
FlexGen 关注的是另一个层面：当模型已经大到单卡显存放不下时，真正限制吞吐的往往不是某个 matmul 有多快，而是数据在 GPU、CPU、磁盘之间如何流动。只要 I/O 策略错了，GPU 算力再强也会被饿住。
它的思路可以概括成三句话：
- 为了吞吐，可以牺牲延迟。
- 为了减少 I/O，可以改变计算顺序。
- 为了扩大 batch，可以把 KV cache 和权重一起纳入 offloading/压缩策略。
这也是它和普通“显存不够就 offload”方案的区别。FlexGen 不是被动地把放不下的东西丢到 CPU/SSD，而是主动搜索：哪些 tensor 应该留在 GPU，哪些应该放到 CPU，什么时候读，什么时候写，batch/block 该多大。
6. 局限与阅读时需要注意的地方
首先，FlexGen 明确面向 latency-insensitive 的批处理场景。它可以为了高吞吐接受很长的单批延迟。论文里 OPT-175B 的高吞吐配置甚至可以到几千秒级 latency，所以不能把它理解成面向实时聊天服务的系统。
其次，实验里的 throughput benchmark 使用 dummy model weights，真实权重主要用于精度评估。这对系统吞吐测试是常见做法，因为计算图和 tensor shape 一样，但如果关心具体模型、具体任务、具体采样策略下的端到端表现，还需要额外实验。
第三，FlexGen 的效果依赖硬件层级。CPU 内存容量、PCIe 带宽、SSD 读写速度都会影响最优策略。论文的 T4 + 208GB RAM + NVMe SSD 是一个代表性资源受限环境，但不是所有机器都会得到同样的 policy。
最后，这篇论文的时代背景是 2023 年。它回答的是“在资源受限条件下，如何通过 offloading 做大模型高吞吐生成”。如果目标是今天的在线 serving，还需要和后来的 KV cache 管理、paged attention、continuous batching、speculative decoding 等方向一起看。
7. 总结
FlexGen 给我的启发是：当硬件资源不足时，系统设计不一定只能追求“尽量模拟多卡显存充足的执行方式”。如果任务允许，我们可以换一个目标函数：不再要求每个请求尽快完成，而是让一大批请求总体完成得更快。
在这个目标下，很多选择都会反过来：
- batch 越大，单批延迟越高，但 I/O 越容易摊薄；
- 权重不一定每次用完就释放，复用权重可能比尽快完成某个 batch 更重要；
- CPU 虽然慢，但如果能避免搬运巨大的 KV cache，反而可能更划算；
- KV cache 不只是运行时缓存，而是高吞吐推理里的主要内存对象之一。
所以我会把 FlexGen 看成一篇“LLM inference memory hierarchy 论文”。它最重要的贡献不是证明单张 T4 可以勉强跑 OPT-175B，而是系统地说明了：在 LLM 推理里，吞吐、延迟、batch size、KV cache、权重 I/O 和存储层级之间是一组必须共同优化的变量。