Implement Expert Swaping and MOE-TokenDedupe for balancing training load and minimizing comms.
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 3h
- Merged PRs (30d)
- 272
Description
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
Often models like DeepSeek have very imbalanced token routing which can make balancing compute across nodes difficult. This has already been studied well by DeepSeek and their inference stack uses the https://github.com/deepseek-ai/EPLB to balance inference of their models at scale. However, the same can also be done at training time to provide lift to e2e training time. This paper describes https://arxiv.org/abs/2508.09591
how this trick along with a hierarchical a2a implementation, can accelerate DeepSeek training by 22%+. Another trick that can be done is deduplicating and reduplicating tokens in the hierarchical a2a.
**Describe the solution you'd like**
A clear and concise description of what you want to happen.
Implement the techniques described in the paper above. Expert swapping would significantly make training large model easier, with higher EP, especially when pretraining from scratch by making the training workload far more balanced across GPUs. TokenDeduplication would also help accelerate the A2A operations, but may be more difficult to implement into DeepEP.
**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.
**Additional context**
Add any other context or screenshots about the feature request here.
Contributor guide
Research direction
Start by reading the linked paper and the DeepSeek EPLB implementation referenced in the issue, then inspect the repository's DeepEP integration. The issue names no files or tests; done would require implementing expert swapping and MOE token deduplication for training and demonstrating improved load balance or communication performance.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 20/100