SP-MoE: Expediting Mixture-of-Experts Training with Optimized Pipelining Planning

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. Extensive experiments conducted on two GPU clusters demonstrate that SP-MoE can effectively identify the optimal combination of pipeline degrees and achieve 16.1% - 34.3% reduction in training latency compared to three state-of-the-art MoE systems.