我读了读苹果的AI大模型论文，发现这几个秘密

今年WWDC上，苹果高调推出了AI战略，还顺带玩了个谐音梗，把自家的AI定义为“Apple Intelligence”。

随着大模型不断Scaling up，推理所需的计算和存储开销也在快速增长。

然而手机的内存资源很有限，这就导致我们很难在手机上直接运行本地大模型。

举个例子，一个7B参数的模型大约需要14GB的内存来加载模型权重。然而，最新的iPhone 15 Pro Max的运行内存只有8GB，远远不够。

如何在有限的内存下实现高效的大模型推理，是端侧AI发展的重要任务。

今年初，苹果发表了一篇可以说对端侧AI有着奠基意义的重要论文，一定程度上解决了这个问题。

这篇论文的标题是《LLM in a flash: Efficient Large Language Model Inference with Limited Memory》，周末我仔细学习了一下，收获很大，分享给大家。

一、大模型推理的背后：闪存与DRAM

通常来说，大模型在推理时需要将所有参数完整加载到DRAM中。随后DRAM再把模型参数按需传递给GPU和CPU，进行推理运算。

DRAM的容量普遍较小，但是带宽高；相比之下，闪存的容量大，但带宽低。

由于需要在DRAM中加载整个大模型的权重，需要占据较大的DRAM空间。

可是手机的DRAM容量普遍很小，这对移动端本地推理构成了巨大挑战。

为了应对这个棘手的问题，需要用到大模型FFN层的稀疏性。

二、核心思想：运用FFN层的稀疏性进行动态加载

大模型的Transformer层可以分为Attention层和FFN层。

Attention层的权重参数一直保存在DRAM中，约占大模型权重的1/3，这部分不涉及参数的动态加载。

FFN层的权重参数是非常稀疏的，比如OPT 6.7B模型在其FFN层有着97%的稀疏性，Falcon 7B模型有着95%的稀疏性。

这意味着，每次推理时，FFN层只有很小一部分神经元被激活，大部分权重参数都是0，没有参与推理运算。

因此，我们只需要把这部分活跃神经元加载到DRAM中即可，不需要加载整个大模型，这就大大降低了对DRAM容量的要求。

三、Predictor：低秩预测器

FFN层的稀疏性主要是由ReLU激活函数引起的。

ReLU激活函数的特点是输出非负值，当输入小于0时输出为0，这导致在FFN层中有超过90%的中间输出为0，产生了高度的稀疏性。

为了不把整个投影矩阵加载到DRAM中，苹果构造了一个低秩预测器来预测哪些中间神经元会被ReLU激活（即输出非零值）。

这样一来，我们只需要把预测为非零输出的神经元的参数，从闪存中加载到DRAM中即可。

四、Windowing：滑窗动态加载

苹果提出了一种动态加载权重参数的方法，称为Windowing。

相比一次性把整个输入序列都给到FFN层，我们可以每次只输入一个定长序列。

这里需要选择一个window size（比如5），然后每次我们只向FFN层输入5个token的输入序列。

如下图所示，我们先把“Once Upon A Time There”这5个token构成的输入序列给到FFN层。

通过低秩预测器，我们可以推断出这5个token对应的活跃神经元有哪些（绿色），然后将这些神经元数据加载到DRAM中。

随后我们将滑窗向右平移1个token，再通过低秩预测器推断出针对“Upon A Time There Was”这个5 token序列而言，哪些神经元是活跃的。

这5个token对应的活跃神经元和上5个token对应的活跃神经元是高度重合的。因此我们只需要加载小部分新权重（蓝色），并且删除一些旧权重（红色）即可。

这种基于滑窗的边际动态加载机制，有效降低了DRAM的存储要求。

五、Bundling：行列捆绑存储

在此基础之上，苹果还提出了行列捆绑存储策略，进一步提高了有限内存资源下的推理效率。

当我们要计算一个神经元数据时，需要用上投影层权重矩阵的对应列乘以下投影层权重矩阵的对应行。

这就需要把行、列数据都从闪存加载到DRAM中，需要读取两次。

为了简化操作，提高吞吐量，我们可以把行、列数据捆绑存储在闪存中，这样只要读取一次就好了。

想法很朴素，效果很明显。

六、实验结果：推理速度提升20-25倍

为了验证以上方法的有效性，苹果用OPT 6.7 B大模型，在M1 Max上进行了测试。

实验表明，当我们把低秩预测器（Predictor）、滑窗动态加载（Windowing）、行列捆绑存储（Bundling）这三种优化方法都用上后，能够在只有6.5 GB DRAM的手机上运行14.3 GB的大模型。（注：6.7 B的16位大模型，对应参数权重大小约为13.4 GB）

换句话说，通过这些方法的组合，能够让手机本地运行两倍于自身DRAM容量的大模型。

与此同时，该方法能够在CPU上实现比传统加载方法快4-5倍，在GPU上快20-25倍的推理速度，I/O延迟从2310毫秒直降到87毫秒，效果显著。

七、结语：端侧AI的头号玩家

对于曾经开创“iPhone时刻”的苹果而言，人们希望看到它在AI时代能够继续引领浪潮，做到“Only Apple Can Do”。

前段时间，Coatue在报告中详细阐述了端侧AI的巨大投资机会。预计未来90%的推理任务会发生在你手机和PC上，对日常工作生活的任务请求进行实时响应。

苹果是端侧AI的头号玩家，期待它能推出下一款划时代意义的作品。