LLM模型之PowerInfer2.0

PowerInfer-2

在智能手机上的快速大语言模型推理 原文链接：PowerInfer-2: Fast Large Language Model Inference on a Smartphone

1 简介

PowerInfer-2是一个为智能手机上大型语言模型(LLM)高速推理而设计的框架，特别适用于大小超过设备内存容量的模型。

1.1 性能提升

PowerInfer-2 是第一个为 TurboSparse-Mixral-47B 模型提供服务的系统，智能手机每秒生成 11.68 个tokens。对于完全在内存中拟合的模型，PowerInfer-2 在保持与 llama.cpp 和 MLC-LLM 相当的推理速度的同时，可以减少大约 40% 的内存使用。

  1. 低推理延迟：最大限度地减少预填充阶段（TTFT）和解码阶段（TBT）的推理延迟；
  2. 低内存占用：减少推理过程中的内存使用，即使模型大小超过设备内存限制，也能实现LLM的低延迟推理；
  3. 灵活性：确保设计能够无缝适应具有不同计算、内存和存储容量的智能手机。

2 概述

2.1 核心

PowerInfer-2 的核心是将 LLM 推理中典型的粗粒度矩阵计算分解为细粒度神经元集群计算。

2.2 神经元簇和架构

PowerInfer-2以神经元簇的粒度进行计算和I/O操作，神经元簇可以在计算过程中动态地由多个激活的神经元组成，神经元的数量由计算单元的计算能力决定，从而可以充分利用具有不同计算能力的XPU。

图 1 展示了PowerInfer-2的整体架构，它分为在线（右部分）和离线（左部分）程序。

  • 在线部分提供神经元簇粒度的推理，包括四个协作组件：多态神经元引擎、内存中神经元缓存、灵活的神经元加载和神经元簇级别 I /O 管道。
  • 离线部分描述在线推理中涉及的每个组件的具体配置并指导在线过程。
图1: PowerInfer-2整体架构

3 在线部分

3.1 多态神经元引擎

PowerInfer-2引入了一个多态神经元引擎，该引擎动态地将神经元组合成神经元簇，以利用LLM推理阶段和异质XPU的不同计算特征。用于预填充和解码阶段的两个计算工作流程如图 2 所示。

图2: PowerInfer-2的多态神经元引擎
    1. 预填充阶段采用以 NPU 为中心的工作流程，利用 NPU 进行计算，利用 CPU 进行准备；
    1. 解码阶段以 CPU 为中心，仅使用 CPU 核心来利用稀疏激活。

3.1.1 以 NPU 为中心的预填充

对于预填充阶段，神经元簇包含权重矩阵中的所有神经元，所有 tokens 都会同时处理。尽管这些tokens中的每一个都显示出高稀疏性并激活不同的神经元，但由于这些激活的聚合，总体稀疏性显着降低。因此，PowerInfer-2在预填充阶段不会使用预测器来计算激活的神经元，而是选择直接将所有神经元合并成一个大的神经元簇。

CPU核心不参与矩阵计算，在预填充阶段只为NPU执行必要的准备任务。

  • 首先，由于内存有限，PowerInfer-2 在预填充阶段依赖 CPU 核心将闪存中存储的权重加载到内存中。
  • 其次，由于当前的 NPU 不支持直接使用量化权重进行计算，PowerInfer-2 在 NPU 计算之前使用 CPU 内核对数据进行反量化。

图 2-a 演示了 CPU 和 NPU 如何协作在 Transformer 层粒度中执行预填充阶段推理。 NPU 计算需要使用与 CPU 共享的有限内存。因此，在 NPU 计算开始之前，CPU 应将所需的矩阵权重预加载到该共享内存中。

  1. 在NPU进行任何矩阵乘法之前，多个CPU中核从神经元缓存中读取量化矩阵权重，并提前将这些矩阵权重反量化为fp16，最终将结果存储在CPU和CPU之间的共享内存中。
  2. 同时，PowerInfer-2 使用另一个大核将下一层的所有矩阵权重异步预加载到神经元缓存中。
  3. 中核的反量化、NPU的计算和大核的I/O操作同时进行，以减少I/O开销。

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

3.1.2 以CPU为中心的解码

与预填充阶段不同，解码阶段在每次迭代期间集中于单个 token，表现出显着的稀疏性，因为权重矩阵中只有一小部分神经元（大约 10%）被激活并参与计算。因此，对于解码阶段，调用预测器来识别在开始计算之前将激活哪些神经元，然后引擎将这些激活的神经元合并成一个小的神经元簇，并利用 CPU 核心动态计算神经元簇:

  • 当批量大小为 1 时(即就一个 token)，CPU 内核上矩阵向量计算的延迟低于 NPU 上的延迟。
  • 由于稀疏性导致激活的神经元数量减少，CPU 内核最适合 XPU 中的这些较轻且稀疏的计算。

具体来说，PowerInfer-2 在解码阶段利用 CPU 内核来计算注意力块和 FFN 块。

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

图 2-b 说明了不同CPU核进行的解码阶段推断，总结如下：

  1. CPU 核心首先从神经元缓存中读取注意力块的权重，并使用输入向量进行计算。
  2. 运行预测器来确定后续权重矩阵中神经元的激活状态。
  3. CPU核心将激活的神经元分成几个簇，每个核心负责用输入向量计算其簇内激活的神经元，并最终聚合结果。
  4. 如果这些神经元位于神经元缓存内，CPU 内核将使用输入向量计算它们。在缓存未命中的情况下，即神经元不在神经元缓存中，CPU 核心上运行的 I/O 线程将神经元异步加载到缓存中，最终通知计算线程完成计算。

3.2 内存中神经元缓存

PowerInfer-2 引入了针对 LLM 内各种数据类型量身定制的分段神经元缓存设计。它将缓存划分为多个区域，每个区域都有特定的预取和逐出策略。

注意力块权重较小且激活较少，在整个运行时预加载并保留。

相比之下，FFN 块容易频繁激活热神经元，对这些神经元使用基于最近最少使用 (LRU) 的动态驱逐策略。这种方法确保热神经元更有可能保留在缓存中，而冷神经元经常被逐出并按需从闪存加载。重要的是，逐出过程不涉及写入存储，而只是丢弃内存中的权重。

PowerInfer-2 利用经典的双队列方法来实现其 LRU 神经元缓存，该缓存以单个神经元的粒度管理 LLM 权重。该系统维护两个双向链表队列，标记为活动和非活动，其中队列中神经元的顺序由它们最近访问的时间决定，最近访问的神经元位于队列的头部。

**LRU 缓存管理器的工作流程如下： 在运行时，所有神经元最初都会加入非活动队列。重新访问时，它们被提升到活动队列的前面。已经在活动队列中的神经元在后续访问时被移动到头部。为了管理缓存容量，当活动队列占满缓存空间的 90% 时，活动队列尾部的神经元将被移至非活动队列，直到活动队列的占用率降至 90% 以下。如果缓存达到容量，非活动队列尾部的神经元将被丢弃，为新条目腾出空间。*

3.3 灵活的神经元加载

PowerInfer-2 通过自适应捆绑和加载神经元来最大限度地减少 I/O 开销，这是由模型的量化决定的。

3.3.1 自适应捆绑

在推理过程仍然不可避免地导致未缓存神经元的 I/O 操作，为了优化 I/O 读取吞吐量并最小化 I/O 操作，PowerInfer-2 还捆绑相关神经元。尽管在单个FFN权重矩阵内的共激活在移除热神经元后变得不常见，但不同矩阵中相应位置的神经元通常会一起激活。因此PowerInfer-2 选择根据神经元粒度而不是矩阵结构来存储神经元权重。

**例如，Gate、Up、Down 矩阵中第 i 个神经元的共激活概率高达 80%，所以将 Gate、Up 和 Down 矩阵中第 i 个神经元的权重连接到单个条目中。*

3.3.2 灵活的神经元加载

考虑到不同模型的量化方法和 UFS I/O 的固有特性，PowerInfer-2 进一步引入了不同模型的不同 I/O 加载策略。

对于没有量化的模型，由于每个神经元占用的存储空间较大，PowerInfer-2使用更大粒度的随机读取来提高I/O带宽。 **例如，Llama-7B-FP16 中的单个神经元占用 8KB，来自 Gate、Up 和 Down 矩阵的神经元的总大小为 24KB。 PowerInfer-2 通过一次随机 I/O 读取，有效地将整个 24KB 激活包传输到内存中。*

对于 4 位量化模型，捆绑包大小设置为 8KB。以 Llama-7B 模型为例，每个神经元被量化为 4 位精度并占用 2.5KB（量化的 int4 值为 2KB，量化组的 FP16 尺度为 0.5KB），组合束大小达到 7.5KB。为了与存储介质 4KB 的最小读取粒度保持一致，PowerInfer-2 在捆绑包中额外补充了 0.5KB，将总数四舍五入为 8KB。然而，PowerInfer-2 并没有在单个 I/O 操作中加载这些 8KB 包，而是选择了 4KB 粒度。因为我们发现两次单独的 4KB 随机读取的带宽超过了单个 8KB 读取的带宽，从而优化了 I/O 读取过程。

此外，考虑到这些神经束内共激活的可能性为 80%，这些成束神经元仍有近 20% 的概率不被共激活。因此，组合两个 4KB 随机读取可能会导致带宽浪费。为了缓解这个问题，对于使用 4 位量化的模型，PowerInfer-2 会延迟第二个 4KB 读取，直到获得门神经元乘法的结果。(PowerInfer-2 使用预测器来确定门矩阵内神经元的激活，并根据此信息启动束第一部分的负载。之后，如果门神经元的输出（通过激活函数）非零，PowerInfer-2 将继续加载捆绑包的第二部分，从而最大限度地减少不必要的 I/O 操作。)

3.4 神经元簇级管道

为了减少 I/O 延迟，PowerInfer-2 引入了一种新颖的管道机制，可以同时处理神经元集群和 I/O 操作(将计算与 I/O 活动重叠来隐藏 I/O 开销)。图 3 是两种类型的管道，结合了矩阵向量乘法和五个核心（4 个计算核心和 1 个 I/O 核心）上的 I/O 操作。

图3: 神经元簇级管道
    1. 矩阵级管道将管道分成孤立的矩阵单元；
    1. PowerInfer-2 中的神经元簇级管道打破了矩阵障碍，并将其计算和 I/O 操作混合在神经元簇粒度中。

3.4.1 矩阵级管道

一种简单的方法是矩阵级重叠，它发出 I/O 命令从存储中检索矩阵神经元，同时处理内存中已有的神经元。当存储中的神经元被加载时，它们会立即被处理。虽然这种矩阵级重叠方法可以在一定程度上隐藏计算过程中I/O操作的成本，但它仍然要求系统等待矩阵内所有神经元的完成，包括从存储中获取的神经元，然后再继续下一个。

如图 3-a 所示，假设一个矩阵包含8个神经元簇，其中4个位于内存中，其余4个位于存储中。一部分 I/O 操作可以隐藏在缓存的神经元计算后面。但由于 I/O 时间较长，仍然存在 CPU 内核必须等待 I/O 完成的情况。

3.4.2 神经元簇级管道

为了消除 I/O 操作的等待时间，PowerInfer-2 引入了神经元簇级管道机制，如图3-b所示。通过将神经元簇作为粒度，可以在来自多个矩阵的神经元簇计算中重叠 I/O 操作。

具体来说，PowerInfer-2打破了矩阵计算之间的障碍；一旦一个神经元簇完成计算，它立即开始计算内存中下一个矩阵中的神经元簇。该机制有效减少了等待气泡。

PowerInfer-2将神经元簇的执行过程分为5个连续的阶段，创建了多个计算线程和一个I/O线程来分别处理这5个阶段的计算和I/O操作：

  1. 通过预测器(Pred)确定Gate、Up和Down矩阵的行/列是否被激活;
  2. 从存储器(GIO)中读取Gate矩阵的行权值;
  3. 计算Gate矩阵的行与输入向量(GC)的乘积;
  4. 从存储(UDIO)中读取Up和Down矩阵的行/列;
  5. 分别计算Up和Down矩阵的行/列与输入向量的乘积(UDC)。

图 4 展示了计算和 I/O 线程如何工作来实现神经元簇管道。

  1. 在每个FFN块开始时，所有神经元最初都处于Pred阶段，被插入到计算队列中。
  2. 计算线程处理这些神经元，仅将那些激活的神经元推进到后续阶段。
  3. 然后这些激活的神经元合并成神经元簇。
  4. 如果神经元簇的 Gate 权重在内存中可用，则神经元簇进入GC阶段并返回到计算队列。如果不是，则将其设置为 GIO 并移至 I/O 队列。
  5. 同时，计算线程继续处理队列中的下一个神经元簇。
  6. 同时，I/O 线程从 I/O 队列中取出神经元，根据需要执行 I/O 任务。
  7. UDIO 和UDC 的执行遵循与GC 和GIO 类似的模式。
图4: 神经元簇级管道工作流

4 离线部分

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

4.1 生成执行计划

PowerInfer-2 的离线规划器通过分析模型权重、用户输入和硬件规格来生成执行计划，输出计算、内存和 I/O 的配置。

  1. 输入配置：
    • 硬件：硬件配置的参数，如CPU FLOPS、I/O吞吐量和内存带宽。
    • 用户：用户自定义的参数，如CPU约束、内存限制、译码速度下界等。
    • Model：离线分析器收集的关于模型的参数，例如模型的大小、稀疏度级别和缓存特征等。
  2. 输出配置：
    • 计算：根据CPU和NPU的计算强度确定不同阶段或层中CPU和NPU的使用比例。
    • 内存：为了实现内存使用和推理性能之间的平衡，规划器允许用户在运行 PowerInfer-2 之前设置所需的推理速度。根据设置的速度，PowerInfer-2 计算所需的最佳缓存大小。
    • I/O：规划器触发分析器来测量模型的稀疏性以及冷热神经元的分布。

附录

LLM in a flash

将模型参数存储在闪存中，但将其按需带入DRAM，解决超过DRAM可用容量的llm的挑战。

  1. 减少数据加载量
    1. 有选择性地把部分参数常驻在DRAM中。而这部分常驻的参数就是Attention参数和Embedding参数，原因是因为它们占比较小。
    2. 借鉴了DejaVu。每次在推理时动态加载MLP预测为激活神经元对应的参数。
    3. 滑动窗口。保留处理过去k个token时的激活神经元所对应的参数在DRAM中，并在处理当前token时只对：1）部分多余的参数进行删除；2）缺少的参数进行加载。
  2. 提高传输数据量(读取大块数据) 未来改进: 把向上投影的第i列和向下投影的第i行捆绑存储。当激活第i个中间神经元时，这两部分数据会同时被使用。通过在闪存中将这些对应的列和行一起存储，可以将数据整合成更大的块进行读取。

原文链接: LLM in a flash: Efficient Large Language Model Inference with Limited Memory 参考 翻译