大语言模型推理优化论文-EdgeMoE

论文地址：https://arxiv.org/pdf/2308.14352

代码仓库：https://github.com/UbiquitousLearning/mllm

研究机构：北京邮电大学、电子科技大学、中关村实验室

问题背景

大型语言模型（LLMs）如 GPT 和 Mixtral-8x7B 在通用 ML 任务上表现出色，但将其从数据中心部署到边缘设备（如手机）面临巨大挑战。

主要挑战在于其庞大的参数量导致在设备上运行时的内存和计算成本难以承受。

解决方案 - EdgeMoE

提出 EdgeMoE，一个专门为混合专家（Mixture-of-Experts, MoE）架构的稀疏大型语言模型设计的设备端推理引擎。

MoE LLMs 的特点是可以通过增加专家数量（参数）来扩展模型能力，同时保持计算复杂度相对恒定（仅激活部分专家）。

EdgeMoE 的核心设计理念是将模型分区存储到不同的存储设备中：

非专家权重（“热权重”）：常驻设备内存（因为它们每个 token 推理都需要使用）。

专家权重（“冷权重”）：存储在外部存储中，仅在被激活时才按需加载到内存中的一个专家缓冲区（expert buffer）。

这个设计基于一个关键观察：专家权重贡献了大部分内存占用，但只贡献了少量计算（由于稀疏激活）。

关键技术优化（降低专家 I/O 开销）

专家级位宽自适应（Expert-wise bitwidth adaptation）

在可接受的精度损失范围内，为不同的专家选择不同的量化位宽（如 INT2, INT4, INT8）。

根据专家对量化精度的敏感度（重要性）进行细粒度调整，优先降低对精度影响小的专家的位宽，从而减小专家尺寸，减少加载数据量。

算法细节：

1.把所有专家逐个量化到INT2，在不同的验证集上分析精度损失，形成一个不同层不同专家重要性的热力图。例子如下：

专家重要性热力图

2.把所有可用位宽（如INT2/4/8、FP32）分别去量化全部专家，得到一个最接近用户容忍损失的位宽上下界。在这个上下界内，调整每个专家的位宽，从热力图中选Top-K哥最不重要的专家开始，将其量化为下界，其余保持上界。测试所得模型精度，若精度损失仍低于用户目标，说明模型可承受更多低位宽专家，则继续沿列表增大 K，直至精度损失达到目标；否则减小 K，通过提升更多专家到高位宽来降低精度损失。

专家预加载（Expert preloading）

在专家被激活（路由函数执行）之前，预测下一个可能被激活的专家。

将预测专家的加载操作与当前层的计算操作流水线化（pipeline），以隐藏 I/O 延迟。

这基于另一个观察：专家激活路径（一个 token 经过各层激活的专家序列）在实践中的分布是高度偏斜的（Power Law），存在跨层相关性。

专家预加载工作流程图：

算法细节：

作者基于观察：连续层的专家激活在统计上相关，即给定前 0…n−1 层的专家激活，可以以高置信度估计第 n 层各专家的激活概率。在 SAMSum 数据集上运行 ST-base-8 模型。如图 8(左) 所示，若已知前两层激活，则第 3 层有 87.1% 概率激活第 5 号专家，即 P(E₃=5 | E₁=3, E₂=1)=87.1%；图 8(右) 进一步统计不同激活路径证实了这一点。

专家激活观察

离线阶段：基于上述观察，EdgeMoE 在多个数据集上执行模型，构建专家激活统计档案。生成一个字典，键为前两连续 MoE 层的专家激活状态，值为下一层各专家激活概率。该统计档案供在线推理使用。

在线阶段：在每个 MoE 层路由前，EdgeMoE 用前层激活作为键查询统计档案，按估计激活概率顺序预加载专家到专家缓冲区（若尚未存在）。预加载持续至路由器完成并获知真实激活专家为止。实践中，EdgeMoE 每层可预加载 1–3 个专家，具体取决于计算-I/O 速度差距。

主要成果

EdgeMoE运行的框架图如下：

EdgeMoE框架图

在流行的 MoE LLMs 和边缘设备（如 Raspberry Pi 4B, Jetson TX2）上进行了评估。

内存节省： 相比将整个模型加载到内存的基线，显著减少了内存占用（摘要中未提具体倍数，正文有详述）。

推理加速： 相比动态加载专家的基线（如 IO-EXP）和 STI 等内存优化基线，实现了显著的推理加速（摘要中未提具体倍数，正文有详述）。

里程碑： 首次在 Jetson TX2 等商用边缘设备上实现了 >10B 参数的 LLMs 的快速推理，且精度损失可忽略（≤2%）。

开源： 代码已公开在 GitHub (https://github.com/UbiquitousLearning/mllm)。

总结

EdgeMoE 是一个创新的设备端推理引擎，它通过将 MoE 模型分区存储（常驻非专家权重 + 按需加载专家权重）并结合专家级位宽自适应和预测性专家预加载两项关键技术，显著降低了专家 I/O 开销，从而在资源受限的边缘设备上实现了大型稀疏 MoE 语言模型的高效（内存+计算）推理。