Research map / Architectures and efficiency / Mixture of experts
Scaling Beyond the GPU Memory Limit for Large Mixture-of-Experts Model Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes ES-MoE, which offloads MoE expert parameters and optimizer states to host memory and SSD with expert-wise pipelined processing, CPU optimization, and dynamic expert placement to train large MoE models on few GPUs.
Problem
Training large Mixture-of-Experts models requires adding GPUs as the number of experts grows, and token load imbalance across experts causes zero-padding waste and straggler GPUs.
Approach
Offload expert parameters and optimizer states per expert to CPU memory/SSD, schedule expert uploads in a pipeline overlapped with token permutation and computation, run CPU optimizer per expert as soon as its backward pass finishes, and dynamically place experts on GPUs by load using greedy makespan scheduling, avoiding batched matmul and zero-padding.
Key results
- ES-MoE trains a 29.3B-parameter MoE-L with 64 experts on 4 A100 GPUs where FairSeq, Tutel and ZeRO-OffloadE hit OOM.
- ES-MoE achieves up to 17.5x throughput over ZeRO-OffloadE and up to 3.16x (2.13x in the conclusion) over MoE-specialized frameworks.
- With 4 TB SSD, ES-MoE scales up to 67x more experts than baselines; fine-tuning Fairseq-MoE-15B on SST-2 reaches 88% accuracy in 6.5 hours on 4 GPUs.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add expert pinning (ablation) | ES-MoE · MoE-L, 16 experts | training throughput: 19,501 (-3.8%) | Helped |
| add optimizer overlapping (ablation) | ES-MoE · MoE-L, 16 experts | training throughput: 17,943 (-8.7%) | Helped |
| use larger batch size (ablation) | ES-MoE · MoE-L, 16 experts | training throughput: 5,959 (-301%) | Helped |
| eliminate zero padding (ablation) | ES-MoE · MoE-L, 16 experts | training throughput: 4,674 (-27.4%) | Helped |
| fine-tune Fairseq-MoE-15B using ES-MoE on 4 GPUs | Fairseq-MoE-15B · SST-2 | accuracy | Helped |
| use ES-MoE instead of Fairseq for MoE-L pretraining | MoE-L, 8 experts · 1B tokens | training duration: 9.47 -> 6.90 | Helped |
Surprising findings
- Expert-wise CPU optimization without delayed updates keeps mathematical equivalence, unlike ZeRO-Offload's delayed update which introduces staleness.
Limitations
- Throughput drops as the number of tokens per expert decreases with more experts, reducing computational efficiency.
- Relies on CPU optimizer which is much slower than GPU optimizer (about 31x slower with Adam), so depends on host resources.
- Evaluated only on GPT-derived MoE models with top-1 gating on a single 4xA100 node.
Experimental setup
Datasets: WikiText-103 SST-2 MNLI BoolQ
Benchmarks: WikiText-103 SST-2 MNLI BoolQ
Metrics: training throughput (words/s) maximum number of experts training loss validation loss GPU utilization fine-tuned accuracy
Baselines: FairSeq Tutel ZeRO-OffloadE MegaBlocks
Builds on: ZeRO-Offload Tutel MegaBlocks GShard Switch Transformer
Closest papers by meaning
- MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts (ICLR 2025 · similarity 0.92)
- Joint MoE Scaling Laws: Mixture of Experts Can Be Memory Efficient (ICML 2025 · similarity 0.91)
- Mixture of Lookup Experts (ICML 2025 · similarity 0.90)
- LSH-MoE: Communication-efficient MoE Training via Locality-Sensitive Hashing (NeurIPS 2024 · similarity 0.90)
- Advancing Expert Specialization for Better MoE (NeurIPS 2025 · similarity 0.90)
- Pushing Mixture of Experts to the Limit: Extremely Parameter Efficient MoE for Instruction Tuning (ICLR 2024 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.