Inferencia Consciente de la Capacidad: Mitigando el Efecto del Rezagado en Mezclas de Expertos

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

March 7, 2025

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

cs.AI

Resumen

La Mezcla de Expertos (MoE, por sus siglas en inglés) es una arquitectura efectiva para escalar modelos de lenguaje grandes al aprovechar la activación dispersa de expertos, optimizando el equilibrio entre rendimiento y eficiencia. Sin embargo, bajo el paralelismo de expertos, MoE sufre ineficiencias en la inferencia debido a una asignación desequilibrada de tokens a expertos, donde algunos expertos están sobrecargados mientras que otros permanecen subutilizados. Este desequilibrio conduce a una mala utilización de recursos y a un aumento de la latencia, ya que el experto más cargado dicta el retraso general, un fenómeno que definimos como el \textit{Efecto Rezagado}. Para mitigar esto, proponemos Inferencia Consciente de la Capacidad, que incluye dos técnicas clave: (1) \textit{Descarte de Tokens Consciente de la Capacidad}, que descarta tokens sobrecargados para regular la latencia máxima de MoE, y (2) \textit{Redirección de Tokens Consciente de la Capacidad}, que reasigna tokens desbordados a expertos subutilizados, equilibrando la distribución de tokens. Estas técnicas optimizan colectivamente la utilización tanto de expertos con alta carga como con baja carga, lo que resulta en una canalización de inferencia de MoE más eficiente. Experimentos extensos demuestran la efectividad de nuestros métodos, mostrando mejoras significativas en la eficiencia de inferencia, por ejemplo, un aumento promedio del 0.2\% en el rendimiento y una aceleración de 1.94 veces en la inferencia en 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.