苹果公司在内存容量有限的设备上实现高效的大规模语言模型推理

大型语言模型 29/01/2024

三个要点
✔️ 提出一种方法，用于在大小超过板载内存（DRAM）的大型语言模型上执行推理
✔️ 提出窗口和行列捆绑建议。行列捆绑建议
✔️ 当一个大型语言模型只有一半的模型参数放在 DRAM 中时，提出的方法在 CPU 上的速度比天真的方法快 4-5 倍，在 GPU 上的速度比天真的方法快 20-25 倍

LLM in a flash: Efficient Large Language Model Inference with Limited Memory
written by Keivan Alizadeh, Iman Mirzadeh, Dmitry Belenko, Karen Khatamifard, Minsik Cho, Carlo C Del Mundo, Mohammad Rastegari, Mehrdad Farajtabar
(Submitted on 12 Dec 2023 (v1), last revised 4 Jan 2024 (this version, v2))
Comments: Published on arxiv.
Subjects: Computation and Language (cs.CL); Artificial Intelligence (cs.AI); Machine Learning (cs.LG)

code：

本文所使用的图片要么来自论文、介绍性幻灯片，要么是参考这些图片制作的。

介绍

近年来，以 ChatGPT 为代表的大规模语言模型（LLM）引起了公众的关注。

随着模型参数数量的增加，LLM 的性能也随之提高，目前正在开发更大的 LLM，以进一步提高大型语言模型的性能。

随着 LLM 越来越大，有必要存储大量的模型参数值，并根据这些模型参数值进行大规模计算。LLM 的运行离不开具有大容量内存的个人电脑，它是存储和计算这些模型参数值的 "车间"。

对许多人来说，不可能在自己的电脑上运行最先进的大型 LLM，而只能在云中作为网络服务运行，如 ChatGPT。而在智能手机等小型移动设备上直接运行 LLM 则更加困难。

与云服务不同，在智能手机等移动设备（即所谓的边缘设备）上运行的人工智能被称为边缘人工智能，由于无需网络即可处理，因此有望提高处理速度，降低个人信息和机密信息泄露的风险。

不过，如前所述，目前大型 LLM 的发展趋势已势不可挡，在边缘设备上安装高性能 LLM 并使其正常工作变得越来越困难。

在本文所述的论文中，以 iPhone 而闻名的苹果公司的研究人员提出了一种技术，可使 LLM 在没有大内存和内存有限的设备上高效运行。*这并不是一种只能在设备上进行训练的技术，而是一种可以根据学习到的模型参数值进行高效推理的技术。

谷歌宣布将于 2023 年 12 月 7 日在 Pixel 8 Pro 中实现 LLM（双子座）。作为智能手机的附加值，极有可能受到各家公司的关注。谷歌似乎是在智能手机中实现 LLM，一开始就减少了模型参数的数量，使其能够适合设备内存，而苹果论文中的技术则是针对苹果论文中的技术是一种旨在高效执行无法放入设备内存的 LLM 的方法。基于该技术的配备 LLM 的 iPhone 可能会在未来上市销售。

让我们来看看它的工作原理和运作方式。

原始方法与拟议方法的推理速度比较

首先，展示了拟议方法的有效性。

图 1 显示了天真方法和建议方法的推理速度对比。图 1 比较了从闪存加载的天真 LLM 模型和建议的方法。

https://aisholar.s3.ap-northeast-1.amazonaws.com/posts/January2024/%E5%9B%B31.png — 图 1：推理速度比较。

在这里，内存和闪存这两个术语可能会引起混淆，但它们指的是不同的概念。

内存通常被称为 PC 内存。从隐喻上讲，这一概念相当于 "工作场所"，在本文中指的是一种名为 DRAM 的存储设备。

在个人电脑上，如果要对数据进行任何处理，都必须将其存储在内存中，以便与工作场所相对应，但如果关闭个人电脑，数据也会丢失。

因此，需要保留的数据必须存储在固态硬盘或 SD 卡等存储设备上。从比喻上讲，这个概念相当于一个 "仓库"。

手机拍摄的图像和其他数据都存储在这个存储器中，通常用于存储的固态硬盘和 SD 卡的存储核心称为闪存。基本上，当设备启动时，学习到的模型参数值也会存储在存储器中。

因此，本文中的内存对应 DRAM，闪存对应闪存。

回到推理速度的比较，本文比较了 LLM 在只有 LLM 模型大小一半的内存可用这一限制条件下的推理处理速度。换句话说，我们考虑的是无法将存储在闪存中的整个 LLM 模型部署到内存中的情况。

基线（图中的 Naive）似乎考虑了这样一种情况：LLM 模型大小的一半已经存在 DRAM 中，另一半从闪存中读取并计算。虽然基线本身的设置允许整个 LLM 模型位于 DRAM 中，但一个标记（文本在 LLM 中处理时的处理单元）却位于 DRAM 中。它的设置似乎是这样的：为了推断一个标记，需要从闪存中加载一半的 LLM 模型（把它想象成把文本分成单词，然后逐词处理）。

如图所示，这次的 LLM 模型是 Falcon 7B 和 Opt 6.7B，这两个模型有大约 70 亿个模型参数，被认为具有模型参数稀疏性，下文将对此进行讨论。

图中，"计算 "是指计算时间，"从闪存加载 "是指将模型参数从闪存加载到 DRAM 所需的时间，"内存管理 "是指内存管理所需的时间。对于这两个模型，从闪存加载模型参数的时间是推理的瓶颈，这一点可以从基线推理时间的百分比看出。

本文提出的方法通过从闪存中高效加载模型参数，实现了两种模型的提速。特别是在图 6.7B 中的 GPU 上进行推理时，传统方法推断（生成）一个标记需要 2 秒以上，而本文提出的方法可以将这一过程加速到 0.1 秒以内。换句话说，速度快了 20 多倍。

闪存与 DRAM 存储容量、传输速率和 LLM 模型大小之间的关系

我们刚才已经解释过，实际执行工作的是内存（DRAM），而闪存是保存数据的仓库，这就造成了往返时间的损失。那么，DRAM 和闪存的特性究竟有何不同呢？关键区别在于数据存储容量和传输速度。图 2 显示了 DRAM 和闪存（Flash Memory）的传输速率（带宽）比较。

闪存的存储容量约为 100 GB，而 DRAM 的存储容量约为 10 GB。因此，一般来说，闪存的存储容量大于 CPU 和 GPU 的内部存储器（DRAM）。闪存可存储的模型参数是 DRAM 的 10 倍。

既然如此，为什么不直接使用闪存呢？不过，闪存的速度比 DRAM 慢。闪存的带宽（数据传输速率）约为 1 GB/秒，而 DRAM 约为 100 GB/秒。换句话说，闪存的存储容量是 DRAM 的 10 倍，但传输速率却是 DRAM 的 1/100。与 CPU 和 DRAM 之间的往返计算相比，CPU 和闪存之间的往返计算对推理时间的影响相当大。

如果 DRAM 的存储容量大于 LLM 的模型大小，则只需将闪存中存储的全部训练模型参数加载到 DRAM 中一次即可进行推理。

虽然从闪存中加载整个模型需要时间，但后续推理可在已加载到 DRAM 中的模型上进行计算。此外，由于 LLM 是计算密集型的，通常会在 GPU 上进行处理，因为 GPU 可以比 CPU 更大规模地执行并行计算，但必要的模型参数会被传输到 GPU 的内存中进行计算。

相反，如果没有存储容量大于 LLM 模型大小的 DRAM，该怎么办呢？这就是本文要解决的问题。

拟议方法的机制

概述

本文提出的 "瞬间 LLM "方法有两大创新点。

第 1 点是减少数据传输量。具体来说，整个模型参数不会一次性从闪存加载到 DRAM。为了只传输当前令牌推理真正需要的模型参数，我们建议利用稀疏性来预测少量非零模型参数，并只传输这些参数；利用窗口分割将输入序列分割到前馈网络，并只传输窗口滑动前后必要模型参数的差值。开窗法的目的是将输入序列分割成前馈网络，以便只按顺序传输开窗前后所需的模型参数差值。

要点 2 是优化从闪存读取数据的块大小（一次传输的数据大小）。闪存的块大小越大，读取（读取）吞吐量（单位时间内的处理量）往往越高。为了提高吞吐量并消除从闪存读取模型参数时的瓶颈，建议采用行列捆绑的方式来改善块大小。

要点 1：减少数据传输量（利用稀疏性和窗口处理）

LLM 以变压器为模型，但变压器的每一层都有一个注意层和一个前馈网络。在本文中，注意力层的模型参数始终保存在内存中。这占总模型参数的 1/3。为了减少数据传输，我们将重点放在前馈网络的模型参数（占模型参数的其余 2/3）上。

所研究的 LLM 模型 OPT 和 Falcon 的前馈网络模型参数非常稀疏（大部分模型参数值为 0）。在前馈网络的模型参数值中，OPT 的 97% 和 Falcon 的 95% 为 0。0 模型参数值对计算没有贡献，因此真正需要的模型参数只有其中的非零模型参数。因此，我们建议采用窗口法，根据需要动态读取这些非零模型参数。如果真的要检查是否稀疏，最终还是要将其读入 DRAM，因此他们也在建立一个预测模型来预测是否稀疏。这似乎是一个已有的想法，但它显示的特点是，如果已知当前层的注意层的输出，就可以预测该层以外的前馈网络的输出是否为零。注意层的模型参数始终保存在内存中这一设置似乎也可以用来进行预测。

可能有人会担心预测是否真的能提前发现模型参数值是否为零，但本文在三个不同的 Zero shot 任务（没有 LLM 微调的评估任务）上进行了精度评估，结果表明，有预测和无预测的 OPT 6.7B 模型的精度几乎相同。结果表明，有预测和无预测的 OPT 6.7B 模型的准确度几乎相同。

图 3 是 WINDOWING 的概念图。

LLM 处理用户输入的单词序列，如图所示。在处理这样的单词序列（输入序列）时，本文设置了一个滑动窗口，并读取窗口中与单词相关的模型参数进行推理。只有预测为非零的模型参数值才会被读入。

现在，我们将注意力进一步转向主动神经元。在本文中，活跃神经元的定义是：对于一个输入标记，前馈网络各层的输出均为正值的神经元。可以说，只有活跃神经元才需要被读入，因为非活跃神经元的输出加起来为零，实际上与计算无关。

窗口化的关键在于，下一个窗口中单词输入的大部分活跃神经元与上一个窗口中的活跃神经元相同。在图 3 中，蓝色是需要加载（模型参数）的神经元，红色是不需要加载（模型参数）的神经元。略深的蓝色为新神经元，与前一个窗口不同，新神经元是需要读取的新神经元。因此，在窗口处理过程中，涉及最后五个标记的活跃神经元都会被存储起来。这就减少了下一个窗口中需要新加载的神经元数量，从而减少了数据传输。

图 4 显示了猎鹰 7B 根据窗口大小所需的神经元（模型参数）传输量。

x 轴表示窗口大小，y 轴表示加载到 DRAM 中的模型参数百分比：在 "聚合使用"（记住前一个窗口，不只加载差值）情况下，窗口大小越大，参与计算的模型参数越多。在增量传输（在窗口中读入新的神经元）情况下，窗口大小越大，与前一个窗口共享的模型参数越多，表明只需读入少量新的差异。

要点 2：优化块大小的行列捆绑

块大小与闪存读取吞吐量之间的关系如图 5 所示。

横轴为数据块大小，纵轴为读取吞吐量，两条线之间的差值为线程数。研究表明，从闪存读取数据时，块大小和线程数越多，从闪存读取数据的速度就越快。为了利用闪存的这一特性，我们采用了行列捆绑技术来增加块大小。

在 OPT 和 Falcon 模型中，计算第 i 个中间神经元似乎需要向上投影的第 i 行和向下投影的第 i 列。看来应该考虑前馈网络中存在的映射矩阵（模型参数）的值。这些相应的矩阵数据（模型参数）可以放在一起读取（行列捆绑），以改善块的大小。

行列捆绑的概念图如图 6 所示。

虽然本文对图 6 没有太多直接解释，但可以认为预测器的输出与前面描述的绕组有关。在图中，预测器输出中的紫色、黑色、蓝色和红色神经元被认为是所需的神经元。换句话说，这表明只读取了四个必要的神经元，而图 6 中原本需要读取八行神经元。右侧显示了所需神经元的向上投影行（Up Proj Columns）和向下投影列（Down Proj Rows），右侧显示了如何将它们从闪存中编译到一起。如果向上投影行和向下投影列的大小各为 d_models，那么将它们放在一起就可以将它们的大小增加到 2d_models，即两倍的块大小。由此可见，建议的方法既减少了从闪存读取的数据量，又尽可能地增加了块的大小，从而提高了传输速度，实现了高效推理。