Piper: Efficient Large-Scale MoE Training via Resource Modeling and Pipelined Hybrid Parallelism

Problem
This paper addresses the challenges associated with training Mixture-of-Experts (MoE) models on high-performance computing (HPC) platforms, particularly focusing on the inefficiencies stemming from large memory footprints, extensive communication overhead, and workload imbalances. The authors highlight that existing frameworks struggle to optimize these factors, which limits the scalability and performance of MoE architectures. This work is presented as a preprint and has not yet undergone peer review.

Method
The authors propose a framework named Piper, which employs resource modeling to optimize the training of MoE models. Piper quantifies the memory, compute, and communication requirements for various MoE configurations through a mathematical model, validated via micro-benchmarking and hardware profiling. The framework incorporates pipeline parallelism with optimized scheduling to enhance efficiency. Key innovations include a novel all-to-all communication algorithm that significantly improves bandwidth and a strategy to mitigate the performance bottlenecks identified, such as all-to-all latency and low GPU utilization due to imbalanced computations. The authors do not disclose specific details regarding the architecture of the MoE models used or the exact compute resources required for training.

Results
Piper demonstrates a substantial improvement in model training efficiency, achieving 2-3.5 times higher Model Forward Utilization (MFU) compared to state-of-the-art frameworks like X-MoE. Additionally, the proposed all-to-all communication algorithm yields a bandwidth increase of 1.2-9 times over the vendor’s implementation. These results indicate a significant enhancement in both computational efficiency and communication throughput, which are critical for scaling MoE models effectively on HPC platforms.

Limitations
The authors acknowledge several limitations, including the potential for the resource modeling to be overly specific to certain HPC configurations, which may limit generalizability. They also note that while Piper improves efficiency, it does not eliminate all sources of latency, particularly in scenarios with extreme workload imbalances. Furthermore, the paper does not explore the implications of these optimizations on the final model performance in terms of accuracy or other downstream tasks, which could be critical for practical applications.

Why it matters
The implications of this work are significant for the future of large-scale model training, particularly in the context of MoE architectures, which are becoming increasingly prevalent in state-of-the-art AI systems. By addressing the inefficiencies in training MoE models, Piper could facilitate the deployment of larger and more complex models on HPC platforms, ultimately leading to advancements in various AI applications. This research lays the groundwork for further exploration into hybrid parallelization strategies and resource-aware training methodologies, which could enhance the scalability and performance of future AI systems.

Authors: Sajal Dash, Feiyi Wang
Source: arXiv:2605.05049
URL: https://arxiv.org/abs/2605.05049v1