NVIDIA / NVIDIA/Megatron-LM

Implement Expert Swaping and MOE-TokenDedupe for balancing training load and minimizing comms.

Open
#1,839 1 comment 0 reactions 0 assignees View on GitHub
community-request enhancement module: moe
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.