Mitigating Contention in Stream Multiprocessors for Pipelined Mixture of Experts: An SM-Aware Scheduling Approach

Sparsely activated Mixture-of-Experts (MoEs) models have become prominent in Large Language Models (LLMs) due to their ability to expand model capacity without proportional increases in computation. MoE layers feature multiple experts, with only a few activated per sample, enhancing model performance across various domains such as natural language generation and translation. The dynamic activation of MoE experts introduces extensive communications in distributed training. However, this dynamic activation creates communication challenges in distributed training. While previous work attempted to pipeline computation and communication through input chunking, we found that these tasks compete for Stream Multiprocessors (SMs) on GPUs, making the scheduling ineffective. In this paper, we update the optimization problem to minimize training time while accounting for SM contentions. We develop performance models for computation and communication tasks to identify MoE layer bottlenecks. By delaying GEMM launching and splitting GEMM operations, we enable communication to preempt SMs, enhancing overall efficiency and more stability.