摘要本文解读ACL 2024论文《LLM in a flash: Efficient Large Language Model Inference with Limited Memory》。该论文提出闪存驻留的大模型推理框架通过融合激活稀疏预测、滑动窗口复用与行列捆绑三大技术让设备仅用一半 DRAM 就能运行两倍大小的模型。其特别之处在于以硬件成本模型为指导把每 token 延迟拆解为 I/O 加载、内存管理与计算针对性优化前两者。实验表明I/O 延迟从 2196ms 降到 87ms约 25 倍端到端推理在 CPU 上提速 4-5 倍、GPU 上提速 20-25 倍并在 5 种模型、3 种后端上验证为端侧大模型部署提供了重要借鉴。视频讲解点击观看 B 站视频摘要论文基本信息背景与动机研究主线从问题到结论基准/方法设计分类全景方法细节低秩预测器预测哪个神经元会被激活滑动窗口Windowing增量加载行列捆绑Row-Column Bundling块大小翻倍DRAM 管理预分配 交换删除实验设计与结果结果对比总结关键发现局限性常见问题FAQLLM in a flash 的核心思路是什么为什么激活稀疏性对闪存推理如此关键滑动窗口windowing解决了什么问题行列捆绑为什么能提速这个方法需要重新训练模型吗与 DejaVu、FlexGen 的区别是什么参考链接论文基本信息项目内容标题英文LLM in a flash: Efficient Large Language Model Inference with Limited Memory标题中文闪存驻留的大模型高效推理模型放得下算得动作者Keivan Alizadeh, Iman Mirzadeh, Dmitry Belenko, S. Karen Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, Mehrdad Farajtabar机构Apple会议ACL 2024arXivhttps://arxiv.org/abs/2312.11514项目网站https://machinelearning.apple.com/research/llm-in-a-flash背景与动机大模型推理对内存的要求极高一个 70 亿参数的模型仅加载半精度权重就需要超过14GBDRAM超出了绝大多数个人设备手机、笔记本的内存上限。传统方案要么把整个模型塞进 DRAM直接限制模型规模要么依赖量化压缩但无法根本解决模型装不下的问题。而闪存SSD的容量通常比 DRAM 大一个数量级——问题只在于它慢。本文的出发点是一条硬件事实闪存顺序大块读的吞吐远高于小块随机读M1 Max 上 1GiB 线性读可达 6 GiB/s 以上而小块随机读受首字节延迟支配且 NAND 闪存带宽比 DRAM 低一个量级。基于此作者建立了推理成本模型每 token 延迟 I/O 加载 DRAM 管理 计算前两项正是本文优化对象。在既有工作中DejaVuICML 2023利用激活稀疏性只计算必要神经元但要求权重常驻 GPU 内存FlexGenICML 2023把权重与 KV 缓存卸载到闪存但在模型完全放不进 DRAM的场景下受限于闪存低吞吐。本文首次面向模型完全无法驻留内存的端点设备用硬件成本模型指导闪存数据搬运的算法设计——这是与所有既有路线的本质差异。图 1仅有一半内存可用时单 token 平均推理延迟按需取权把反复加载 flash 的代价与计算重叠延迟大幅下降图 2硬件特性(a) flash 容量远大于 DRAM 但带宽低一个量级(b) 随机读吞吐随顺序块大小与线程数增大研究主线从问题到结论图 9研究主线DRAM 装不下 → flash 动机 → 成本模型 → windowing / 捆绑 → 实验 → 2 倍模型 20 倍加速Mermaid 流程图基准/方法设计方法的核心是三条互补策略全部由硬件成本模型驱动减少数据量ReLU 激活让 FFN 层超过 90% 的神经元输出为零。OPT 6.7B 的 FFN 层稀疏率高达97%ReLUfication 后的 Falcon 7B 为95%FATReLU 化的 Llama 2 为90%。因此每个 token 真正需要的权重只是一小部分。增大块大小闪存顺序读的吞吐远高于随机读把读取块变大能显著摊薄首字节延迟。优化 DRAM 管理预分配矩阵 指针结构把换入换出代价从重排整个矩阵降为 $O(c \times d_{model})$ 的交换改写。选择性持久策略embeddings 与注意力矩阵常驻 DRAM约占模型 1/3FFN 部分按需动态加载兼顾计算效率与内存占用。分类全景图 10高效推理三条路线与本文定位模型压缩、选择性执行、权重卸载本文为 flash 按需取权Mermaid 分类图方法细节低秩预测器预测哪个神经元会被激活既然只有少数神经元被激活就需要在加载权重之前知道哪些神经元会被激活。本文为每一层训练一个低秩预测器输入只需当前层注意力模块的输出比 DejaVu 用上一层 FFN 输出更晚、更准输出该层将被 ReLU 激活的神经元集合。训练数据为 C4 的 1 万样本、2 个 epoch每层在 A100 上约 4 小时。零样本评测显示预测器几乎不损失精度Arc Easy 66.1→66.2Arc Challenge 30.6→30.6HellaSwag 50.3→49.8。图 3滑动窗口保留过去 $k5$ 个 token 的活跃神经元处理新 token 只需增量加载少量新权重滑动窗口Windowing增量加载激活神经元的集合在相邻 token 间高度重叠。论文维护一个滑动窗口$k5$把窗口内 token 激活过的神经元权重缓存在 DRAM处理新 token 时只增量加载新出现的神经元。对 OPT 6.7B每 token 平均只需加载2.4%的 FFN 权重无窗口时约 10%窗口整体占用 DRAM 为模型大小的 52.1%预测器 1.25% embeddings 3% 注意力 32.3% FFN 15.5%。行列捆绑Row-Column Bundling块大小翻倍FFN 中第 $i$ 个中间神经元的激活同时对应 up-projection 的第 $i$ 列与 down-projection 的第 $i$ 行。把这两部分在闪存中同址连续存储读取块从 $d_{model}$ 翻倍为 $2d_{model}$32 位模型约 32 KiB首字节延迟只付一次实测吞吐从 1.25 GB/s 提升到 2.25 GB/s。图 4行列捆绑把 up-projection 第 $i$ 列与 down-projection 第 $i$ 行打包存储单次读入两倍大小的连续块DRAM 管理预分配 交换删除数据进入 DRAM 后如果每次重排矩阵会引入约 25% FFN 常驻数据的改写开销。本文预分配每层最大需求矩阵$\texttt{Req}i \times 2d{model}$用 pointer 记录原始神经元索引、num_used 追踪已用行数删除时与末行交换$O(c \times d_{model})$新增时直接追加到尾部。推理时矩阵前半作 up-projection后半转置作 down-projection——中间神经元顺序不影响输出。图 5DRAM 管理先与末行交换再追加新权重保持连续占用、减少不必要的数据移动图 6低秩预测器(a) 预测器判定为正的元素与 up-projection 真实激活高度一致(b) 小型低秩网络找出将被激活的中间神经元实验设计与结果评测协议单序列推理C4 验证集子集128 个 prompt token 生成 256 个 tokenDRAM 预算约为模型一半Phi-2 为 65%。硬件覆盖 Apple M1 Max1TB SSD、M2 Ultra2TB SSD与 NVIDIA RTX 409024GB读取用 32 线程并行并绕过操作系统缓存。基线为朴素按需加载与半驻留的 hybrid 方案。主结果OPT 6.7B 16bitM1 Max——四步技术叠加的 I/O 延迟分解方法组合DRAM(GB)Flash→DRAM吞吐(GB/s)I/O 延迟(ms)朴素加载013.4 GB6.102196Hybrid6.76.7 GB6.101090预测器4.80.9 GB1.25738窗口6.50.2 GB1.25164捆绑全部6.50.2 GB2.2587每 token 搬运量从 13.4GB 降到 0.2GB67 倍I/O 延迟从 2196ms 压到 87ms约25 倍。稀疏读的吞吐本身更低1.25 vs 6.1 GB/s但行列捆绑把块大小翻倍让吞吐回升到 2.25 GB/s。跨模型端到端延迟ms模型后端NaiveHybridAll本文OPT 6.7BCPU3182--669OPT 6.7BGPU2218--84Falcon 7BCPU30951947706Persimmon 8BCPU380624951041Phi-2 2.7BCPU1287711546Llama 2 7BCPU30951903994全部模型在约一半内存预算下获得 3-5 倍CPU与 26 倍GPU端到端加速OPT 在 GPU 上叠加投机解码后进一步降到60ms。消融与鲁棒性生成 1000 个 token 时闪存加载延迟保持平稳无热节流衰退Nucleus 采样$p0.9$不会让长生成退化量化不改变激活稀疏模式OPT 量化前后平均活跃率 3.30% vs 3.27%因此 4-bit 加载下 7B 模型的 DRAM 需求可从 3.5GB 降到2GB 以下——直接指向手机端部署。图 7OPT 6.7B 权重加载延迟随生成长度变化1000 tokens 内保持平稳无热节流衰退图 8内存-延迟权衡OPT 6.7BGPU窗口越大、DRAM 占用越高每 token 需加载的权重越少、延迟越低结果对比总结图 11结果对比朴素 2196ms → 预测器 738ms → 窗口 164ms → 捆绑 87ms → 端到端 GPU 84msMermaid 流程图关键发现数据量维度预测器 滑动窗口把每 token 搬运量从 13.4GB 降到 0.2GB67 倍缩减。吞吐维度行列捆绑让稀疏读吞吐从 1.25 GB/s 提升到 2.25 GB/sI/O 延迟 164ms→87ms。端到端OPT 6.7B 在 GPU 上从 2218ms 降到 84ms26 倍CPU 上 3182ms→669ms4.8 倍。跨模型泛化Falcon 7B / Persimmon 8B / Phi-2 / Llama 2 均获得 3-5 倍 CPU 加速无需改架构。精度无损低秩预测器对零样本指标几乎无影响HellaSwag 仅 50.3→49.8。与投机解码正交OPT GPU 端到端 84ms 进一步降到 60ms。局限性单批推理评测聚焦单序列生成prompt 处理与多批batch场景未展开。功耗与热限制未系统量化稀疏模型瞬时功率更低但生成时长更长总能耗反而更高。依赖稀疏化网络方法建立在 ReLU / FATReLU 稀疏模型之上非稀疏模型需要先重训练。内存假设DRAM 预算约为模型一半更大 / 更小预算下的延迟-精度权衡有待探索。作者展望了非稀疏网络的选择性加载、按 prompt 动态检索权重、以及与投机解码、多批推理的组合——这些方向后续被 PowerInfer 等系统工作直接继承。常见问题FAQLLM in a flash 的核心思路是什么把模型参数存在闪存SSD中推理时利用激活稀疏性只把当前 token 需要的权重按需加载到 DRAM从而在 DRAM 只有模型一半的设备上运行 2 倍大小的模型。为什么激活稀疏性对闪存推理如此关键ReLU 系激活让 FFN 层 90% 以上的神经元输出为零OPT 6.7B 达 97%每个 token 真正需要的权重只是极小一部分稀疏预测让按需取权成为可能。滑动窗口windowing解决了什么问题相邻 token 激活的神经元高度重叠。窗口保留最近 $k5$ 个 token 激活过的权重新 token 只增量加载新出现的神经元每 token 搬运量从 10% FFN 降到 2.4%。行列捆绑为什么能提速FFN 第 $i$ 个神经元对应 up-projection 第 $i$ 列与 down-projection 第 $i$ 行把它们在闪存中相邻存储后读取块从 $d_{model}$ 翻倍为 $2d_{model}$摊薄了闪存随机读的首字节延迟吞吐从 1.25 提升到 2.25 GB/s。这个方法需要重新训练模型吗预测器需要训练每层约 4 小时 A100C4 数据 1 万样本但主模型只需做一次 ReLU 稀疏化微调Persimmon 8B 原生使用平方 ReLU甚至无需微调。与 DejaVu、FlexGen 的区别是什么DejaVu 要求权重常驻 GPU 内存FlexGen 受限于闪存低吞吐本文面向模型完全放不进 DRAM 的端点设备用硬件成本模型同时优化数据量与读取块大小实现 4-5 倍 CPU、20-25 倍 GPU 加速。参考链接arXiv 论文https://arxiv.org/abs/2312.11514Apple ML Research 项目页https://machinelearning.apple.com/research/llm-in-a-flashDejaVuICML 2023https://arxiv.org/abs/2310.17157FlexGenICML 2023https://arxiv.org/abs/2303.06865PowerInferSOSP 2024https://arxiv.org/abs/2404.07965给大家推荐一款自用写文献综述、无虚构文献的 AI复旦大学 FudanNLP 团队自研 切问学术官网qiewenpaper.com覆盖3.6 亿篇可溯源真实中英文文献能自动整合文献观点生成规范综述还能挖掘研究创新点、复现实验配合视频教学新手快速上手文献综述写作后记博客的关键词集中在编程、算法、机器人、人工智能、数学等等持续高质量输出中。讨论QQ群白拾的小屋 (750365700)⭐B站账号白拾的物理AI组会活跃于知识区和动画区✨GitHub主页YhbCode000工程文件
)
