Kapazitätsbewusste Inferenz: Minderung des Straggler-Effekts in Mixture of Experts

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

March 7, 2025

Autoren: Shwai He, Weilin Cai, Jiayi Huang, Ang Li

cs.AI

Zusammenfassung

Das Mixture of Experts (MoE) ist eine effektive Architektur zur Skalierung großer Sprachmodelle, die durch die Nutzung spärlicher Expertenaktivierung den Kompromiss zwischen Leistung und Effizienz optimiert. Unter Expertenparallelismus leidet MoE jedoch an Inferenzineffizienzen aufgrund unausgeglichener Token-zu-Experten-Zuweisung, bei der einige Experten überlastet sind, während andere unterausgelastet bleiben. Dieses Ungleichgewicht führt zu einer schlechten Ressourcennutzung und erhöhter Latenz, da der am stärksten belastete Experte die Gesamtverzögerung bestimmt – ein Phänomen, das wir als \textit{Straggler-Effekt} bezeichnen. Um dies zu mildern, schlagen wir Capacity-Aware Inference vor, das zwei Schlüsseltechniken umfasst: (1) \textit{Capacity-Aware Token Drop}, das überlastete Token verwirft, um die maximale Latenz von MoE zu regulieren, und (2) \textit{Capacity-Aware Token Reroute}, das überlaufende Token unterausgelasteten Experten neu zuweist und so die Tokenverteilung ausgleicht. Diese Techniken optimieren gemeinsam sowohl die Nutzung hoch- als auch niedrig belasteter Experten und führen zu einer effizienteren MoE-Inferenzpipeline. Umfangreiche Experimente demonstrieren die Wirksamkeit unserer Methoden und zeigen signifikante Verbesserungen in der Inferenzeffizienz, z. B. eine durchschnittliche Leistungssteigerung von 0,2\% und eine 1,94-fache Beschleunigung der Inferenz bei Mixtral-8x7B-Instruct.

English

The Mixture of Experts (MoE) is an effective architecture for scaling large language models by leveraging sparse expert activation, optimizing the trade-off between performance and efficiency. However, under expert parallelism, MoE suffers from inference inefficiencies due to imbalanced token-to-expert assignment, where some experts are overloaded while others remain underutilized. This imbalance leads to poor resource utilization and increased latency, as the most burdened expert dictates the overall delay, a phenomenon we define as the \textit{Straggler Effect}. To mitigate this, we propose Capacity-Aware Inference, including two key techniques: (1) \textit{Capacity-Aware Token Drop}, which discards overloaded tokens to regulate the maximum latency of MoE, and (2) \textit{Capacity-Aware Token Reroute}, which reallocates overflowed tokens to underutilized experts, balancing the token distribution. These techniques collectively optimize both high-load and low-load expert utilization, leading to a more efficient MoE inference pipeline. Extensive experiments demonstrate the effectiveness of our methods, showing significant improvements in inference efficiency, e.g., 0.2\% average performance increase and a 1.94times inference speedup on Mixtral-8times7B-Instruct.