Инференция с учетом пропускной способности: смягчение эффекта отстающих в смеси экспертов

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

March 7, 2025

Авторы: Shwai He, Weilin Cai, Jiayi Huang, Ang Li

cs.AI

Аннотация

Смесь экспертов (Mixture of Experts, MoE) представляет собой эффективную архитектуру для масштабирования больших языковых моделей за счёт использования разреженной активации экспертов, что позволяет оптимизировать баланс между производительностью и эффективностью. Однако при использовании параллелизма экспертов MoE сталкивается с проблемами неэффективности вывода из-за несбалансированного распределения токенов между экспертами, когда некоторые эксперты перегружены, а другие остаются недоиспользованными. Этот дисбаланс приводит к низкой утилизации ресурсов и увеличению задержек, поскольку наиболее загруженный эксперт определяет общую задержку — явление, которое мы определяем как \textit{Эффект Отстающего}. Для смягчения этой проблемы мы предлагаем подход \textit{Вывод с учётом ёмкости}, включающий два ключевых метода: (1) \textit{Отбрасывание токенов с учётом ёмкости}, которое удаляет перегруженные токены для регулирования максимальной задержки MoE, и (2) \textit{Перенаправление токенов с учётом ёмкости}, которое перераспределяет избыточные токены на недоиспользованных экспертов, балансируя распределение токенов. Эти методы совместно оптимизируют использование как высоконагруженных, так и малонагруженных экспертов, что приводит к более эффективному процессу вывода в MoE. Эксперименты демонстрируют эффективность наших методов, показывая значительное улучшение эффективности вывода, например, увеличение средней производительности на 0.2\% и ускорение вывода в 1.94 раза на модели 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.