Inférence tenant compte de la capacité : Atténuation de l'effet de traînard dans les mélanges d'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

Résumé

Le Mixture of Experts (MoE) est une architecture efficace pour mettre à l'échelle les grands modèles de langage en exploitant l'activation parcimonieuse des experts, optimisant ainsi le compromis entre performance et efficacité. Cependant, dans un contexte de parallélisme des experts, le MoE souffre d'inefficacités lors de l'inférence en raison d'une répartition déséquilibrée des tokens entre les experts, où certains experts sont surchargés tandis que d'autres restent sous-utilisés. Ce déséquilibre entraîne une mauvaise utilisation des ressources et une latence accrue, car l'expert le plus sollicité dicte le délai global, un phénomène que nous définissons comme l'\textit{effet Straggler}. Pour atténuer ce problème, nous proposons l'inférence basée sur la capacité, incluant deux techniques clés : (1) le \textit{Capacity-Aware Token Drop}, qui supprime les tokens en surcharge pour réguler la latence maximale du MoE, et (2) le \textit{Capacity-Aware Token Reroute}, qui réalloue les tokens en excès vers des experts sous-utilisés, équilibrant ainsi la distribution des tokens. Ces techniques optimisent conjointement l'utilisation des experts à charge élevée et faible, conduisant à un pipeline d'inférence MoE plus efficace. Des expériences approfondies démontrent l'efficacité de nos méthodes, montrant des améliorations significatives en termes d'efficacité d'inférence, par exemple une augmentation moyenne des performances de 0,2\% et une accélération de l'inférence de 1,94 fois sur 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.