Efficient Mixture-of-Experts Training with Pipelining Optimization

Sparsely activated Mixture-of-Experts (MoE) has emerged as a key technique to expand the size of Transformer-based large language models (LLMs) while maintaining low computational costs. However, MoE layers require to route the input data to distributed devices, incurring significant communication latency. Existing studies have primarily focused on alleviating this problem by overlapping computation and communication tasks within a single MoE layer, which fails to achieve sufficient overlap and results in limited performance gains. In this work, we introduce an orthogonal partitioning dimension from existing task-parallel methods by leveraging the autoregressive nature of causal Transformer-based LLMs, i.e. partitioning tasks along the sequence dimension. This provides more flexible and efficient overlaps among tasks from both non-MoE and MoE layers. To this end, we propose an efficient MoE training approach, SP-MoE, with two innovative designs. 1) It incorporates non- MoE layers into the overlapping with not only the current MoE layer but also the preceding MoE layer, thereby facilitating more efficient training; 2) It identifies the optimal combination of pipeline degrees for non-MoE and MoE layers and devises the best scheduling plans for load-imbalanced non-MoE and uniform MoE layers to achieve the goal of minimizing the total training latency. Furthermore, we extend our algorithm to incorporate pipeline parallelism (PP) for training larger MoE models, which introduces peer-to-peer (P2P) communications between adjacent PP stages. To mitigate this problem, we propose S3P-MoE with two optimizations: 1) it partitions P2P communication into chunks and each P2P chunk is transmitted immediately to be overlapped with expert computation; 2) it allows the subsequent PP stage to commence upon receiving the corresponding P2P chunk. Moreover, we refine our algorithm to jointly optimize pipeline degrees for both MoE and non-MoE layers, as well as P2P communication, to minimize overall training latency. Extensive experiments demonstrate that SP-MoE and S3P-MoE effectively identify optimal pipeline degrees, reducing training time by up to 34:3% and 30:61%, respectively, compared to stateof- the-art methods.