Capaciteitsbewuste Inferentie: Het Verminderen van het Straggler-effect in Mixture of Experts

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

March 7, 2025

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

cs.AI

Samenvatting

Het Mixture of Experts (MoE) is een effectieve architectuur voor het schalen van grote taalmodelen door gebruik te maken van sparse expert-activering, waardoor de afweging tussen prestaties en efficiëntie wordt geoptimaliseerd. Echter, onder expert-parallelisme lijdt MoE aan inefficiënties tijdens inferentie als gevolg van een onbalans in token-to-expert-toewijzing, waarbij sommige experts overbelast zijn terwijl anderen onderbenut blijven. Deze onbalans leidt tot slechte resourcebenutting en verhoogde latentie, aangezien de meest belaste expert de totale vertraging bepaalt, een fenomeen dat we definiëren als het \textit{Straggler Effect}. Om dit te mitigeren, stellen we Capacity-Aware Inference voor, inclusief twee belangrijke technieken: (1) \textit{Capacity-Aware Token Drop}, die overbelaste tokens verwijdert om de maximale latentie van MoE te reguleren, en (2) \textit{Capacity-Aware Token Reroute}, die overbelaste tokens herverdeelt naar onderbenutte experts, waardoor de tokenverdeling wordt gebalanceerd. Deze technieken optimaliseren gezamenlijk zowel het gebruik van zwaar belaste als licht belaste experts, wat resulteert in een efficiëntere MoE-inferentiepijplijn. Uitgebreide experimenten tonen de effectiviteit van onze methoden aan, met significante verbeteringen in inferentie-efficiëntie, zoals een gemiddelde prestatieverbetering van 0,2\% en een inferentieversnelling van 1,94 keer op 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.