Skip to content

feat(fsdp): support expert parallelism in MFSDP v2 adapter - #6450

Merged
wujingyue merged 9 commits into
NVIDIA:mainfrom
wujingyue:agent/mfsdp-v2-ep
Aug 13, 2026
Merged

feat(fsdp): support expert parallelism in MFSDP v2 adapter#6450
wujingyue merged 9 commits into
NVIDIA:mainfrom
wujingyue:agent/mfsdp-v2-ep

Conversation

@wujingyue

@wujingyue wujingyue commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Summary

  • shard MoE expert parameters on the expert-DP mesh in the MFSDP v2 MCore adapter
  • apply the EP gradient divisor before expert-DP averaging
  • add an EP=2 training integration test against a full-batch EP=1 reference, including loss-curve logging. This is similar to test_expert_parallel.py but on the MCore adapter level.

Addresses #5656

Testing

  • python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/mfsdp_v2/test_mcore_adapter.py::TestMcoreAdapterExpertParallel::test_build_train_and_step

@copy-pr-bot

copy-pr-bot Bot commented Aug 11, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@copy-pr-bot

copy-pr-bot Bot commented Aug 11, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@svcnvidia-nemo-ci svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Aug 12, 2026
@wujingyue
wujingyue enabled auto-merge August 12, 2026 17:57
@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

❌ Claude fix stopped because exact-SHA CI did not complete in time. View exact-SHA CI.

svcnvidia-nemo-ci and others added 2 commits August 13, 2026 05:49
Signed-off-by: svcnvidia-nemo-ci <svcnvidia-nemo-ci@nvidia.com>
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31684995835

@svcnvidia-nemo-ci

Copy link
Copy Markdown
Contributor

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31685596036

Merged via the queue into NVIDIA:main with commit 281075d Aug 13, 2026
185 checks passed
@wujingyue
wujingyue deleted the agent/mfsdp-v2-ep branch August 13, 2026 12:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Approved All necessary approvals have been made complexity: low Run tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants