摘要本文解读 ICML 2023 论文《FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU》。该论文提出FlexGen 高吞吐生成推理引擎通过融合GPU/CPU/磁盘三级内存聚合、线性规划卸载策略搜索与4-bit 权重与 KV cache 压缩让 1750 亿参数的 OPT-175B 模型在一张 16GB 消费级显卡上完成高吞吐生成其特别之处在于把「在哪存、何时算、谁来算」形式化为可证明的优化问题。实验表明 FlexGen 在 OPT-175B 上首次达到1 token/s 生成吞吐有效批次 144最大吞吐较 DeepSpeed ZeRO-Inference 提升100 倍以上为资源受限场景下的大模型推理提供了重要借鉴。视频讲解点击观看 B 站视频摘要论文基本信息为什么单卡跑大模型这么难研究主线从问题到结论基准/方法设计分类全景FlexGen 的三板斧方法细节实验设计与结果结果对比总结关键发现局限性常见问题FAQFlexGen 如何做到单卡运行 175B 模型为什么 FlexGen 吞吐比 DeepSpeed 高这么多4-bit 量化会损失多少精度FlexGen 与 llama.cpp 有什么关系CPU 也能参与计算吗FlexGen 开源吗参考链接论文基本信息项目内容标题英文FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU标题中文FlexGen单 GPU 上的大语言模型高吞吐生成推理作者Ying Sheng, Lianmin Zheng, Binhang Yuan, Zhuohan Li, Max Ryabinin, Daniel Y. Fu, Zhiqiang Xie, Beidi Chen, Clark Barrett, Joseph E. Gonzalez, Percy Liang, Christopher Ré, Ion Stoica, Ce Zhang机构斯坦福大学 · 加州大学伯克利分校 · 苏黎世联邦理工学院会议ICML 2023arXivhttps://arxiv.org/abs/2303.06865项目网站https://github.com/FMInference/FlexGen为什么单卡跑大模型这么难大模型生成推理对算力和内存的要求极高GPT-175B 仅模型权重就需要 325GB 显存至少要 5 张 A10080GB配合复杂并行策略才能运行。而 FlexGen 关注的是另一类更普遍的需求——基准评测、信息抽取、数据整理这类延迟不敏感、需要大批量处理的任务。在这类任务里延迟可以牺牲吞吐才是目标这为降低硬件门槛提供了空间。瓶颈并不只在权重。KV cache 是新的瓶颈当批次为 512、输入序列 512、输出序列 32 时KV cache 峰值达到1.2TB是模型权重的 3.8 倍直接限制了大批次的可行性。现有卸载系统DeepSpeed ZeRO-Inference、Hugging Face Accelerate把训练期的调度策略直接搬来推理逐行遍历计算图、反复加载权重I/O 开销巨大OPT-175B 的批次只能做到 1–2吞吐远低于硬件极限。Petals 为代表的去中心化协作推理则受制于网络带宽无法支撑大批次。FlexGen 的思路是把 GPU、CPU、磁盘当作统一的三级内存池用大批次摊薄跨层级的 I/O 成本并与计算重叠。研究主线从问题到结论图 6FlexGen 研究主线问题 → 吞吐导向动机 → 计算图形式化 → LP 策略搜索 → 4-bit 压缩 → 1 token/s 与 100 倍吞吐Mermaid 流程图基准/方法设计FlexGen 把生成推理建模为计算图遍历问题图中每个方块是一次「某层 × 某 token」的 GPU 批次计算同色方块共享同一层权重合法路径必须满足依赖关系、设备内存约束目标是最小化总执行时间。在此基础上FlexGen 从三个维度构造搜索空间计算调度从行优先改为zig-zag 块调度——按列遍历、块内多 GPU 批次复用权重I/O 复杂度被证明不超过最优方案的2 倍。张量放置用 9 个比例变量刻画权重、激活、KV cache 在 GPU/CPU/磁盘上的分布如 20% 权重在 GPU、80% 在 CPU。计算委派KV cache 在 CPU 时把 attention score 的计算交给 CPU只需搬运激活I/O 量缩小为原来的1/ss 为序列长度。图 1单张 NVIDIA T416GB上 OPT-175B 与 OPT-30B 的延迟—吞吐帕累托前沿FlexGen 最大吞吐较现有系统高 100 倍其他系统受 OOM 限制无法继续提升分类全景FlexGen 的三板斧图 7FlexGen 三板斧zig-zag 卸载调度、代价模型与 LP 策略搜索、4-bit 量化与稀疏注意力外加多卡流水并行扩展Mermaid 分类图方法细节zig-zag 块调度是 FlexGen 的核心按列遍历计算图让权重常驻 GPU 复用只换入换出激活与 KV cache块大小GPU 批次 × 块内批次即有效批次大小。同时用 6 个逻辑线程把「加载下一层权重、存取前后批的 cache/激活、当前计算」重叠起来最后一个同步点收尾。代价模型与策略搜索单块延迟 $T T_{pre}\cdot l T_{gen}\cdot (n-1)\cdot l$$l$ 为层数、$n$ 为生成步数每个阶段取各类 I/O 与计算的最大值再叠加 GPU/CPU/NVMe 三处峰值内存约束。固定 $(bls, gbs)$ 后9 个放置比例构成一个线性规划问题目标 $\min T/bls$硬件参数先 profiling 拟合LP 秒级可解。图 2LLM 推理的计算图方块为某层 × 某 token 的 GPU 批次计算同色方块共享同一层权重图 3行优先调度每步重新加载权重、I/O 巨大zig-zag 块调度按列复用权重仅换入换出激活与 KV cache4-bit 近似方法按 64 个元素一组计算 min/max 后量化到 4-bit 整数、计算前反量化回 FP16权重沿输出通道、KV cache 沿隐藏维度分组零重训零校准稀疏注意力则每个 query 只保留 K cache 的Top-10% 索引只加载对应 V cache 子集。实验设计与结果评测硬件为NVIDIA T416GB Xeon CPU208GB NVMe SSD1.5TB模型为 OPT 6.7B/30B/175B提示长度 512/1024、每提示生成 32 token生成吞吐 生成 token 数 /prefill 解码时间。基线为 DeepSpeed ZeRO-Inference、Hugging Face Accelerate 与按 GPU 数归一化的 Petals。端到端生成吞吐token/s序列长度 512系统6.7B30B175BAccelerate25.120.620.01DeepSpeed9.280.600.01Petals8.252.840.08FlexGen25.267.320.69FlexGen (4-bit)29.128.701.12在 OPT-175B 上FlexGen4-bit达到1.12 token/s较 DeepSpeed 提升 112 倍30B 上 7.32 vs 0.60提升 12 倍。压缩版把所有权重与 KV cache 驻留 CPU有效批次 144彻底避开磁盘交换。同 5000 秒延迟约束下FlexGen 吞吐高40 倍允许更高延迟并开启压缩可达100 倍。图 4不同网络延迟与带宽下 FlexGen单 T4与 Petals 集群4×T4per-GPU 归一化的延迟与吞吐卸载在高吞吐场景优于长流水协作推理多卡流水并行解码吞吐 token/sFlexGen(1)→FlexGen(4) 在 175B 上为 0.83→3.86超线性因为每卡内存压力下降后可切到 CPU 驻留并扩大批次端到端吞吐未线性增长源于 prefill 流水气泡与仅生成 32 token 的设置。精度4-bit 量化Lambada acc / WikiText ppl配置Lambada 30BWikiText 30BLambada 175BWikiText 175BFP160.72512.720.75810.824-bit0.72412.900.75610.944-bit 稀疏0.71812.900.75610.94逐技术消融吞吐 token/s30B 全量 7.32无 CPU 计算 4.03、无重叠 5.86、移植 DeepSpeed 策略仅 1.57175B 全量 0.69无策略搜索 0.27、无磁盘直接 OOM——策略质量决定数量级差异。FlexGen 还完成了 HELM 集成OPT-IML-30B 在单卡上21 小时跑完 7 个代表性子场景。图 5diagonal 块调度灰色为一次性预热区每轮计算一条对角线的多个子对角线各层权重加载一次覆盖整条对角线理论上块大小可扩近 2 倍、平均完成延迟减半需非连续 KV cache 注意力未实现结果对比总结图 8结果对比总结175B 吞吐 0.01→0.69→1.12 token/s、解码 0.83→3.86 超线性、4-bit 精度几乎无损Mermaid 流程图关键发现首次在 16GB 单卡上以 1 token/s 运行 OPT-175B有效批次 144突破此前卸载系统的吞吐上限。4-bit 权重 KV cache 压缩零重训零校准OPT-175B 的 Lambada 准确率仅从 0.758 降到 0.756几乎无损。策略搜索是吞吐的核心OPT-175B 关闭策略搜索后吞吐从 0.69 掉到 0.27移植 DeepSpeed 策略仅 0.01。I/O 重叠与 CPU 委派各贡献显著30B 上去掉重叠 7.32→5.86去掉 CPU 计算 7.32→4.03。流水并行实现解码吞吐超线性175B 解码从单卡 0.83 到四卡 3.86 token/s。HELM 评测民主化OPT-IML-30B 的 7 个子场景在单卡 21 小时完成显著降低基准评测门槛。局限性面向延迟不敏感任务批处理场景适用交互式聊天等低延迟场景不适用。强依赖 CPU 内存容量208GB CPU DRAM 是 175B 卸载的关键CPU 内存小的机器如 3090 配 125GB吞吐反而低于 T4 配置。最优调度未落地I/O 最优的 diagonal 调度需要非连续 KV cache 注意力实际实现的是次优的 zig-zag 块调度。近似方法有边界3-bit 量化无法保持精度启用量化后 CPU 委派因解压开销过大而关闭。常见问题FAQFlexGen 如何做到单卡运行 175B 模型FlexGen 把权重、激活与 KV cache 按比例分布到 GPU、CPU 和磁盘三级内存用 zig-zag 块调度按列复用权重配合 4-bit 压缩让数据主要驻留 CPU从而绕开慢速磁盘在 16GB 显存上跑通 OPT-175B。为什么 FlexGen 吞吐比 DeepSpeed 高这么多DeepSpeed 继承训练期逐行调度每生成一步都要重新加载整层权重FlexGen 按列遍历复用权重并用线性规划搜索最优放置与批次大小把有效批次从 1–2 提升到 144I/O 成本被大批次摊薄。4-bit 量化会损失多少精度几乎无损OPT-175B 的 Lambada 准确率从 FP16 的 0.758 到 4-bit 的 0.756WikiText 困惑度从 10.82 到 10.943-bit 则无法保持精度。FlexGen 与 llama.cpp 有什么关系FlexGen2023是首个形式化卸载搜索的系统其「GPU/CPU 协同、层级卸载、批处理摊薄 I/O」的思想被 llama.cpp、PowerInfer 等本地推理生态广泛沿用。CPU 也能参与计算吗可以。当 KV cache 在 CPU 上时FlexGen 把 attention score 的计算委派给 CPU只搬运激活I/O 量缩小为 1/s但启用量化后 CPU 解压开销过大该委派被关闭。FlexGen 开源吗开源代码在 https://github.com/FMInference/FlexGen 支持 OPT 系列模型可配置 T4/3090 等不同硬件。参考链接arXiv 论文页https://arxiv.org/abs/2303.06865项目主页代码https://github.com/FMInference/FlexGenDeepSpeed-Inference基线https://github.com/microsoft/DeepSpeedPetals 协作推理基线https://github.com/bigscience-workshop/petalsHELM 评测基准https://crfm.stanford.edu/helm/vLLM后续 KV cache 管理代表作https://github.com/vllm-project/vllm给大家推荐一款自用写文献综述、无虚构文献的 AI复旦大学 FudanNLP 团队自研 切问学术官网qiewenpaper.com覆盖3.6 亿篇可溯源真实中英文文献能自动整合文献观点生成规范综述还能挖掘研究创新点、复现实验配合视频教学新手快速上手文献综述写作后记博客的关键词集中在编程、算法、机器人、人工智能、数学等等持续高质量输出中。讨论QQ群白拾的小屋 (750365700)⭐B站账号白拾的物理AI组会活跃于知识区和动画区✨GitHub主页YhbCode000工程文件