Pro-Prophet: A Systematic Load Balancing Method for Efficient Parallel Training of Large-Scale MoE Models

The size of deep learning models has been increasing to enhance model quality. The linear increase in training computation budgets with model size means that training an extremely large-scale model is exceedingly time-consuming. Recently, the Mixture of Experts (MoE) has drawn significant attention as it can scale models to extra-large sizes with a near-stable computation budget. However, inefficient distributed training of large-scale MoE models hinders their broader application. Specifically, a considerable dynamic load imbalance occurs among devices during training, significantly reducing throughput. Several load-balancing works have been proposed to address the challenge. System-level solutions draw more attention for their hardware affinity and non-disruption of model convergence compared to algorithm-level ones. However, they are troubled by high communication costs and poor communication-computation overlap. To address these challenges, we propose a systematic load-balancing method, Pro-Prophet, which consists of a planner and a scheduler for efficient parallel training of large-scale MoE models. To adapt to the dynamic load imbalance, we have profiled training statistics and utilized them to design Pro-Prophet. For lower communication volume, the Pro-Prophet planner determines a series of lightweight load-balancing strategies and efficiently searches for a communication-efficient one for training based on the statistics. For sufficient overlapping of communication and computation, the Pro-Prophet scheduler schedules the data-dependent operations based on the statistics and operation characteristics, further improving the training throughput. We conduct extensive experiments in various clusters and MoE models. The results indicate that Pro-Prophet achieves up to 2.66x speedup on MoE-GPT models compared to two popular MoE frameworks, namely Deepspeed-MoE and FasterMoE. Furthermore, Pro-Prophet has demonstrated a load-balancing improvement of up to 11.01x and speedups on modern MoE models up to 1.22x compared to a representative load-balancing work, FasterMoE.