苹果这项新工作将为未来 iPhone 加入大模型的能力带来无限想象力。
近年来，GPT-3、OPT 和 PaLM 等大型语言模型（LLM）在广泛的 NLP 任务中表现出了强大的性能。不过，这些能力伴随着大量计算和内存推理需求，毕竟大型语言模型可能包含数千亿甚至万亿参数，使得高效加载和运行变得有挑战性，尤其是在资源有限的设备上。
当前标准的应对方案是将整个模型加载到 DRAM 中进行推理，然而这种做法严重限制了可以运行的最大模型尺寸。举个例子，70 亿参数的模型需求 14GB 以上的内存才能加载半精度浮点格式的参数，这超出了大多数边缘设备的能力。
为了解决这种局限性，苹果的研究者提出在闪存中保存模型参数，至少比 DRAM 大了一个数量级。接着在推理中，他们直接并巧妙地从闪存加载所需参数，不再需求将整个模型拟合到 DRAM 中。
这种步骤基于最近的工作构建，这些工作表明 LLM 在前馈网络（FFN）层中表现出高度希罕性，其中 OPT、Falcon 等模型的希罕性更是超过 90%。因此，研究者利用这种希罕性， 有选择地仅从闪存中加载具有非零输入或预测具有非零输出的参数。
论文地址：https://arxiv.org/pdf/2312.11514.pdf
具体来讲，研究者讨论了一种受硬件启发的利润模型，其中包括闪存、DRAM 和计算核心（CPU 或 GPU）。接着引入两种互补技术来最小化数据传输、最大化闪存吞吐量：
窗口：仅为前几个 token 加载参数，并重用最近计算的 tokens 的激活。这种滑动窗口步骤缩小了加载权重的 IO 请求数量；
行列捆绑：保存上投影和下投影层的串联行和列，以读取闪存的更大连续块。这将通过读取更大的块来增加吞吐量。
为了进一步最小化从闪存传输到 DRAM 的权重数量，研究者还设法预测 FFN 希罕性并避免加载归零参数。结合使用窗口和希罕性预测可以为每一个推理查询仅加载 2% 的闪存 FFN 层。他们还提出了静态内存预分配，最大限度缩小了 DRAM 内的传输并缩小了推理耽误。
本文的闪存加载利润模型在加载更好数据与读取更大块之间取得了平衡。与 CPU 和 GPU 中的 naive 实现相比，优化该利润模型并有选择地按需加载参数的闪存策略可以运行两倍于 DRAM 容量的模型，并将推理速度分别提升 4-5 倍和 20-25 倍。
有人评价称，这项工作会让 iOS 开发更加有趣。
闪存和 LLM 推理
带宽和能量限制
虽然现代 NAND 闪存提供了高带宽和低耽误，但仍达不到 DRAM 的性能水准，尤其是在内存受限的系统中。下图 2a 说明了这些差异。
依赖 NAND 闪存的 naive 推理实现可能需求为每一个前向传递重新加载整个模型，这一过程非常耗时，即使是压缩模型也需求几秒工夫。此外将数据从 DRAM 传输到 CPU 或 GPU 内存需求耗费更多能量。
在 DRAM 充足的场景中，加载数据的利润有所降低，这时模型可以驻留在 DRAM 中。不过，模型的初始加载仍然耗能，尤其是在第一个 token 需求快速响应工夫的情况下。本文的步骤利用 LLM 中的激活希罕性，通过有选择地读取模型权重来解决这些挑战，从而缩小了工夫和耗能利润。
读取吞吐量
闪存系统在大量连续读取场景下表现最佳，例如配备 2TB 闪存的 Apple MacBook Pro M2 基准测试表明，对未缓存文件进行 1GiB 线性读取的速度超过 6GiB/s。然而由于这些读取固有的多阶段性质，包括操作系统、驱动程序、中端处理器和闪存控制器，较小的随机读取无法复制这种高带宽。每一个阶段都会出现耽误，从而对较小读取造成较大的影响。
为了规避这些限制，研究者提倡两种主要策略，它们可以同时使用。
第一种策略是读取较大的数据块。虽然吞吐量的增长不是线性的（较大的数据块需求较长的传输工夫），但初始字节的耽误在总请求工夫中所占的比例较小，从而提高了数据读取的效率。图 2b 描述了这一原理。一个与直觉相反但却有趣的观察结果是，在某些情况下，读取比需求更多的数据（但数据块较大）然后丢弃，比只读取需求的部分但数据块较小更快。
第二种策略是利用保存堆栈和闪存控制器固有的并行性来实现并行读取。研究结果表明，在标准硬件上使用多线程 32KiB 或更大的随机读取，可以实现适合希罕 LLM 推理的吞吐量。
最大化吞吐量的关键在于权重的保存方式，因为提高平均块长度的布局可以显著提高带宽。在某些情况下，读取并随后丢弃多余的数据，而不是将数据分割成更小的、效率更低的数据块，可能是有益的。
从闪存加载
受上述挑战的启发，研究者提出了优化数据传输量和提高读取吞吐量的步骤，以显著提高推理速度。本节将讨论在可用计算内存远远小于模型巨细的设备上进行推理所面临的挑战。
分析该挑战，需求在闪存中保存完整的模型权重。研究者评估各种闪存加载策略的主要指标是耽误，耽误分为三个不同部分：从闪存加载的 I/O 利润、管制新加载数据的内存支出以及推理操作的计算利润。
苹果将在内存限制条件下缩小耽误的解决方案分为三个战略领域，每一个领域都针对耽误的特定方面：
1、缩小数据负载：旨在通过加载更少的数据来缩小与闪存 I/O 操作相关的耽误。
2、优化数据块巨细：通过增加加载数据块的巨细来提高闪存吞吐量，从而缩小耽误。
以下是研究者为提高闪存读取效率而增加数据块巨细所采用的策略：
捆绑列和行
基于 Co-activation 的捆绑
3、有效管制加载的数据：简化数据加载到内存后的管制，最大限度地缩小支出。
虽然与访问闪存相比，在 DRAM 中传输数据的效率更高，但会产生不可忽略的利润。在为新神经元引入数据时，由于需求重写 DRAM 中的现有神经元数据，重新分配矩阵和添加新矩阵可能会导致巨大的支出。当 DRAM 中的前馈网络（FFN）有很大一部分（约 25%）需求重写时，这种代价尤其高昂。
为了解决这个问题，研究者采用了另一种内存管制策略。这包括预分配所有必要的内存，并建立相应的数据结构进行有效管制。如图 6 所示，该数据结构由指针、矩阵、偏置、已用数和 last_k_active 等元素组成：
图 6：内存管制，首先将最后一个元素复制到删除神经元，以保持内存块的连续性，然后将所需元素堆栈到最后，这样可以避免多次复制整个数据。
值得注意的是，重点并不在过程的计算方面，因为这与本文工作核心无关。这种划分使研究者能够专注于优化闪存交互和内存管制，从而在内存受限的设备上实现高效推理。
实验结果
OPT 6.7B 模型的结果
预测器。如图 3a 所示，本文的预测器能准确识别大多数激活的神经元，但偶尔也会误识数值接近于零的非激活神经元。值得注意的是，这些接近零值的假负类神经元被排除后，并不会明显改变最终输出结果。此外，如表 1 所示，这样的预测准确度水平并不会对模型在零样本任务中的表现产生不利影响。
耽误分析。当窗口巨细为 5 ，每一个 token 需求访问 2.4% 的前馈网络（FFN）神经元。对于 32 位模型，每次读取的数据块巨细为 2dmodel × 4 字节 = 32 KiB，因为它涉及行和列的连接。在 M1 Max 上，从闪存加载每一个 token 的耽误为 125 毫秒，内存管制（包括神经元的删除和添加）的耽误为 65 毫秒。因此，与内存相关的总耽误不到每一个 token 190 毫秒（参见图 1）。相比之下，基线步骤需求以 6.1GB/s 的速度加载 13.4GB 的数据，导致每一个 token 的耽误约为 2330 毫秒。因此，与基线步骤相比，本文的步骤有了很大改进。
对于 GPU 机器上的 16 位模型，闪存加载工夫缩短至 40.5 毫秒，内存管制工夫为 40 毫秒，由于从 CPU 向 GPU 传输数据的额外支出，工夫略有增加。尽管如此，基线步骤的 I/O 工夫仍然超过 2000 毫秒。
表 2 提供了每种步骤对性能影响的详细比较。
Falcon 7B 模型的结果
耽误分析。在本文的模型中使用巨细为 4 的窗口，每一个 token 需求访问 3.1% 的前馈网络（FFN）神经元。在 32 位模型中，这相当于每次读取的数据块巨细为 35.5 KiB（按 2dmodel ×4 字节计算）。在 M1 Max 设备上，从闪存加载这些数据所需的工夫约为 161 毫秒，内存管制过程又增加了 90 毫秒，因此每一个 token 的总耽误工夫为 250 毫秒。相比之下，基线耽误工夫约为 2330 毫秒，本文的步骤大约快 9 到 10 倍。