Exploiting Inter-Layer Expert Affinity for Accelerating Mixture-of-Experts Model Inference

📄 arXiv: 2401.08383v2 📥 PDF

作者: Jinghan Yao, Quentin Anthony, Aamir Shafi, Hari Subramoni, Dhabaleswar K., Panda

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

发布日期: 2024-01-16 (更新: 2024-01-17)

💡 一句话要点

提出ExFlow以加速混合专家模型推理

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

关键词: 混合专家模型 推理加速 分布式系统 层间亲和性 整数规划 自然语言处理 高性能计算

📋 核心要点

  1. 现有的混合专家模型在分布式系统上推理时面临严重的通信瓶颈，影响高性能计算资源的有效利用。
  2. 本文提出ExFlow，通过利用层间专家亲和性，优化了专家路由和聚合过程，减少了通信开销。
  3. 实验结果表明，ExFlow在推理吞吐量上较现有方法提升了高达2.2倍，并显著降低了跨GPU路由延迟。

📝 摘要（中文）

在大型语言模型如生成预训练变换器中，混合专家范式已成为增强模型表现力和准确性的有效技术。然而，在分布式系统上并行推理时，专家路由和聚合所需的广泛Alltoall通信造成了显著的挑战，成为通信瓶颈。本文提出了一种名为ExFlow的轻量级优化技术，旨在显著加速这些MoE模型的推理。我们通过利用层间专家亲和性来减轻通信开销，提出了一种上下文一致的专家并行方案，显著减少了所需的Alltoall通信次数。通过设计高效的整数规划模型，我们证明了在适当放置专家的情况下，可以减少多达67%的跨GPU路由延迟，并在推理吞吐量上实现高达2.2倍的提升。

🔬 方法详解

问题定义：本文解决的是在分布式系统上推理混合专家模型时，由于专家路由和聚合所需的Alltoall通信造成的性能瓶颈。现有方法通常需要两次Alltoall通信，导致效率低下。

核心思路：论文的核心思路是通过利用层间专家亲和性，设计一种上下文一致的专家并行方案，从而将所需的Alltoall通信次数减少到一次，进而降低通信开销。

技术框架：整体架构包括专家路由模块、聚合模块和优化模块。专家路由模块负责根据输入token的条件概率进行专家选择，聚合模块则负责将各个专家的输出进行汇总，优化模块则通过整数规划模型优化专家的放置。

关键创新：最重要的技术创新在于提出了利用层间专家亲和性来优化推理过程的思路，与现有方法相比，ExFlow显著减少了通信次数和延迟。

关键设计：在设计中，采用了高效的整数规划模型来捕捉专家亲和性，并通过合理的GPU专家放置策略来实现性能提升。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，ExFlow在推理吞吐量上较现有最先进的混合专家实现提升了高达2.2倍，同时在跨GPU路由延迟上减少了多达67%。这些结果表明，ExFlow在优化混合专家模型推理方面具有显著优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、机器翻译和对话系统等大型语言模型的推理任务。通过加速混合专家模型的推理过程，ExFlow能够提高这些系统的响应速度和处理能力，进而提升用户体验和系统效率。

📄 摘要（原文）

In large language models like the Generative Pre-trained Transformer, the Mixture of Experts paradigm has emerged as a powerful technique for enhancing model expressiveness and accuracy. However, deploying GPT MoE models for parallel inference on distributed systems presents significant challenges, primarily due to the extensive Alltoall communication required for expert routing and aggregation. This communication bottleneck exacerbates the already complex computational landscape, hindering the efficient utilization of high-performance computing resources. In this paper, we propose a lightweight optimization technique called ExFlow, to largely accelerate the inference of these MoE models. We take a new perspective on alleviating the communication overhead by exploiting the inter-layer expert affinity. Unlike previous methods, our solution can be directly applied to pre-trained MoE models without any fine-tuning or accuracy degradation. By proposing a context-coherent expert parallelism on distributed systems, our design only uses one Alltoall communication to deliver the same functionality while previous methods all require two Alltoalls. By carefully examining the conditional probability in tokens' routing across multiple layers, we proved that pre-trained GPT MoE models implicitly exhibit a strong inter-layer expert affinity. We then design an efficient integer programming model to capture such features and show that by properly placing the experts on corresponding GPUs, we can reduce up to 67% cross-GPU routing latency. Our solution beats the cutting-edge MoE implementations with experts from 8 to 64, with up to 2.2x improvement in inference throughput. We further provide a detailed study of how the model implicitly acquires this expert affinity at the very early training stage and how this affinity evolves and stabilizes during training.