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

TL;DR

Piper is a framework that improves large-scale MoE model training efficiency on HPC platforms by using resource modeling and pipelined hybrid parallelism, leading to significant performance gains.

Contribution

It introduces a resource-aware training strategy with pipeline parallelism for MoE models, addressing memory, communication, and workload imbalance challenges.

Findings

Piper achieves 2-3.5X higher MFU than X-MoE.

A new all-to-all algorithm delivers 1.2-9X bandwidth improvements.

The framework effectively mitigates communication and workload bottlenecks.

Abstract

Frontier models increasingly adopt Mixture-of-Experts (MoE) architectures to achieve large-model performance at reduced cost. However, training MoE models on HPC platforms is hindered by large memory footprints, frequent large-scale communication across heterogeneous networks, and severe workload imbalance. To characterize these challenges, we develop a mathematical model that quantifies memory, compute, and communication requirements for MoE configurations under various parallelization schemes, verified through micro-benchmarking, code instrumentation, and hardware profiling. Our analysis identifies performance bottlenecks: all-to-all latency at scale from expert parallelism, insufficient compute-communication overlap, low GPU utilization from imbalanced skinny GEMMs, and the absence of platform-aware hybrid parallelization strategies. To address these, we introduce Piper, a framework that leverages resource modeling to identify efficient training strategies for MoE models on target HPC platforms, applying pipeline parallelism with optimized schedules. Piper achieves 2-3.5X higher MFU than state-of-the-art frameworks such as X-MoE, and a novel all-to-all algorithm delivers 1.2-9X bandwidth over vendor implementation.