Mixture-of-Parallelisms: Towards Memory-Efficient Training Stack for Mixture-of-Experts Models
This work addresses the memory and throughput bottlenecks in training large-scale MoE models, enabling practitioners to train trillion-parameter models with long contexts on limited hardware.
The paper presents a memory-efficient training stack for Mixture-of-Experts (MoE) models that combines parallelism techniques to achieve high throughput and memory efficiency, enabling training of trillion-parameter models with million-token context lengths on fewer than 12 8x H200 GPU nodes, delivering 4.7x–8.2x higher per-GPU throughput than a strong FSDP2 baseline and sustaining training at context lengths up to 1M tokens where the baseline fails beyond 64–128K.
This paper showcases a memory-efficient training stack for Mixture-of-Experts (MoE) models. It is a training paradigm that combines and specializes various existing and novel parallelism techniques at different layers and stages of the Mixture-of-Experts (MoE) model training pipeline. It leverages these techniques to achieve maximal efficiency given the physical constraints of CPU, CPU memory, GPU HBM memory, and the CPU-GPU, GPU-GPU, and node-node communication bandwidth of the GPU cluster. It also contains a novel strategy for the optimizer step to achieve high throughput and memory efficiency, enabling practitioners to conduct lossless pre-training/fine-tuning of trillion-parameter scale models, at a million context length, with just under 12 8x H200 GPU nodes, with state-of-the-art throughput and memory efficiency. In our experiments, MoP delivers 4.7x--8.2x higher per-GPU throughput than a strongly-tuned FSDP2 baseline (with the gap widening at larger scale) and sustains training at context lengths up to 1M tokens, where the baseline runs out of memory beyond 64--128K.