MegaScale-Infer: Serving Mixture-of-Experts at Scale with Disaggregated Expert Parallelism

本文提出MegaScale-Infer系统，通过分离注意力模块和FFN模块的并行策略以及高效M2N通信库，优化大规模MoE模型的推理效率，实现高达1.90倍的吞吐量提升。

Large Language Model, Efficiency, Multimodality, Pre-training, Fine-tuning, Human-AI Interaction

Ruidong Zhu, Ziheng Jiang, Chao Jin, Peng Wu, Cesar A. Stuardo, Dongyang Wang, Xinlei Zhang, Huaping Zhou, Haoran Wei, Yang Cheng, Jianzhe Xiao, Xinyi Zhang, Lingjun Liu, Haibin Lin, Li-Wen Chang, Jianxi Ye, Xiao Yu, Xuanzhe Liu, Xin Jin, Xin Liu

ByteDance Seed, Peking University

Generated by grok-3-mini-latest

Background Problem

混合专家（MoE）模型在扩展大型语言模型（LLM）时展示了巨大的潜力，能够提升性能并降低计算复杂度。然而，在实际推理场景中，MoE的稀疏激活架构导致前馈网络（FFN）从计算密集型转变为内存密集型，从而显著降低GPU利用率并增加运营成本。具体问题包括：推理过程中注意力模块的内存密集性导致低利用率，而FFN模块在MoE稀疏性下无法充分利用GPU计算能力；此外，批量大小受限于GPU内存和响应延迟约束，进一步恶化了效率问题。本文的工作起点是针对这些挑战，优化大规模MoE模型的推理效率，以减少不必要的计算成本。

Method

• 核心思想： 通过分离注意力模块和FFN模块（即专家模块），实现独立的缩放和异构部署，最大化GPU利用率并降低成本。分离策略允许注意力模块使用数据并行，FFN模块使用专家并行。
• 如何实现： 采用分离专家并行策略，包括ping-pong管道并行，将请求批次分割为多个微批次，并在注意力模块和专家模块之间交替处理以隐藏通信开销。具体步骤如下：
  1. 分离部署： 将注意力模块和FFN模块部署在不同的GPU上，注意力模块使用数据并行复制，FFN模块使用专家并行分发。
  2. ping-pong管道并行： 将全局批次分割为$m$个微批次，确保注意力模块和FFN模块交替工作，满足条件$T_a \approx T_e$、$T_c < T_f$和$m \times T_f \ge 2 \times (T_f + T_c)$，其中$T_a$和$T_e$分别是注意力模块和专家模块的计算时间，$T_c$是通信时间，$T_f = \max(T_a, T_e)$。
  3. 部署计划优化： 基于性能模型搜索最佳部署计划，包括张量并行大小、微批次数量和批量大小，使用算法枚举可行方案并模拟性能。
  4. 高性能M2N通信： 开发自定义通信库，消除不必要的GPU-to-CPU数据拷贝、组初始化开销和GPU同步，通过RDMA write with immediate和流量优化（如高优先级ACK和拥塞控制微调）减少延迟。
• 关键步骤： 不修改模型本身，只在推理阶段调整并行策略和通信机制，确保在保持低延迟的同时提高吞吐量。

Experiment

• 实验设置： 使用Mixtral-8×22B、DBRX和Scaled-MoE（参数规模132B到317B）模型，数据类型为bfloat16。工作负载基于真实生产数据，输入和输出序列长度中位数分别为571和159个token。实验在同构（NVIDIA Ampere 80GB GPU）和异构（H20和L40S GPU）集群上进行，延迟要求设置为每输出token 150毫秒。比较基线包括vLLM和TensorRT-LLM，支持FlashAttention和连续批处理。
• 为什么这样设计： 实验旨在验证分离策略在不同模型和硬件下的有效性，评估吞吐量、延迟和成本效率。异构部署将注意力模块分配到内存带宽更优的GPU（如H20），FFN模块分配到计算能力更优的GPU（如L40S），以最大化资源利用。
• 结果分析： 结果符合预期，MegaScale-Infer在同构集群下实现高达1.90倍的每GPU吞吐量提升，在异构集群下实现1.86倍的每成本吞吐量提升。M2N通信库比NCCL减少68.2%中位延迟和92.9%尾延迟，提高4.2倍吞吐量。去除ping-pong管道并行时吞吐量下降，证明了其在隐藏通信开销方面的作用。实验显示，优化部署计划显著减少GPU空闲时间，吞吐量随数据并行度增加而线性提升，直至平衡计算时间。

Further Thoughts

这项工作突出了资源分离在AI推理中的潜力，或许可以扩展到其他稀疏模型或结合联邦学习场景，以进一步减少跨设备通信开销；同时，异构部署策略可能启发边缘计算中的LLM服务优化，平衡延迟和成本；此外，与现有的推理加速框架（如vLLM）整合，可能实现更全面的性能提升，但需注意在动态工作负载下的负载均衡挑战。