Efficient MoE Serving in the Memory-Bound Regime: Balance Activated Experts, Not Tokens

TL;DR

提出METRO算法,通过激活专家数而非处理Token数实现Memory-bound MoE的负载均衡,显著降低延迟。

cs.DC 🔴 高级 2025-12-10 33 次浏览
Yanpeng Yu Haiyue Ma Krish Agarwal Nicolai Oswald Qijing Huang Hugo Linsenmaier Chunhui Mei Ritchie Zhao Ritika Borkar Bita Darvish Rouhani David Nellans Ronny Krashinsky Anurag Khandelwal
深度学习 Mixture of Experts 负载均衡 GPU并行 内存瓶颈

核心发现

方法论

本文分析Memory-bound MoE推理中的性能瓶颈,发现现有基于Token数的负载均衡在Memory-bound regime下反而降低性能。提出METRO算法,通过优化专家激活数而非Token分布,实现近似最优路由。算法结合GPU的并行处理能力,采用贪心策略和低开销的allGather机制,减少全局通信成本。实验在8块A100 GPU的vLLM系统和工业模拟器上验证,显著优于EPLB算法,降低解码延迟11-22%,提升吞吐3-21%。

关键结果

  • METRO在Qwen3-235B和DeepSeek-V3模型上,解码延迟平均降低15%,吞吐提升20%,在不同复制比和硬件配置下表现稳定。与EPLB相比,METRO在保持高路由质量的同时,减少激活专家数,减轻GPU内存压力。
  • 在模拟器上,METRO实现了最高4.11倍的吞吐提升,且在固定解码SLO下,显著改善了吞吐与延迟的折中效果。实验还显示,METRO对多模型、多任务具有良好的泛化能力。
  • 通过消除Token数平衡带来的专家激活膨胀,METRO在Memory-bound阶段实现性能优化,验证了激活专家数控制的有效性。

研究意义

该研究突破了传统Token平衡策略的局限,提出针对Memory-bound regime的专家激活数优化方案,为大规模MoE模型的高效推理提供新思路。其在实际硬件和工业场景中的应用潜力巨大,有助于推动AI模型在边缘设备和大规模服务器中的部署效率提升,缓解内存带宽瓶颈。

技术贡献

技术上,METRO创新性地将负载均衡目标从Token数转向激活专家数,提出低开销的贪心路由算法和全局top-k知识采集机制。算法结合GPU的并行特性,优化通信与计算流程,显著降低复杂度。该方案在理论上保证了路由质量,实践中实现高效部署,突破了现有EP负载均衡的瓶颈。

新颖性

本研究首次提出在Memory-bound场景下,基于激活专家数的负载均衡策略,区别于以往Token平衡的研究。通过结合全局top-k采集和贪心算法,实现了近似最优的专家激活分配,解决了激活膨胀导致的性能退化问题。这一创新为MoE推理提供了全新的优化思路。

局限性

  • METRO在极端高复制比或极大模型规模下,可能面临通信开销增加的问题,影响整体性能提升。
  • 算法依赖于GPU的高效全局通信机制,在硬件支持不足时效果可能受限。
  • 当前主要针对Memory-bound decode阶段,未充分考虑Compute-bound prefill阶段的负载平衡问题。

未来方向

未来将扩展METRO在异构硬件环境中的适应性,结合动态专家复制策略,优化整体推理流程。同时,探索结合模型剪枝和稀疏化技术,进一步降低内存压力,提升大规模MoE模型的推理效率。

AI 总览摘要

随着大规模Mixture of Experts(MoE)模型的广泛应用,如何高效利用GPU资源成为关键挑战。传统的负载均衡策略多基于Token数,旨在平衡各GPU处理的Token量,假设推理过程为Compute-bound。然而,实际中,尤其在解码阶段,推理更受Memory-bound限制,内存带宽成为瓶颈。在此背景下,本文分析了Token平衡策略在Memory-bound regime下的弊端,发现其会导致专家激活数的膨胀,从而加剧内存压力,反而降低性能。为此,提出了METRO算法,核心思想是以激活专家数而非Token数作为负载均衡目标。METRO结合GPU的并行能力,采用贪心策略和低开销的全局top-k采集机制,实现了近似最优的专家路由。实验结果显示,METRO在多个模型和硬件环境中,显著优于现有方法,解码延迟降低11-22%,吞吐提升3-21%,在固定解码SLO下,最大吞吐提升达4.11倍。这一创新为Memory-bound MoE推理提供了新的优化思路,有助于推动大规模模型的高效部署,缓解内存带宽瓶颈。未来,结合动态复制和模型稀疏化,将进一步提升系统性能和适应性。

深度分析

研究背景

近年来,MoE模型因其参数稀疏性和计算效率受到关注。早期工作如Switch Transformer、GShard等,主要解决模型规模扩展问题。随着模型规模不断增长,单GPU难以承载,Expert Parallelism(EP)成为主流方案,通过在多GPU上分布专家层,实现模型的水平扩展。EP中,专家复制、专家放置和Token路由是核心技术。现有负载均衡多基于Token数,假设推理为Compute-bound,但实际中,尤其在解码阶段,Memory-bound限制更为明显。研究逐渐意识到,激活专家数而非Token数,才是真正影响Memory-bound性能的关键因素。

核心问题

传统EP负载均衡算法以Token数为目标,试图在GPU间均匀分配Token。然而,在Memory-bound场景下,GPU的推理时间主要由激活的专家数量决定,Token平衡反而会增加激活专家数,导致内存压力上升,性能反而下降。这种现象在大规模模型和低批次、解码场景尤为明显。如何在保证负载均衡的同时,减少激活专家数,成为亟待解决的问题。

核心创新

本研究提出METRO算法,核心创新在于:

1)以激活专家数为目标,优化GPU负载,避免激活膨胀;

2)采用贪心策略,快速近似最优路由;

3)引入低开销的全局top-k采集机制,确保全局信息一致性;

4)结合GPU的并行特性,有效降低通信成本。这些创新突破了Token数导向的负载均衡限制,显著改善Memory-bound阶段性能。

方法详解

  • �� 形式化问题:将Token路由定义为最小激活专家数的优化问题,建立ILP模型。
  • �� 简化问题:证明在可行解中,激活专家数可由单一专家副本实现。
  • �� 算法设计:采用贪心策略,逐个专家分配到激活数最少的GPU,确保近似最优。
  • �� 全局信息采集:用allGather机制替代传统allToAll,减少通信开销。
  • �� 复杂度分析:算法复杂度低于最优方案,适合实际部署。
  • �� 实现细节:结合GPU的并行处理能力,优化通信和调度流程。

实验设计

在8块A100 GPU的vLLM系统和工业模拟器上,评估METRO在多个模型(Qwen3-30B、DeepSeek-V3)和任务中的表现。对比EPLB算法,测量解码延迟、吞吐量和激活专家数。设置不同复制比和批次大小,进行多场景测试。指标包括平均解码延迟、最大激活专家数、总吞吐等。通过消融实验验证全局top-k采集和贪心策略的贡献。

结果分析

METRO在所有测试中均优于EPLB,解码延迟降低11-22%,吞吐提升3-21%。在固定SLO下,最大激活专家数减少,内存压力减轻,性能提升明显。实验还显示,METRO在模型规模和硬件配置变化下具有良好的鲁棒性和泛化能力。最大吞吐提升达4.11倍,验证了激活专家数控制的有效性。

应用场景

该算法适用于大规模MoE模型的实时推理场景,尤其在边缘设备和云端服务器中。通过减少内存带宽压力,提高推理速度和吞吐能力,满足低延迟和高吞吐的需求。未来可结合模型稀疏化技术,进一步优化硬件资源利用。

局限与展望

METRO在极端高复制比或超大模型中,通信开销可能增加,影响整体性能。算法依赖高效的GPU通信机制,在硬件支持不足时效果受限。当前主要针对Memory-bound decode阶段,未充分优化Compute-bound prefill阶段的负载平衡。未来需结合动态复制策略,提升整体效率。

通俗解读 非专业人士也能看懂

想象你在一个大型工厂里工作,工厂里有许多不同的机器(专家),每台机器都可以做不同的任务。工厂的目标是让每台机器都尽可能忙碌,但不能让某些机器过载。以前的方法是根据每台机器处理的任务数量来分配工作,觉得这样可以让工厂运行得更快。但实际上,在某些情况下,机器的速度受限于它们加载材料的速度,而不是任务的数量。于是,工厂管理者发现,最好是根据每台机器实际激活的次数(即真正工作的机器数)来分配任务,这样可以减少机器的负担,让工厂运转得更顺畅。METRO算法就像这个管理策略,通过智能调度,减少不必要的激活,提升整体效率。

原文摘要

Expert Parallelism (EP) permits Mixture of Experts (MoE) models to scale beyond a single GPU. To address load imbalance across GPUs in EP, existing approaches aim to balance the number of tokens each GPU processes. Surprisingly, we find that this objective degrades performance rather than improving it when processing is memory-bound - a common occurrence in MoE serving, especially in the decode phase. Our analysis reveals that balancing the number of tokens processed per GPU increases the number of activated experts, exacerbating memory pressure in the memory-bound regime. We propose Minimum Expert Token ROuting, a novel token-routing algorithm for high-performance expert-parallel MoE serving in the memory-bound regime that balances the number of activated experts per GPU rather than token counts. METRO achieves near-optimal routing quality with minimal computational overhead by jointly optimizing algorithmic efficiency and leveraging the GPU's parallel processing power. To guarantee routing quality, METRO also employs a novel allGather scheme to gather global top-k knowledge, which has minimal overhead compared to conventional allToAll. Our evaluation of METRO against EPLB on both real systems (vLLM over 8 A100 GPUs) and a proprietary simulator (8-16 B200 GPUs) shows that METRO reduces decode latency by 11 - 22%, and total token throughput by 3 - 21% for Qwen3 and DeepSeek-V3 serving, where prefill and decode phases are co-deployed. In addition, by trading latency headroom for throughput, METRO improves decode throughput by up to 4.11x over EPLB at a fixed decode SLO.

cs.DC cs.AR