Feature Request: ScatterMoE (Triton-based Sparse MoE with fused scatter/gather GEMMs)
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Summary
Request to explore [ScatterMoE](https://arxiv.org/abs/2403.08245) techniques in Megatron Core's MoE implementation. ScatterMoE fuses grouped GEMMs with scattered read/write operations via a Triton kernel, eliminating the need to materialize padded tensors in HBM.
**Update:**
While ScatterMoE provides the foundational approach, SonicMoE further optimizes for Hopper/Blackwell architectures by reducing activation memory and improving tile utilization. See https://github.com/NVIDIA/Megatron-LM/issues/2709
## Motivation
Padded MoE implementations copy inputs into padded tensors to handle variable-length expert assignments. This overhead grows with expert count and granularity, and is amplified during training where intermediates are retained for the backward pass.
ScatterMoE pads **indices** rather than **tensors**. A `scatter2scatter` Triton kernel loads tiles using padded indices directly into SRAM. Reported results (8×A100, 1.5B model):
- ~38% higher training throughput vs MegaBlocks
- ~34% lower training memory, ~46% lower inference memory
- Correctness validated via Mixtral 8x7B conversion (≤0.006 error across 11 benchmarks)
## Requested Feature
Investigate adding a ScatterMoE-style backend as a configurable option in `megatron.core.transformer.moe`, compatible with existing routers and expert parallelism.
## References
- [ScatterMoE Paper (Tan, Shen, Panda, Courville)](https://arxiv.org/abs/2403.08245)
- [GitHub: shawntan/scattermoe](https://github.com/shawntan/scattermoe)
Contributor guide
Assessment
This issue has not been assessed yet.