Mixture-of-experts models run faster with expert coactivation improvements
Efficient Mixture-of-Experts with Speculative Decoding via Expert Coactivation
Machine LearningComputation and Language
Summary
Large language models can get slower because they use many expert parts that need to be loaded from memory during processing. The authors found that changing how the model chooses which expert parts to work together, a process called expert coactivation, can speed up the model’s response time. By training the model with specific methods that encourage experts to activate together, the system runs about 21% faster without losing accuracy. This helps reduce the memory transfer bottleneck and makes the overall process more efficient.
What this means in practice
- •For machine learning engineers: Enhance inference speed of large Mixture-of-Experts models by redesigning expert routing during training for higher coactivation.
- •For cloud service operators: Run large language models with Mixture-of-Experts architectures more efficiently to reduce latency and improve throughput.
Authors
Kumari Nishu, Han-Byul Kim, Santosh Chilkunda, Maxwell Horton, Arnav Kundu, Mohammad Samragh, Lauren Hannah, Mohammad Sekhavat, Nikhil Bhendawade, Manuel Ciosici, Iman Mirzadeh, Keivan Alizadeh Vahid, David Harrison, Irina Belousova, Mehrdad Farajtabar, Minsik Cho
Abstract
Mixture-of-Experts (MoE) models are increasingly deployed alongside Speculative Decoding (SD) to accelerate inference, but combining the two is challenging. SD improves the inference speed of dense models by verifying groups of tokens in parallel. However, the inference speedup for SD with MoEs depends heavily on the number of tokens being verified. Using more verification tokens results in more experts being transferred from DRAM to the Neural Processing Unit (NPU), which increases the memory transfer cost. This negatively impacts model runtime, as memory transfer is typically the bottleneck in inference. In this work, we investigate the impact of MoE router design during training on the speed of MoEs with SD. We find that routers with high degrees of expert coactivation result in much faster runtimes, mitigating the impact of using more verification tokens. Motivated by this observation, we assess the impact of various router design choices on expert coactivation and runtime using billion-parameter transformer models. We find that combining a global load-balancing loss, shared experts, a consistency loss, and an autoregressive expert selection mechanism during training results in significantly stronger expert coactivation. This increased coactivation translates into higher overall runtime throughput: our exploration yields a model that improves throughput by 21% over MoE baselines, while maintaining on-par accuracy with the baseline MoE.