Speculative MoE: Communication Efficient Parallel MoE Inference with Speculative Token and Expert Pre-scheduling

📄 arXiv: 2503.04398v3 📥 PDF

作者: Yan Li, Pengfei Zheng, Shuang Chen, Zewei Xu, Yuanhao Lai, Yunfei Du, Zhengang Wang

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

发布日期: 2025-03-06 (更新: 2025-03-19)

💡 一句话要点

Speculative MoE：通过推测Token和专家预调度，提升MoE模型通信效率

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

关键词: 混合专家模型 并行推理 通信优化 推测执行 专家并行

📋 核心要点

  1. 现有MoE推理框架，如DeepSpeed-MoE，在专家并行（EP）中存在大量all-to-all通信，成为性能瓶颈。
  2. Speculative MoE通过推测token路由路径和预调度专家，减少EP中的通信量，从而提升推理效率。
  3. 实验表明，Speculative MoE能显著提升DeepSpeed-MoE和SGLang在不同互连环境下的性能。

📝 摘要（中文）

MoE（混合专家模型）作为一种神经架构，能够将基于Transformer的现代LLM（大型语言模型）扩展到前所未有的规模。然而，大型MoE对计算能力、内存容量和内存带宽的巨大需求使得可扩展服务成为一项根本性挑战，而高效的并行推理已成为在延迟约束下获得足够吞吐量的必要条件。DeepSpeed-MoE是一种最先进的MoE推理框架，它采用包括EP（专家并行）、TP（张量并行）和DP（数据并行）的3D并行范式。然而，我们的分析表明，DeepSpeed-MoE的推理效率在很大程度上受到EP的瓶颈限制，EP是通过代价高昂的all-to-all集合通信来实现token激活路由的。我们的工作旨在通过一种名为Speculative MoE的技术来战略性地减少EP的通信开销，从而提升DeepSpeed-MoE。Speculative MoE具有两种推测并行化方案，即推测token shuffling和推测专家分组，它们预测突出token的专家路由路径，并在设备之间预先调度token和专家，以无损地减少EP的通信量。除了DeepSpeed-MoE之外，我们还将Speculative MoE构建到一种流行的MoE推理引擎SGLang中。实验表明，Speculative MoE可以显著提升快速同构和慢速异构互连上的最先进的MoE推理框架。

🔬 方法详解

问题定义：现有MoE模型推理框架，特别是DeepSpeed-MoE，在专家并行（EP）过程中，需要进行大量的all-to-all通信来路由token激活，这导致了显著的通信开销，成为推理效率的瓶颈。尤其是在大规模MoE模型和异构计算环境下，这个问题更加突出。现有方法缺乏对token路由路径的有效预测和优化，导致通信效率低下。

核心思路：Speculative MoE的核心思路是通过推测token的专家路由路径，并在实际计算之前预先调度token和专家，从而减少EP过程中的通信量。这种“推测执行”的思想类似于CPU中的分支预测，旨在通过提前预测和准备，减少等待时间，提高并行效率。

技术框架：Speculative MoE包含两个主要的推测并行化方案：推测token shuffling和推测专家分组。推测token shuffling预测哪些token将被路由到哪些专家，并提前将token移动到相应的设备上。推测专家分组则将可能被多个token访问的专家预先分组到同一设备上，减少跨设备通信。这两个方案协同工作，共同减少EP的通信开销。

关键创新：Speculative MoE的关键创新在于其推测性的并行化策略。与传统的MoE推理框架不同，Speculative MoE不是简单地按照固定的路由规则进行token路由，而是通过预测token的路由路径，提前进行token和专家的调度。这种推测性的方法能够有效地减少通信量，提高并行效率。

关键设计：Speculative MoE的具体实现细节包括：token路由路径的预测算法（具体算法未知），token和专家的预调度策略（如何平衡计算负载和通信开销未知），以及如何处理预测错误的情况（论文中提到是losslessly trim EP's communication volume，具体实现未知）。此外，Speculative MoE还考虑了不同互连环境下的优化策略，以适应不同的硬件条件。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Speculative MoE在DeepSpeed-MoE和SGLang上进行了实验验证，结果表明其能够显著提升推理性能。具体的性能提升幅度取决于硬件环境和模型规模，但总体而言，Speculative MoE能够有效地减少通信开销，提高并行效率，尤其是在异构互连环境下表现更佳。具体的性能数据未知。

🎯 应用场景

Speculative MoE可应用于各种需要大规模MoE模型推理的场景，例如自然语言处理、计算机视觉和推荐系统。通过提高推理效率，Speculative MoE能够降低部署成本，提升用户体验，并促进MoE模型在资源受限环境中的应用。该技术对于构建更高效、更可扩展的AI系统具有重要意义。

📄 摘要（原文）

MoE (Mixture of Experts) prevails as a neural architecture that can scale modern transformer-based LLMs (Large Language Models) to unprecedented scales. Nevertheless, large MoEs' great demands of computing power, memory capacity and memory bandwidth make scalable serving a fundamental challenge and efficient parallel inference has become a requisite to attain adequate throughput under latency constraints. DeepSpeed-MoE, one state-of-the-art MoE inference framework, adopts a 3D-parallel paradigm including EP (Expert Parallelism), TP (Tensor Parallel) and DP (Data Parallelism). However, our analysis shows DeepSpeed-MoE's inference efficiency is largely bottlenecked by EP, which is implemented with costly all-to-all collectives to route token activation. Our work aims to boost DeepSpeed-MoE by strategically reducing EP's communication overhead with a technique named Speculative MoE. Speculative MoE has two speculative parallelization schemes, speculative token shuffling and speculative expert grouping, which predict outstanding tokens' expert routing paths and pre-schedule tokens and experts across devices to losslessly trim EP's communication volume. Besides DeepSpeed-MoE, we also build Speculative MoE into a prevailing MoE inference engine SGLang. Experiments show Speculative MoE can significantly boost state-of-the-art MoE inference frameworks on fast homogeneous and slow heterogeneous interconnects.