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

arXiv: https://arxiv.org/abs/2411.10003 · HTML: https://mlsanity.com/paper/2411.10003/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Placed in Efficient MoE training and inference systems, one of 12 directions in Mixture-of-experts.
- Builds on 21 of its 60 references that are in the map.
- First posted on arXiv 15 Nov 2024.

## Placed in

- Mixture-of-experts / [Efficient MoE training and inference systems](https://mlsanity.com/moe/efficient-moe-training-and-inference-systems/)

## Builds on (most cited first, 21 of 21)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel](https://mlsanity.com/paper/2304.11277/) (2023)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [ST-MoE: Designing Stable and Transferable Sparse Expert Models](https://mlsanity.com/paper/2202.08906/) (2022)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints](https://mlsanity.com/paper/2212.05055/) (2022)
- [Sequence Parallelism: Long Sequence Training from System Perspective](https://mlsanity.com/paper/2105.13120/) (2021)
- [Taming Sparsely Activated Transformer with Stochastic Experts](https://mlsanity.com/paper/2110.04260/) (2021)
- [FastMoE: A Fast Mixture-of-Expert Training System](https://mlsanity.com/paper/2103.13262/) (2021)
- [Mixture of A Million Experts](https://mlsanity.com/paper/2407.04153/) (2024)
- [HetuMoE: An Efficient Trillion-scale Mixture-of-Expert Distributed Training System](https://mlsanity.com/paper/2203.14685/) (2022)
- [One Student Knows All Experts Know: From Sparse to Dense](https://mlsanity.com/paper/2201.10890/) (2022)
- [PipeTransformer: Automated Elastic Pipelining for Distributed Training of Transformers](https://mlsanity.com/paper/2102.03161/) (2021)
- [Maximizing Parallelism in Distributed Training for Huge Neural Networks](https://mlsanity.com/paper/2105.14450/) (2021)
- [MiCS: Near-linear Scaling for Training Gigantic Model on Public Cloud](https://mlsanity.com/paper/2205.00119/) (2022)

## Built on by (most cited first, 4 of 4)

- [STEM: Scaling Transformers with Embedding Modules](https://mlsanity.com/paper/2601.10639/) (2026)
- [Input Domain Aware MoE: Decoupling Routing Decisions from Task Optimization in Mixture of Experts](https://mlsanity.com/paper/2510.16448/) (2025)
- [Fine-grained MoE Load Balancing with Linear Programming](https://mlsanity.com/paper/2511.16947/) (2025)
- [Function Over Form: Distributional Orthogonalization in Mixture-of-Experts with Replica Expert Mechanism](https://mlsanity.com/paper/2609.32398/) (2026)

## Abstract (arXiv)

The size of deep learning models has been increasing to enhance model quality. The linear increase in training computation budget with model size means that training an extremely large-scale model is exceedingly time-consuming. Recently, the Mixture of Expert (MoE) has drawn significant attention as it can scale models to extra-large sizes with a 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 overlapping. 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 profile training statistics and use them to design Pro-Prophet. For lower communication volume, 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, Pro-Prophet scheduler schedules the data-dependent operations based on the statistics and operation features, further improving the training throughput. Experimental results indicate that Pro-Prophet achieves up to 2.66x speedup compared to Deepspeed-MoE and FasterMoE. Additionally, Pro-Prophet achieves a load-balancing enhancement of up to 11.01x when compared to FasterMoE.