Inferência com Consciência de Capacidade: Mitigando o Efeito Straggler em Mistura de Especialistas

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

Resumo

A Mistura de Especialistas (MoE, do inglês Mixture of Experts) é uma arquitetura eficaz para escalar modelos de linguagem de grande porte ao aproveitar a ativação esparsa de especialistas, otimizando o equilíbrio entre desempenho e eficiência. No entanto, sob o paralelismo de especialistas, o MoE sofre com ineficiências de inferência devido à atribuição desequilibrada de tokens para especialistas, onde alguns especialistas ficam sobrecarregados enquanto outros permanecem subutilizados. Esse desequilíbrio leva a uma má utilização de recursos e a um aumento da latência, uma vez que o especialista mais sobrecarregado determina o atraso geral, um fenômeno que definimos como o \textit{Efeito Straggler}. Para mitigar isso, propomos a Inferência Consciente da Capacidade, que inclui duas técnicas principais: (1) o \textit{Descarte de Tokens Consciente da Capacidade}, que descarta tokens sobrecarregados para regular a latência máxima do MoE, e (2) o \textit{Redirecionamento de Tokens Consciente da Capacidade}, que realoca tokens excedentes para especialistas subutilizados, equilibrando a distribuição de tokens. Essas técnicas, em conjunto, otimizam a utilização tanto de especialistas com alta carga quanto com baixa carga, resultando em um pipeline de inferência do MoE mais eficiente. Experimentos extensivos demonstram a eficácia de nossos métodos, mostrando melhorias significativas na eficiência de inferência, como um aumento médio de desempenho de 0,2\% e uma aceleração de inferência de 1,94 vezes no 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.