EasyBalance: Cross-Layer Load Balancing in Distributed MoE Inference
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 \textbf{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 from other layers can be viewed as naturally redundant, and (2) expert workloads of multiple layers can be jointly executed. Based on these observations, EasyBalance greedily schedules a subset of cross-layer workloads at each MoE-computation stage to run, while deferring the remaining workloads for future balancing opportunities. Extensive experiments across different models, tasks, and parallelism configurations demonstrate that EasyBalance consistently accelerates distributed MoE inference, reducing GPU idling by mostly over 40\%.
Lay Summary
Load Balancing is a critical problem when running Mixture-of-Experts (MoE) models on multiple GPUs. The problem is that the workloads of devices are uneven, and all the devices must wait for the slowest to finish. Existing load balancing methods must incur replication or migration of experts in the model, leading to additional overheads and limiting the flexibility (as expert migration is not flexible). This paper proposes EasyBalance, a cross-layer load balancing strategy. It solves the bottleneck by executing multiple sequences at different layers, and let the routing distributions of different layers compensate the imbalance with each other. Therefore, it enables instant adaptability across tasks and incurring essentially no additional overhead, providing strong flexibility and scalability. Experiments across different models, tasks, and parallelism configurations demonstrate that EasyBalance consistently accelerates expert-parallel inference.