LLM端侧部署系列 | PowerInfer-2助力AI手机端侧部署47B大模型 (论文解读)

• 引言

• 简介

• PowerInfer-2 概述

• 神经元感知的运行时推理

  • 多态神经元引擎

  • 内存中的神经元缓存

  • 灵活的神经元加载

  • Neuron-Cluster-Level Pipeline

• 生成执行计划

• 执行

• 总结

0. 引言

一雨池塘水面平，淡磨明镜照檐楹。东风忽起垂杨舞，更作荷心万点声。

小伙伴们好，我是《小窗幽记机器学习》的小编：卖冰棍的小男孩。紧接前文：LLM端侧部署系列 | 手机上运行47B大模型?上交推理框架PowerInfer-2助力AI手机端侧部署，今天这篇小作文从论文的角度，详细解读PowerInfer-2的技术细节。需要强调的是，PowerInfer-2是首个在智能手机上运行高达470亿参数LLM的高速推理框架。此外，如果小伙伴们想跟小编进一步交流，也可以公众号上添加小编微信。

论文地址： 

https://arxiv.org/abs/2406.06282

GitHub地址：

https://github.com/SJTU-IPADS/PowerInfer

1. 简介

为啥大模型要做端侧部署？除了常说的端侧数据安全隐私，其实还有系统上的原因。PowerInfer-2是一个专为智能手机如何高速推理大型语言模型（LLM）而设计的框架，特别适用于模型大小超过设备内存容量的场景。PowerInfer-2的关键思路是将传统的矩阵计算分解为细粒度的神经元集群计算。具体而言，在PowerInfer-2中设计多态神经元引擎，能够根据LLM推理的不同阶段自适应采用不同的计算策略。此外，引入了分段神经元缓存(neuron caching) 和细粒度神经元集群级流水线(fine-grained neuron-cluster-level pipelining)技术，有效地减少I/O操作引起的开销。

从PowerInfer-2的实验结果表明，它能够支持多种LLM模型在两款智能手机(OnePlus 12和Ace 2)上运行，并在速度上比最先进的框架快29.2倍。值得注意的是，PowerInfer-2是第一个能够在智能手机上以每秒11.68个token的生成速度为TurboSparse-Mixtral-47B模型提供服务的系统。对于完全适应内存的模型，PowerInfer-2在保持与llama.cpp和MLC-LLM相当的推理速度的同时，内存使用量减少了约40%。

2. PowerInfer-2概述

传统的低功耗语言模型推理通常依赖于矩阵计算作为推理的基本单元，这种方式在智能手机的异构硬件环境中所需要的计算量和I/O开销都比较大。这种粗粒度的计算不能有效地利用XPUs灵活的计算能力。更糟糕的是，如果矩阵权重的一部分存储在存储设备上，那么在进行矩阵计算之前，必须等待这些权重被加载到内存中，这会带来相当长的I/O等待时间。

PowerInfer-2这一专为智能手机端侧部署而生的高速LLM推理框架，实现了三个目标：

• 低推理延迟：在预填充阶段和解码阶段都尽量减少推理延迟；

• 低内存占用：减少推理过程中的内存使用量，即使当模型大小超过设备内存限制时也能实现低延迟推理；

• 灵活性：确保能够适应具有不同算力、内存和存储能力的智能手机。

2.1 神经元集群和架构

在PowerInfer-2中提出了一种称为神经元集群的计算抽象，专门用于异构计算场景中的LLM推理。PowerInfer-2以神经元集群为粒度执行计算和I/O操作，神经元集群在计算过程中可以动态组合多个被激活的神经元，神经元的数量由计算单元的计算能力决定。例如，在解码阶段，当CPU核执行计算时，分配给每个CPU核的神经元集群的大小小于预填充阶段中处理的神经元集群的大小。通过使用这种抽象，PowerInfer-2可以充分利用具有不同计算能力的XPUs，有效降低I/O开销。

图2示例了PowerInfer-2的整体架构，该架构分为在线部分（右侧）和离线部分（左侧）。在线部分以神经元集群为粒度进行推理，并包括四个协同组件：多态神经元引擎、内存中的神经元缓存、灵活的神经元加载和神经元集群级别的I/O管道。

多态神经元引擎在预填充阶段和解码阶段使用完全不同的计算模式。对于预填充阶段，神经元集群包含来自权重矩阵的所有神经元，并主要依赖NPU来处理大规模矩阵乘法。对于解码阶段，它调用一个预测器来识别在启动计算之前将会被激活的神经元。然后，引擎将这些激活的神经元合并成一个小的神经元集群，并利用CPU核动态计算神经元集群，从而大大减少了运行时的计算量和内存使用量。

在开始推理计算之前，计算引擎从神经元缓存中检索神经元权重，神经元缓存经过优化，以充分利用LLM推理中观察到的神经元级访问的局部性。如果缓存缺失，PowerInfer-2将启动一个I/O命令，从存储中获取未缓存的神经元权重。为了减少I/O延迟，PowerInfer-2引入了一种并发处理神经元集群和I/O操作的新型流水线机制。此外，PowerInfer-2通过自适应地打包和加载神经元来最小化I/O开销，这由模型的量化确定。

为了自动适应不同的模型或智能手机，在在线推理开始之前开始，每个模型都在智能手机上执行一次离线过程。该过程涉及接收三种类型的输入：模型权重、用户输入和硬件规格。它输出一个执行计划，描述了在线推理中涉及的每个组件的配置，并指导在线过程。

具体而言，离线规划器输出计算、内存和I/O的配置。对于计算，规划器基于它们的计算能力确定在不同阶段或层中CPU和NPU的使用比例。在内存配置方面，为了在内存使用和推理性能之间取得平衡，规划器允许用户在运行PowerInfer-2之前设置所需的推理速度。根据这个速度设置，PowerInfer-2计算所需的最佳缓存大小。对于I/O配置，规划器触发一个分析器来测量模型的稀疏性和热点神经元的分布。

3. 神经元感知的运行时推理

3.1 多态神经元引擎

PowerInfer-2引入了一种多态神经元引擎，可以动态地将神经元组合成神经元集群，以充分利用LLM推理阶段和异构XPUs的不同计算特性。

3.1.1 以NPU为中心的预填充

在预填充阶段，同时处理所有提示的token。尽管每个token都显示出很高的稀疏性并激活不同的神经元，但由于这些激活的聚合，整体上的稀疏性显著降低。因此，PowerInfer-2不使用预测器来计算激活的神经元，而是选择直接将所有神经元合并成一个大的神经元集群。鉴于与CPU核相比，NPUs在处理大规模矩阵乘法方面表现出色，PowerInfer-2利用NPUs进行预填充阶段的计算。

尽管CPU核不参与矩阵计算，但PowerInfer-2利用它们在预填充阶段执行NPU的基本准备任务。首先，由于内存有限，PowerInfer-2依赖一个CPU核在预填充阶段将存储在闪存中的权重加载到内存中。其次，由于当前的NPUs不支持使用量化权重进行直接计算，PowerInfer-2使用CPU核在NPU进行计算之前对数据进行反量化。

图3-a演示了CPU和NPU在transformer层粒度上如何协同进行预填充阶段的推理。

NPU计算需要使用有限的内存，并与CPU共享该内存。因此，在NPU计算开始之前，CPU应将所需的矩阵权重预加载到这个共享内存中。在特定的LLM层中，在NPU执行任何矩阵乘法之前，多个CPU mid-core从神经元缓存中读取量化矩阵权重，并提前将这些矩阵权重反量化为fp16，最终将结果存储在CPU和NPU之间的共享内存中。与此同时，PowerInfer-2使用另一个big-core异步预加载下一层的所有矩阵权重到神经元缓存中。mid-core的反量化、NPU的计算和big-core的I/O操作并发进行，都可以减少I/O开销。

值得注意的是，由于预填充阶段涉及稠密矩阵而不是稀疏计算，通过I/O进行权重加载可以利用顺序读取将大块数据加载到内存中，从而最大化UFS的I/O带宽的使用。

3.1.2 以CPU为中心的解码

与预填充阶段不同，解码阶段在每次迭代中集中处理单个token，表现出显著的稀疏性，因为权重矩阵中仅有很小一部分的神经元（约10%）被激活并参与计算。因此，当从预填充阶段过渡到解码阶段时，多态神经元引擎将权重矩阵的计算划分为小的神经元群集，其中的元素通过预测器被标识为激活状态。通过实验观察到，当batch size为1时，CPU核上的矩阵-向量计算的延迟低于NPUs上的延迟。此外，由于稀疏性导致激活的神经元数量减少，CPU核非常适用于这些轻量且稀疏的计算任务。因此，PowerInfer-2 在解码阶段中专门使用CPU核进行神经元群集计算。

具体而言，PowerInfer-2在解码阶段利用CPU核计算注意力和FFN块。尽管注意力块没有表现出稀疏性，但当输入只是单个向量时，CPU核仍提供较低的计算延迟。对于FFN块，PowerInfer-2 首先将FFN块的输入向量传递给预测器，预测需要激活的FFN权重矩阵中的神经元，并将它们合并为一个神经元群集。然后，每个CPU核获取一个群集，并计算群集中的这些神经元和输入向量。

图3-b说明了解码阶段不同CPU核执行的推理过程。CPU核首先从神经元缓存中读取注意力块的权重，并与输入向量进行计算。然后，它们运行预测器来确定后续权重矩阵中神经元的激活状态。在第三步中，CPU核将激活的神经元划分为多个群集，每个核负责计算其群集中的激活神经元和输入向量，并最终在边界处聚合结果。如果这些神经元在神经元缓存中，CPU核将使用输入向量进行计算。在缓存未命中的情况下，即神经元不在神经元缓存中，运行在CPU核上的I/O线程异步加载神经元到缓存中，最终通知计算线程完成计算。

3.2 内存中的神经元缓存

高效的缓存设计可以防止昂贵的存储I/O操作，从而优化端到端的推理性能。缓存系统的有效性取决于推理过程中的局部性。然而，传统的LLM推理需要遍历所有权重以生成每个token，没有局部性，因此任何缓存设计都无效。

参考LLM in a Flash中提出在推理过程利用稀疏激活选择性加载权重的做法，这种做法还将共同激活的神经元捆绑在一起，一起从闪存中加载，以减少I/O操作。然而，该方法忽视了神经元激活的偏斜分布，即少数热门神经元更频繁地激活，并与大多数其他神经元高度连接。这对于设计有效的缓存策略带来了挑战。首先，这些热门神经元位于不同的神经元捆绑中，并且会从闪存中重复加载，浪费了I/O带宽。其次，研究结果显示，去除这些热门神经元会使剩余神经元之间的共同激活概率降低到20%以下，使捆绑机制在减少I/O操作方面失效。

为了解决这个问题，PowerInfer-2引入了适用于LLM中各种数据类型的分段神经元缓存设计。它将缓存分为多个区域，每个区域具有特定的预取和驱逐策略。注意力块的权重较小且稀疏激活较少，因此在运行时预加载并保留。相比之下，FFN块由于热门神经元的频繁激活，采用基于最近最少使用（LRU）的动态驱逐策略。这种方法确保热门神经元更有可能保留在缓存中，而冷门神经元则经常被驱逐并按需从闪存中加载。重要的是，驱逐过程不涉及写入存储，只是从内存中直接丢弃权重。

PowerInfer-2 利用经典的双队列方法实现了其LRU神经元缓存，以单个神经元为粒度管理LLM权重。系统维护两个双向链表队列，标记为活动队列和非活动队列，队列中的神经元顺序由它们最近访问的时间确定，最近访问的神经元位于队列的头部。在运行时，所有神经元最初加入非活动队列。在重新访问时，它们被提升到活动队列的前面。已经在活动队列中的神经元在后续访问时移动到队列头部。为了管理缓存容量，当活动队列填满缓存空间的90%时，从活动队列的尾部将神经元移动到非活动队列，直到活动队列的占用率降至90%以下。如果缓存达到容量上限，非活动队列尾部的神经元将被丢弃，为新条目腾出空间。

3.3 灵活的神经元加载

配备了高效存储活跃神经元的神经元缓存后，推理过程仍然不可避免地需要对未缓存神经元进行I/O操作。为了优化I/O读取吞吐量并最小化I/O操作，PowerInfer-2还将相关神经元进行捆绑。尽管一旦热门神经元被移除，单个FFN权重矩阵内的共同激活变得不频繁，但是不同矩阵中相应位置的神经元经常同时激活。例如，Gate、Up和Down矩阵中的第i个神经元的共同激活概率高达80%。因此，PowerInfer-2选择基于神经元而不是矩阵结构来存储神经元权重，将Gate、Up和Down矩阵中第i个神经元的权重拼接成一个单独的条目。

PowerInfer-2还针对不同模型引入了不同的I/O加载策略，考虑到它们的量化方法和UFS I/O的固有特性。对于没有量化的模型，由于每个神经元占用的存储空间较大，PowerInfer-2使用更大的粒度进行随机读取以提高I/O带宽。例如，Llama-7B-FP16中的单个神经元占用8KB，Gate、Up和Down矩阵中的神经元总大小为24KB。PowerInfer-2通过单个随机I/O读取将整个24KB的激活捆绑有效地传输到内存中。

对于4bit量化的模型，捆绑大小设置为8KB。以Llama-7B模型为例，每个神经元量化为4bit精度，占用2.5KB（2KB用于量化的int4值，0.5KB用于量化组的FP16 scales值）， 组合捆绑大小为7.5KB。为了与存储介质的最小读取粒度4KB对齐，PowerInfer-2额外补充了0.5KB的捆绑，使总大小达到8KB。然而，PowerInfer-2不是通过单个I/O操作加载这些8KB的捆绑，而是选择4KB的粒度。

此外，考虑到这些捆绑内部80%的共同激活概率，仍然有近20%的概率这些捆绑的神经元没有共同激活。因此，将两个4KB的随机读取组合起来可能会导致带宽浪费。为了减轻这个问题，对于使用4bit量化的模型，PowerInfer-2延迟等待第二个4K的读取，直到获得Gate神经元乘法的输出结果。具体而言，PowerInfer-2使用预测器确定Gate矩阵内神经元的激活，根据此信息启动捆绑的第一部分加载。然后，如果Gate神经元的输出（通过激活函数）非零，PowerInfer-2继续加载捆绑的第二部分，从而最小化不必要的I/O操作。

3.4 Neuron-Cluster-Level Pipeline

PowerInfer-2通过将计算与I/O活动重叠，隐藏部分I/O开销的功能。一种直接的方法是矩阵级重叠，在处理已经在内存中的神经元的同时，发出I/O命令从存储中检索矩阵神经元。当从存储加载神经元时，它们立即进行处理。虽然这种矩阵级重叠方法可以在计算过程中在一定程度上隐藏I/O操作的成本，但它仍然需要系统在继续下一步之前等待矩阵中的所有神经元完成，包括从存储中获取的神经元。如图4-a所示，假设一个矩阵包含8个神经元集群，其中4个位于内存中，剩下的4个位于存储中。一部分I/O操作可以隐藏在缓存的神经元计算之后。但由于较长的I/O时间，仍然会有CPU核必须等待I/O完成的情况。

为了消除I/O操作的等待时间，PowerInfer-2引入了神经元集群级流水线机制。这种机制基于一个观点：通过将神经元集群作为粒度，可以将I/O操作与多个矩阵的神经元集群计算重叠。具体来说，PowerInfer-2打破了矩阵计算之间的障碍；一旦一个神经元集群完成计算，它立即开始计算在内存中的下一个矩阵中的神经元集群。这种机制有效地减少了等待时间，如图4-b所示。

PowerInfer-2将神经元集群的执行过程分为5个顺序阶段：

• 确定门、上行和下行矩阵的行/列是否通过预测器激活（Pred）

• 从存储器中读取门矩阵的行权重（GIO）

• 计算门矩阵的行与输入向量的乘积（GC）

• 从存储器中读取上行和下行矩阵的行/列（UDIO）

• 分别计算上行和下行矩阵的行/列与输入向量的乘积（UDC）。

PowerInfer-2创建多个计算线程和一个I/O线程，分别处理这5个阶段的计算和I/O操作。这些线程的具体数量和它们将在哪些核上执行是由离线计划器确定的。

图5显示了计算和I/O线程如何工作以实现神经元集群流水线。在每个FFN块的开始时，所有神经元最初处于预测阶段，并被插入计算队列。计算线程处理这些神经元，仅推进那些被激活到后续阶段的神经元。然后，这些被激活的神经元被合并成神经元集群。如果一个神经元集群的门权重在内存中可用，该神经元集群进入GC阶段并返回计算队列。如果不可用，则设置为GIO并移动到I/O队列。同时，计算线程继续处理队列中的下一个神经元集群。与此同时，I/O线程从I/O队列中取出神经元，并根据需要执行I/O任务。UDIO和UDC的执行与GC和GIO类似的模式。

4. 生成执行计划

现在的智能手机配备了各种硬件规格，例如不同的CPU能力、I/O吞吐量和DRAM大小。在这些设备上部署LLM的用户也有不同的目标。有些人可能将生成速度和内存使用率之间的平衡放在首位，而有些人则旨在最大化硬件利用率以提高速度。此外，模型本身在权重数量、结构和稀疏性水平上也有所不同。为了管理这种复杂性，PowerInfer-2引入了一个离线计划器，用于对不同硬件配置的不同要求设计最优的执行计划。

5. 执行

PowerInfer-2是在PowerInfer的基础上开发的，PowerInfer是一个专为稀疏激活的LLM设计的先进服务框架，通过将额外的12K行C++代码集成到PowerInfer中实现。这些增强功能涵盖了几个关键领域，包括多态神经元引擎、神经元缓存、灵活的神经元加载和神经元集群级别的输入/输出流水线。

由于PowerInfer-2依赖于系统API（例如，将页面锁定在内存中的mlock），需要root权限，我们在Android平台上构建了它。即使不需要修改系统内核，已获取root权限的Android系统仍然为开发和调试系统时提供了相当大的灵活性。此外，PowerInfer-2本质上是设计为无需修改内核，因此可以轻松移植到其他操作系统，包括iOS平台。PowerInfer-2的当前实现支持各种不同模型大小的LLM，包括Llama-2系列（7B、13B）、TurboSparse-Mistral（7B）和TurboSparse-Mixtral（47B）。

6. 总结

本文介绍了PowerInfer-2，这是一个支持在智能手机上进行LLM的高速推理的框架，特别适用于超出设备内存容量的模型。PowerInfer-2的关键贡献是利用智能手机异构资源，将矩阵计算转化为更易管理的神经元集群计算。在两款智能手机上的评估结果表明，PowerInfer-2相对于现有的SOTA系统实现了高达29.2倍的加速，并且是第一个能够在智能手机上高效运行像TurboSparse-Mixtral-47B这样超大语言模型的推理框架。