Fast Inference of Mixture-of-Experts Language Models with Offloading

📄 arXiv: 2312.17238v1 📥 PDF

作者: Artyom Eliseev, Denis Mazur

分类: cs.LG, cs.AI, cs.DC

发布日期: 2023-12-28

备注: Technical report

💡 一句话要点

提出一种加速MoE语言模型推理的Offloading策略，可在消费级硬件上运行Mixtral-8x7B。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 MoE 参数Offloading 推理加速 资源受限设备

📋 核心要点

  1. 现有MoE模型虽然推理速度快，但模型体积庞大，对硬件要求高，难以在资源受限的设备上运行。
  2. 论文提出一种新的Offloading策略，利用MoE模型的特性加速参数Offloading，降低内存需求。
  3. 实验表明，该策略可以在消费级硬件和Google Colab上运行Mixtral-8x7B模型，具有实际应用价值。

📝 摘要（中文）

随着大型语言模型（LLM）的广泛应用，许多深度学习从业者都在寻找更高效地运行这些模型的策略。一种策略是使用稀疏混合专家模型（MoE），这种模型架构的特点是对于任何给定的输入，只有一部分模型层是激活的。这种特性使得基于MoE的语言模型比其稠密模型能更快地生成token，但也因为拥有多个专家而增加了模型大小。不幸的是，这使得最先进的MoE语言模型在没有高端GPU的情况下难以运行。本文研究了在加速器内存有限的消费级硬件上运行大型MoE语言模型的问题。我们在参数Offloading算法的基础上，提出了一种新的策略，该策略利用MoE LLM的内在属性来加速Offloading。使用这种策略，我们可以在桌面硬件和免费的Google Colab实例上运行混合量化的Mixtral-8x7B。

🔬 方法详解

问题定义：论文旨在解决大型MoE语言模型在资源受限的硬件上部署和推理的问题。现有方法，如直接加载整个模型，需要大量的GPU内存，使得在消费级硬件上运行这些模型变得不可行。参数Offloading是一种可能的解决方案，但传统的Offloading方法没有充分利用MoE模型的特性，效率有待提高。

核心思路：论文的核心思路是利用MoE模型中专家（Experts）的稀疏激活特性，只Offloading当前激活的专家参数。由于每个token只需要激活一部分专家，因此可以显著减少需要Offloading的参数量，从而加速推理过程。这种方法避免了频繁地在CPU和GPU之间传输大量不必要的参数。

技术框架：整体框架包括以下几个主要阶段：1) 模型加载：将MoE模型的参数存储在CPU内存中。2) 路由计算：根据输入token，确定需要激活的专家。3) 参数Offloading：将激活的专家的参数从CPU内存Offloading到GPU内存。4) 前向传播：在GPU上执行激活专家的前向传播计算。5) 结果返回：将计算结果返回给CPU。这个过程循环进行，直到生成完整的输出序列。

关键创新：论文的关键创新在于提出了针对MoE模型的自适应Offloading策略。与传统的Offloading方法相比，该策略能够根据token的路由结果，动态地选择需要Offloading的专家参数，从而显著减少了Offloading的数据量和时间。这种自适应性是利用MoE模型稀疏激活特性的关键。

关键设计：关键设计包括：1) 专家选择机制：高效地确定每个token需要激活的专家。2) Offloading调度策略：优化CPU和GPU之间的数据传输，减少延迟。3) 混合精度量化：使用混合精度量化进一步压缩模型大小，降低内存需求。4) 针对特定硬件的优化：针对CPU和GPU的架构特点，进行针对性的性能优化。

📊 实验亮点

实验结果表明，该方法可以在桌面硬件（如配备消费级GPU的PC）和免费的Google Colab实例上成功运行Mixtral-8x7B模型。通过混合精度量化和自适应Offloading策略，显著降低了内存需求和推理延迟，使得在资源受限的环境中部署大型MoE模型成为可能。具体的性能数据（如推理速度、内存占用等）未在摘要中给出，属于未知信息。

🎯 应用场景

该研究成果可应用于各种需要部署大型MoE语言模型的场景，例如：在本地设备上运行AI助手、在资源受限的边缘设备上进行自然语言处理、以及在云服务器上降低MoE模型的推理成本。该方法使得更多用户能够在消费级硬件上体验到先进的AI技术，加速了LLM的普及。

📄 摘要（原文）

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.