Research map / Architectures and efficiency / Mixture of experts
FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts 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 FlowMoE, a pipeline scheduling framework for distributed MoE training that jointly pipelines MHA, gating, expert computing, A2A communication, and chunked all-reduce with Bayesian-optimized chunk size, reporting 13%-57% training time reduction over prior MoE frameworks.
Problem
Distributed MoE training under expert parallelism leaves MHA computing, gating, and all-reduce communication unoverlapped, and existing pipeline schedulers only handle expert computation and A2A communication within the MoE layer.
Approach
Build a unified pipeline across all tasks of a transformer block (MHA, gating, expert compute, dispatch/combine A2A, all-reduce), and schedule all-reduce as priority-lowered tensor chunks inserted into gaps between A2A communications, with the chunk size auto-tuned by Bayesian optimization.
Key results
- FlowMoE reduces per-iteration time on 16-GPU Cluster 1 relative to ScheMoE by 14%-31% across four real-world MoE models (e.g., DeepSeek-V2-S 3205.3 ms vs 4093.7 ms).
- On 490 + 393 customized MoE layer cases, FlowMoE is on average 26% faster than ScheMoE.
- FlowMoE saves 10%-16% energy relative to ScheMoE and up to 7% memory relative to ScheMoE.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use FlowMoE's unified task pipeline and priority-based all-reduce chunk scheduling | distributed MoE training · real-world MoE models, 16 GPUs | iteration time: 1.13×-1.82× speedup | Helped |
| add MHA computing and gating pipelining (ablation) | FlowMoE with MoE/A2A pipelining · customized MoE layer, 16 GPUs | iteration time: 1115.2 -> 1012.6 | Helped |
| add all-reduce chunk pipelining with priority scheduling (ablation) | FlowMoE with MoE/A2A pipelining · customized MoE layer, 16 GPUs | iteration time: 1115.2 -> 895.3 | Helped |
| use Bayesian optimization instead of a fixed 1MB all-reduce chunk size (ablation) | FlowMoE-AR · customized MoE layer, 16 GPUs | iteration time: 971.5 -> 895.3 | Helped |
| add both MHA/gating and all-reduce pipelining (ablation) | FlowMoE · customized MoE layer, 16 GPUs | iteration time: 1115.2 -> 796.1 | Helped |
| increase pipelining degree from R=2 to R=4 (ablation) | FlowMoE · DeepSeek-V2-S, 16 GPUs | iteration time: 3205.3 -> 3113.8 | Helped |
| increase pipelining degree from R=4 to R=8 (ablation) | FlowMoE · DeepSeek-V2-S, 16 GPUs | iteration time: 3113.8 -> 3295.9 | Hurt |
| replace Bayesian optimization with grid search for chunk-size tuning (ablation) | FlowMoE · DeepSeek-V2-S, 16 GPUs | iteration time: 3205.3 -> 3498.8 | Hurt |
| replace Bayesian optimization with random chunk-size selection (ablation) | FlowMoE · DeepSeek-V2-S, 16 GPUs | iteration time: 3205.3 -> 3902.75 | Hurt |
| use a fixed chunk size instead of BO-selected chunk size (ablation) | FlowMoE · GPT2-Tiny-MoE, 16 GPUs | iteration time: 95.6 -> 104.7 | Hurt |
| use a fixed chunk size instead of BO-selected chunk size (ablation) | FlowMoE · BERT-Large-MoE, 16 GPUs | iteration time: 351.9 -> 362.2 | Hurt |
| use a fixed chunk size instead of BO-selected chunk size (ablation) | FlowMoE · LLaMA2-MoE, 16 GPUs | iteration time: 1124.0 -> 1167.9 | Hurt |
Surprising findings
- Tutel is faster than FasterMoE in most cases despite FasterMoE's dedicated design.
- In bandwidth-limited settings with more GPUs, pipeline scheduling benefits less and FlowMoE's advantage over other frameworks becomes less pronounced.
Limitations
- Advantage shrinks in bandwidth-limited scenarios as GPU count grows.
- Requires BO re-tuning of the all-reduce chunk size when hardware changes.
- Evaluated only on two small clusters (up to 16 RTX3090/2080Ti GPUs) with models trained only for iteration-time measurement.
Experimental setup
Datasets: OpenWebText wikitext-103 customized MoE layer configurations
Benchmarks: OpenWebText wikitext-103
Metrics: per-iteration time speedup energy consumption memory usage
Baselines: vanilla expert parallelism FasterMoE Tutel FSMoE ScheMoE
Builds on: Tutel PipeMoE FasterMoE ScheMoE FSMoE
Closest papers by meaning
- Shortcut-connected Expert Parallelism for Accelerating Mixture of Experts (ICML 2025 · similarity 0.89)
- Advancing Expert Specialization for Better MoE (NeurIPS 2025 · similarity 0.89)
- FlashMoE: Fast Distributed MoE in a Single Kernel (NeurIPS 2025 · similarity 0.89)
- Scaling Beyond the GPU Memory Limit for Large Mixture-of-Experts Model Training (ICML 2024 · similarity 0.88)
- MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts (ICLR 2025 · similarity 0.88)
- BigMac: A Communication-Efficient Mixture-of-Experts Model Structure for Fast Training and Inference (AAAI 2025 · similarity 0.88)
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.