OmniPipe: Efficient, Flexible and Scalable Pipeline Parallelism for Large Model Training

Abstract

Mixture-of-Experts (MoE) has become the de facto architecture for scaling large language models, offering expanded capacity with manageable compute. Pipeline parallelism (PP) is indispensable for distributed MoE training, but state-of-the-art PP schemes face three major limitations: large pipeline bubbles, high per-stage latency due to insufficient overlap of Expert Parallelism (EP) communication, and limited flexibility due to rigid configuration constraints. As a result, training efficiency degrades severely in large-scale clusters. We propose OmniPipe, a flexible bidirectional multi-pipeline parallelism scheme for unified dense and MoE LLM training. It relaxes the rigid constraints of prior bidirectional pipelines by supporting configurable pipeline replica counts and flexible micro-batch scaling. A flexible scheduling strategy further achieves better overlap between EP All-to-All communication and computation, effectively minimizing pipeline bubbles and intra-stage latency to reduce end-to-end pipeline execution time. Implemented within Megatron-LM, OmniPipe enables full 5D hybrid parallelism. On NVIDIA A800 GPU clusters, OmniPipe consistently outperforms the best configurations of existing PP schemes, achieving a geometric-mean 1.12 × speedup on MoE workloads with up to 1.30 × , while also achieving a 1.10 × average speedup on dense models. The results demonstrate that OmniPipe minimizes the pipeline bubble ratio while effectively overlapping EP communication with computation, enabled by the flexible and scalable parallelism scheme of bidirectional pipelines.