🤖 AI Summary
This work addresses the inefficiency and device idling in distributed Mixture-of-Experts (MoE) inference caused by imbalanced expert loads. The authors propose a cross-layer load balancing strategy that leverages the inherent redundancy among experts across different MoE layers. Without altering expert-to-device mappings or requiring expert replication or migration, the method achieves dynamic load balancing with zero additional overhead by jointly scheduling computation tasks across layers. A greedy algorithm selects the subset of subtasks to execute at each step while deferring others to exploit future balancing opportunities, seamlessly integrating into existing distributed inference frameworks. Experiments demonstrate that the approach consistently improves inference efficiency across diverse models, tasks, and system configurations, reducing GPU idle time by over 40% on average.
📝 Abstract
Load Balancing has emerged as a critical problem in expert-parallel distributed inference of Mixture-of-Experts (MoE) models. As routing distributions are typically skewed across experts, devices hosting lighter-loaded experts must idle to wait for the heaviest during expert computing, leading to inefficiency. Existing load-balancing approaches primarily rely on expert replication or migration within each layer, which introduce additional overhead and limit their flexibility and scalability. To address this problem, we propose EasyBalance, a cross-layer load balancing strategy that requires no modifications to the expert-device mapping, enabling instant adaptability and incurring essentially no additional overhead. Our key insights are that (1) experts of other layers can be viewed as naturally redundant for the current layer, and (2) cross-layer MoE workloads can be jointly executed to mitigate their individual imbalance. Based on these observations, EasyBalance greedily schedules a subset of cross-layer workloads to run at each MoE step and defers the remaining workloads for future balancing opportunities, effectively leveraging cross-layer imbalance mitigation. Extensive experiments across models, tasks, and configurations demonstrate that EasyBalance consistently accelerates distributed MoE inference, reducing GPU idling by mostly over 40%. Code is available at https://github.com/yize-wu/EasyInfra.