arXiv:2503.05066cs.LGcs.AI2025-03被引 17

通过动态调整专家负载,显著提升MoE模型推理速度与效率。

Capacity-Aware Inference: Mitigating the Straggler Effect in Mixture of Experts

  • 根据专家容量动态丢弃过载令牌,缓解负载不均问题。
  • 引入扩展丢弃策略,利用闲置专家提升整体资源利用率。
  • 在多个语言和多模态模型上实现1.85倍加速,性能几乎无损。

Mixture of Experts (MoE) 通过稀疏专家激活实现了大语言模型的高效扩展,但在专家并行下,由于令牌分配不均,部分专家过载而其他专家空闲,导致整体推理延迟由最慢专家决定,这一现象称为‘长尾效应’。为此,我们提出容量感知令牌丢弃(Capacity-Aware Token Drop),通过限制专家容量、丢弃过载令牌,有效降低负载不均,仅造成0.9%性能下降即实现30%加速。进一步提出容量感知扩展丢弃(Capacity-Aware Expanded Drop),允许令牌在本地加入额外专家候选,再施加严格容量约束,提升低负载专家利用率。在语言与多模态MoE模型上的实验表明,该方法显著提升专家利用率、模型性能与推理效率;例如,在Mixtral-8×7B-Instruct上实现0.2%平均性能提升和1.85倍推理加速。代码已开源:https://github.com/CASE-Lab-UMD/Capacity-Aware-MoE。

原文摘要 · Abstract (English)

The Mixture of Experts (MoE) is an effective architecture for scaling large language models by leveraging sparse expert activation to balance performance and efficiency. However, under expert parallelism, MoE suffers from inference inefficiencies due to imbalanced token-to-expert assignment, where underloaded experts complete computations early but must wait for overloaded experts, leading to global delays. We define this phenomenon as the \textbf{\textit{Straggler Effect}}, as the most burdened experts dictate the overall inference latency. To address this, we first propose \textit{\textbf{Capacity-Aware Token Drop}}, which enforces expert capacity limits by discarding excess tokens from overloaded experts, effectively reducing load imbalance with minimal performance impact (e.g., $30\%$ speedup with only $0.9\%$ degradation on OLMoE). Next, given the presence of low-load experts remaining well below the capacity threshold, we introduce \textit{\textbf{Capacity-Aware Expanded Drop}}, which allows tokens to include additional local experts in their candidate set before enforcing strict local capacity constraints, thereby improving load balance and enhancing the utilization of underused experts. Extensive experiments on both language and multimodal MoE models demonstrate the effectiveness of our approach, yielding substantial gains in expert utilization, model performance, and inference efficiency, e.g., applying Expanded Drop to Mixtral-8$\times$7B-Instruct yields a {0.2\%} average performance improvement and a {1.85$\times$} inference speedup. The code is released at: https://github.com/CASE-Lab-UMD/Capacity-Aware-MoE.

MoE推理优化负载均衡大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。