Inferenza Consapevole della Capacità: Mitigazione dell'Effetto Straggler nei Modelli a Miscela di Esperti

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

March 7, 2025

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

cs.AI

Abstract

La Mixture of Experts (MoE) è un'architettura efficace per scalare modelli linguistici di grandi dimensioni sfruttando l'attivazione sparsa degli esperti, ottimizzando il compromesso tra prestazioni ed efficienza. Tuttavia, in condizioni di parallelismo degli esperti, la MoE soffre di inefficienze nell'inferenza a causa di un'assegnazione sbilanciata dei token agli esperti, dove alcuni esperti sono sovraccaricati mentre altri rimangono sottoutilizzati. Questo squilibrio porta a una scarsa utilizzazione delle risorse e a un aumento della latenza, poiché l'esperto più carico determina il ritardo complessivo, un fenomeno che definiamo come \textit{Effetto Straggler}. Per mitigare questo problema, proponiamo l'Inferenza Consapevole della Capacità, che include due tecniche chiave: (1) \textit{Token Drop Consapevole della Capacità}, che scarta i token sovraccaricati per regolare la latenza massima della MoE, e (2) \textit{Token Reroute Consapevole della Capacità}, che riassegna i token in eccesso agli esperti sottoutilizzati, bilanciando la distribuzione dei token. Queste tecniche ottimizzano collettivamente sia l'utilizzo degli esperti ad alto carico che quelli a basso carico, portando a una pipeline di inferenza MoE più efficiente. Esperimenti estensivi dimostrano l'efficacia dei nostri metodi, mostrando miglioramenti significativi nell'efficienza dell'inferenza, ad esempio un aumento medio delle prestazioni dello 0,2\% e un'accelerazione dell'inferenza di 1,94 volte su 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.