Skip to content

feat(training): support MFSDP V2 expert parallelism - #5521

Open
wujingyue wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
wujingyue:agent/mfsdp-v2-ep
Open

feat(training): support MFSDP V2 expert parallelism#5521
wujingyue wants to merge 3 commits into
NVIDIA-NeMo:mainfrom
wujingyue:agent/mfsdp-v2-ep

Conversation

@wujingyue

@wujingyue wujingyue commented Aug 11, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Summary

  • enable the guarded MFSDP V2 MoE path for TP=PP=CP=1 with EP>1
  • add a small, two-GPU generic MoE EP=2 smoke exercising MFSDP V2 training

Upstream dependency

This PR depends on the MCore adapter changes in NVIDIA/Megatron-LM#6450.

Tracking

Validation

  • uv run pre-commit run --all-files
  • focused two-GPU test_fsdp_v2_moe_ep2_pretrain_smoke

@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.

@wujingyue
wujingyue force-pushed the agent/mfsdp-v2-ep branch 6 times, most recently from 5c85c39 to 6a3ab12 Compare August 12, 2026 02:51
@wujingyue
wujingyue marked this pull request as ready for review August 12, 2026 02:52
@wujingyue
wujingyue requested review from cuichenx and yaoyu-33 August 12, 2026 02:52
wujingyue
wujingyue commented Aug 12, 2026
View reviewed changes
Comment thread 3rdparty/Megatron-LM
@yaoyu-33 yaoyu-33 added area:training Training loop, callbacks, and runtime integration blocked Work cannot move forward until an external dependency is cleared feature New capabilities, enhancements, or enablement work full-test-suite high-complexity Harder to merge: prone to conflicts and needs additional test coverage needs-more-tests Requires additional L0 and L1 test coverage before merge labels Aug 12, 2026
@wujingyue wujingyue mentioned this pull request Aug 12, 2026
Draft
cuichenx
cuichenx previously approved these changes Aug 13, 2026
View reviewed changes

@cuichenx cuichenx left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please check the test failure, otherwise LGTM

@wujingyue
wujingyue enabled auto-merge (squash) August 13, 2026 21:16
@wujingyue
feat(training): support MFSDP V2 expert parallelism
974afe6 
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@wujingyue
fix(test): inherit MLA config in MFSDP V2 smoke
2523f54 
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@wujingyue
test(fsdp): simplify MFSDP V2 EP smoke
bc1d14f 
Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>
@wujingyue

wujingyue commented Aug 24, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test bc1d14f

@wujingyue
wujingyue requested a review from cuichenx August 24, 2026 23:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yaoyu-33 yaoyu-33 Awaiting requested review from yaoyu-33
@cuichenx cuichenx Awaiting requested review from cuichenx

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

area:training Training loop, callbacks, and runtime integration blocked Work cannot move forward until an external dependency is cleared feature New capabilities, enhancements, or enablement work full-test-suite high-complexity Harder to merge: prone to conflicts and needs additional test coverage needs-more-tests Requires additional L0 and L1 test coverage before merge

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@wujingyue @cuichenx @yaoyu-33