[论文解读] Fast Inference of Mixture-of-Experts Language Models with Offloading

Artyom Eliseev, D. Peter Mazur|arXiv (Cornell University)|Dec 28, 2023

Topic Modeling被引用 7

一句话总结

本文提出了一种面向MoE的卸载策略，通过利用专家访问模式和预测性预加载，加速了在消费级硬件上大型混合专家（MoE）语言模型的推理。结合混合量化（专家使用2–3比特，注意力层使用4–16比特）与LRU缓存及推测性专家加载，该方法使在桌面级GPU和免费版Google Colab实例上实现Mixtral-8x7B-Instruct的交互式推理速度达到每秒2–3个标记。

ABSTRACT

With the widespread adoption of Large Language Models (LLMs), many deep learning practitioners are looking for strategies of running these models more efficiently. One such strategy is to use sparse Mixture-of-Experts (MoE) - a type of model architectures where only a fraction of model layers are active for any given input. This property allows MoE-based language models to generate tokens faster than their dense counterparts, but it also increases model size due to having multiple experts. Unfortunately, this makes state-of-the-art MoE language models difficult to run without high-end GPUs. In this work, we study the problem of running large MoE language models on consumer hardware with limited accelerator memory. We build upon parameter offloading algorithms and propose a novel strategy that accelerates offloading by taking advantage of innate properties of MoE LLMs. Using this strategy, we build can run Mixtral-8x7B with mixed quantization on desktop hardware and free-tier Google Colab instances.

研究动机与目标

• 在GPU显存有限的消费级硬件上，实现类似Mixtral-8x7B-Instruct的大型MoE语言模型的交互式推理。
• 解决由于专家访问模式不规则和高内存带宽需求，导致的朴素参数卸载在MoE模型中效率低下的问题。
• 通过利用专家重用和早期门控预测等内在访问规律，为MoE架构优化卸载策略。
• 通过量化与智能卸载相结合，在低端硬件（包括配备T4 GPU的免费版Google Colab）上实现实际性能表现。
• 提供一种可部署的开源解决方案，使用户无需高端数据中心GPU即可在本地运行最先进的MoE模型。

提出的方法

• 设计一种面向MoE的卸载流水线，利用LRU缓存减少频繁使用专家在GPU与RAM之间的重复传输。
• 基于门控函数输出实现推测性专家预加载，将专家加载与计算重叠，减少空闲时间。
• 集成混合量化：使用2–3比特的HQQ量化处理MoE专家，注意力层、嵌入层和归一化层使用4–16比特精度。
• 采用混合推理引擎，动态将专家权重从GPU卸载到CPU内存，通过页锁定内存和异步传输实现高效内存管理。
• 利用早期层隐藏状态可预测后续层所需专家的事实，实现预测性加载策略。
• 通过最小化卸载频率并最大化数据重用，优化PCIe带宽受限系统，提升缓存与预加载效率。

Figure 1 : An example of expert loading pattern in Mixtral-8x7B-Instruct for select layers. Blue cells indicate that a certain expert was active when encoding a certain token; deeper blue indicates higher gating weight. Small gray squares show which experts are cached with an LRU cache for $k{=}2$ .

实验结果

研究问题

• RQ1 如何利用MoE模型中的专家访问模式，减少内存受限硬件上的卸载开销并提升推理速度？
• RQ2 在消费级GPU上运行时，MoE层的量化精度与模型性能之间的最优权衡是什么？
• RQ3 与朴素卸载相比，LRU缓存与推测性预加载在MoE推理中能多大程度减少GPU-RAM通信？
• RQ4 是否可以通过仅使用卸载与量化，在免费版Colab或桌面级GPU上实现Mixtral-8x7B-Instruct的交互式推理（2–3个标记/秒）？
• RQ5 在低显存环境下，MoE专用卸载方法与标准卸载框架（如Hugging Face Accelerate）相比性能如何？

主要发现

• 所提出的MoE专用卸载策略在RTX 3060（12GB）、RTX 3080移动版（16GB）和T4（16GB）GPU上，实现了Mixtral-8x7B-Instruct每秒2–3个标记的推理速度。
• 在配备T4的免费版Google Colab上，使用2比特量化专家并结合完整卸载、预加载与LRU缓存，模型实现每秒2.092个标记的生成速度。
• 与朴素卸载相比，LRU缓存与推测性预加载的组合在RTX 3060上将推理时间最多减少15%。
• 使用3比特HQQ量化处理专家，4比特处理注意力层，可在性能与精度之间取得良好平衡，困惑度下降可忽略。
• 在低端硬件上，该方法相比朴素卸载（如通过Hugging Face Accelerate）最高可提升2.5倍速度，尤其在PCIe受限系统中表现更优。
• 尽管显存较低，RTX 3060桌面版的性能几乎与更高配的RTX 3080移动版相当，表明主机到设备的带宽是关键瓶颈。

Figure 2 : (left) LRU cache hit ratio for different cache size $k$ ; (right) speculative loading recall when pre-loading a different number of experts. Regular lines represent loading 1 layer ahead; dashed line stands for 2 layers ahead; dotted line is 10 layers ahead.