Research map / Architectures and efficiency / Mixture of experts
Occult: Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
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.
Occult reduces all-to-all communication in expert-parallel MoE training and inference by optimizing expert placement from co-activation profiles, deduplicating token replicas via a custom sparse kernel, and optionally pruning cross-device collaborations.
Problem
All-to-all communication in expert-parallel MoE training and inference consumes a large share of runtime (over 40% in large-scale training), limiting scalability.
Approach
Define collaborative communication between co-activated expert pairs (intra- vs inter-device), reschedule expert placement using profiled co-activation statistics, build a bidirectional re-index matrix (BRIM) and Triton sparse matmul kernels that send one replica per token per device, and prune routing to at most N_d devices (routing-score or expert-similarity based) with modified fine-tuning.
Key results
- On prefilling with 2^14 tokens, Occult (two-device pruning) is 8.66x faster than Tutel and 1.70x faster than MegaBlocks on OLMoE.
- For batch size 8 training, Occult (two-device pruning) achieves 1.54x speedup over MegaBlocks, 2.65x over FSDP, ~9x over DeepSpeed, and ~20x over Tutel.
- Two-device collaboration pruning matches or exceeds standard SFT quality on several benchmarks, while single-device pruning degrades quality.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Reschedule expert placement using profiled collaboration | Occult · OLMoE, prefilling (2^14 prompt tokens) | latency (s): 15.93 -> 12.53 | Helped |
| Reduce collaboration pruning to 2 devices (ablation) | Occult · MoE fine-tuning and evaluated NLP tasks | model performance | Helped |
| Reduce collaboration pruning to 1 device (ablation) | Occult · MoE fine-tuning and evaluated NLP tasks | model performance | Hurt |
| Use expert-similarity-based rather than router-score-based pruning (ablation) | Occult collaboration pruning · NLP evaluation benchmarks | accuracy | Mixed |
| Prune collaboration to 1 GPU instead of 2 GPUs (ablation) | Occult · OLMoE, prefilling (2^14 prompt tokens) | latency (s): 9.22 -> 5.82 | Helped |
| Prune collaboration to 2 GPUs instead of trivial placement (ablation) | Occult · OLMoE, prefilling (2^14 prompt tokens) | latency (s): 15.93 -> 9.22 | Helped |
| Reschedule expert placement instead of trivial placement (ablation) | Occult · OLMoE, prefilling (2^14 prompt tokens) | latency (s): 15.93 -> 12.53 | Helped |
| Use Occult, 2-device pruning rather than MegaBlocks | MoE fine-tuning · Batch size 8 | speedup: 1.54× | Helped |
| Use Occult, 2-device pruning rather than MegaBlocks | OLMoE · Prefilling (2^14 prompt tokens) | speedup: 1.70× | Helped |
| Use Occult, 2-device pruning rather than MegaBlocks | Qwen-MoE · Prefilling (2^14 prompt tokens) | speedup: 0.71× | Hurt |
| Use Occult, 2-device pruning rather than MegaBlocks | DeepSeek-MoE · Prefilling (2^14 prompt tokens) | speedup: 0.66× | Hurt |
Surprising findings
- Similarity-based pruning consistently outperforms router-score-based pruning on the reported benchmarks.
- Two-device pruning sometimes yields accuracy above standard fine-tuning (e.g., WSC 86.45 vs 86.08 for OLMoE with router-based 2-GPU pruning).
Limitations
- Expert placement is fixed from a profiling dataset and does not address workload balancing, which the authors leave unsolved.
- Single-device pruning can hurt model quality, and quality gains depend on task and model.
- Experiments run on a single node with PCIe-connected A6000 GPUs, so generalization to NVLink or multi-node clusters is untested.
Experimental setup
Datasets: Alpaca Winogrande WSC PIQA RACE MathQA RTE
Benchmarks: Winogrande WSC PIQA RACE MathQA RTE MMLU HellaSwag BoolQ COQA MRPC SST-2
Metrics: accuracy F1 score exact match score latency throughput
Baselines: vLLM HuggingFace PyTorch FSDP DeepSpeed Tutel MegaBlocks
Builds on: MegaBlocks MC-MoE Triton
Closest papers by meaning
- MoE++: Accelerating Mixture-of-Experts Methods with Zero-Computation Experts (ICLR 2025 · similarity 0.89)
- Shortcut-connected Expert Parallelism for Accelerating Mixture of Experts (ICML 2025 · similarity 0.89)
- BigMac: A Communication-Efficient Mixture-of-Experts Model Structure for Fast Training and Inference (AAAI 2025 · similarity 0.89)
- Mixture of Lookup Experts (ICML 2025 · similarity 0.88)
- Scaling Beyond the GPU Memory Limit for Large Mixture-of-Experts Model Training (ICML 2024 · similarity 0.88)
- Unchosen Experts Can Contribute Too: Unleashing MoE Models’ Power by Self-Contrast (NeurIPS 2024 · 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.
