EdgeMoE: Fast On-Device Inference of MoE-based Large Language Models

type

status

date

slug

summary

1 简介

基于一个观察：大多数计算依靠的小部分权重可以放在设备内存，大多数权重只参与小部分计算。这种特点很适用于分级存储结构。热门的权重放在内存，内存剩余的部分作为冷门权重的 buffer。

有了 buffer，只需要在需要的时候从存储中加载专家到内存即可。但是 I/O 开销又成为一个新问题。解决这个问题有两个方法：

通过量化直接减少 IO 数据量

用流水线隐藏延迟

但是，两个方法都面临挑战：

复杂的量化带来高压缩率的同时也带来了大量预处理开销；普通的量化又不能有效减少数据量。

MoE 系统在路由函数之前无法获取到前面的信息，所以专家预取是个挑战。如果不考虑专家的激活信息，仅仅用 cache 保存那些经常被用到的专家来提高命中率的话，收益又很有限，因为在训练的时候对不同专家的激活频率做了平衡。

EdgeMoE 解决上述问题的两个设计：专家位宽自适应和内存内专家管理

1.1 专家位宽自适应

一种轻量化预处理量化方法，具有专家级别自适应位宽的逐通道线性量化。

基于一个重要观察：不同层的专家甚至同一层的不同的专家在量化之后对模型精度的影响各不相同。

因此，EdgeMoE 用一个细粒度的、专家级别的位宽自适应来充分利用模型冗余。EdgeMoE 逐渐降低最鲁棒的那一小部分专家的位宽，直到精度下降到可容忍的阈值。同时也需要考虑那些进一步量化的专家下降到多小的位宽能够有效加速推理。

EdgeMoE 最终得到一个能够达到目标精度的尺寸尽可能小的混合精度模型。

1.2 内存内的专家管理

为了能实现流水线，EdgeMoE 在路由函数之前预测即将被激活的专家。

基于一个观察：在实际用用中，专家激活路径是不平衡的。这表明在激活的专家之间存在显著的相关性。

因此，在离线阶段，EdgeMoE 基于之前层的激活情况（还是激活值？）建立一个静态模型来估计当前层激活的专家的概率分布。

在在线推理阶段，EdgeMoE 查询这个模型并预加载最可能的专家。

此外，EdgeMoE 对专家缓冲区还设计了一个新的 cache 逐出策略：同时使用当前执行中的激活频率和他们的相对位置。

预测和加载以及逐出策略的目的都是为了最大化命中率

2 EdgeMoE设计

离线专家量化。根据用户在给定任务上指定的准确度损失（例如 5%），EdgeMoE 首先离线预处理预训练模型：分析专家重要性，然后根据评估的重要性将专家量化为不同的位宽。由此产生的模型包含一组具有不同位宽的专家，即使对于同一层中的专家也是如此。

在线专家管理。在运行时，EdgeMoE 实例化预加载/计算管道，并通过专家缓冲区动态管理设备内存和磁盘之间的专家。通过利用专家激活的统计配置文件，EdgeMoE 预先确定哪些专家在其路由器功能之前从磁盘获取，以及哪些专家在缓冲区已满时被逐出。

2.1 量化

2.1.1 选择量化算法

使用通道线性量化，因为它具有良好的准确性和快速的解压时间

通道方式线性量化使用缩放因子将量化的整数值统一映射到原始浮点值。每个通道的缩放因子由该通道内的最大绝对值和位宽可表示的范围确定。

2.1.2 分析专家重要性

对于专家，我们将它们量化为不同的位宽，例如INT2/4/8。理由是专家对模型精度表现出不同的重要性；我们希望最重要的专家具有高位宽，从而对模型精度做出更积极的贡献

为此，EdgeMoE 枚举所有专家，将每个专家量化为 INT2，并分析不同验证集（例如 ST-base-8）上产生的准确度损失。结果显示为热图：

例如，在 1 𝑠𝑡 变压器块将 1 𝑠𝑡 专家量化为 INT2 会使精度降低 0.44%，而将 2 𝑛𝑑 专家量化到相同精度会导致 0.59% 的降低。因此，2𝑛𝑑专家对量化更敏感，对模型精度更重要

EdgeMoE 获取专家重要性列表，该列表按照相应专家量化为 INT2 时的模型精度排序。该列表将用于构建运行时模型。

2.1.3 选择专家位宽

首先，EdgeMoE 决定模型所有专家的位宽范围，作为 EdgeMoE 进一步调整的基线。为此，EdgeMoE 枚举所有专家的可用位宽（例如 INT2/4/8 和 FP32）并分别测量模型精度。然后，EdgeMoE 将位宽的下限和上限设置为精度非常接近可容忍精度损失的位宽。

其次，EdgeMoE 根据位宽的下限和上限单独调整专家的位宽。它从之前获得的列表中的 topK 专家开始。由于它们不太重要，EdgeMoE 将它们量化为较低位宽（即 INT2），同时保持其余较高位宽（例如 INT4）。然后，EdgeMoE 测量所得模型的准确性。如果其精度损失仍然低于期望的目标，这意味着模型可以负担更多较低位宽的专家，EdgeMoE 会按照列表逐渐增加参数 𝐾，直到精度损失达到目标。否则，EdgeMoE 会降低参数 𝐾 以减少准确性损失，即通过将更多专家提升到更高的位宽。

2.2 内存专家管理

2.2.1 Experts re-batching for encoder

在编码器内执行 MoE 层期间，专家被逐个令牌激活，其中多个令牌可能会选择同一专家。

如果没有系统视图，在令牌调用时天真地拉动专家权重会导致重复加载，从而导致严重的延迟和 I/O 带宽浪费。

为了解决这个问题，EdgeMoE 通过重新批处理来重新组织输入令牌。

具体来说，EdgeMoE 重新排序编码器计算，将其从逐个令牌的方式转变为逐个专家的方式。根据专家的顺序，EdgeMoE 加载每个专家的权重，然后连接所有令牌，选择该专家进行计算。该策略确保每个专家的权重仅加载一次，从而防止冗余加载。

2.2.2 预加载和流水线

为了将专家加载与权重计算重叠，我们必须事先预测专家激活，而不是被动等待先前 MoE 层的路由器输出。

先验估计专家激活

如何预测专家激活？我们利用一个关键观察结果：连续层的专家激活在统计上是相关的。也就是说，给定 1..𝑛 − 1 层专家激活的先验知识，我们可以很有把握地估计每个专家在第 𝑛 层的激活概率，公式为 𝑃 (𝐸𝑛 = 𝑖|𝐸0, 𝐸1 .., 𝐸𝑛−1) 其中 𝑖 是专家的索引，𝑛 是层索引。

机会性预加载

EdgeMoE 利用之前的观察来机会性地预加载专家权重并执行管道：在离线阶段，根据之前的观察，EdgeMoE 在多个数据集上执行模型，以构建专家激活的统计概况。为此，EdgeMoE 生成一个字典，其中键表示连续两个 MoE 层中专家的激活状态，值表示单个专家在后续 MoE 层中被激活的概率。然后存储统计概况以供在线推理使用。

在在线阶段，在每个 MoE 层路由之前，EdgeMoE 采用前几层专家的激活状态作为查询统计概况的关键。然后，它按顺序将专家预加载到专家缓冲区（如果不存在），并按照估计的激活概率进行优先级排序。预加载停止，直到路由器完成，从而知道真正激活的专家。实际上，EdgeMoE 能够根据计算 I/O 加速差距为管道的每一层预加载 1-3 个专家。

流水线调度

EdgeMoE 实例化预加载/计算管道以并发执行计算

图 9 阐明了预测成功/失败情况下的管道调度。当 EdgeMoE 准确预测下一个 MoE 层专家的激活（这是一种常见情况）时，它可以显着减少端到端推理延迟，隐藏计算下的加载时间。作为最坏的情况（我们从未观察到），当所有预测都失败时，推理时间仅与按需加载专家一样长。

缓存驱逐策略

经典的缓存策略如 FIFO/LRU/LFU 是为操作系统设计的，主要基于数据访问历史和频率。 EdgeMoE 还利用了专家激活频率，但又结合了另一个独特的机会：由于 LLM 是使用顺序堆叠的transformer层构建的，因此每个专家的激活时间与其位置（即层索引）相关。具体来说，如果专家居住在即将执行的层中，则不会被驱逐，将被赋予较高的分数。

3 实验评估

EdgeMoE 与四个基线进行比较：

IO-FREE 假设所有模型权重都保存在内存中，因此不需要交换 I/O。这是计算效率最高的方法，但由于内存限制而无法扩展。

IO-EXP 将内存视为专家的缓存，并在激活后动态加载它们，类似于EdgeMoE。

IO-QEXP将上述方法与MoQE[43]相结合，将专家权重量化为INT4，并在推理过程中动态加载它们。与 EdgeMoE 类似，量化权重需要转换回 FP32，以便在设备处理器上进行快速推理。

STI 通过模型分片最小化推理精度，并在紧张的内存预算下实例化 IO/计算管道。它不区分专家和非专家的权重。为了公平比较，我们调整预加载分片的缓冲区大小，以便 STI 和 EdgeMoE 具有相同的内存占用