feat(fsdp): expose HSDP and HFSDP through the MFSDP v2 adapter by wujingyue · Pull Request #6873 · NVIDIA/Megatron-LM

wujingyue · 2026-08-26T02:13:37Z

Summary

MFSDP v2's core already supports hybrid sharding, but the MCore adapter never built a
second data-parallel axis: num_distributed_optimizer_instances was ignored and every
parameter was sharded over the full DP domain. This wires it up.

Dense parameters take a ("dp_outer", "dp_shard") mesh whenever there is more than one
instance. Expert parameters deliberately do not — they stay ZeRO-3 over the whole
expert-DP domain, since sharding the large parameters maximally is worth more than the
intra-instance collectives hybrid sharding buys.

No HSDP/HFSDP special case

Each axis takes the placements of its own sharding strategy, per
mfsdp_design.md:

strategy	parameter	gradient	optimizer
`no_shard`	replicate	partial	replicate	DDP
`optim`	replicate	partial	shard	ZeRO-1
`optim_grads`	replicate	shard	shard	ZeRO-2
`optim_grads_params`	shard	shard	shard	ZeRO-3

So HSDP is no_shard outer over ZeRO-3 inner, and HFSDP is ZeRO-1 outer over ZeRO-3
inner — whose optimizer is then sharded on both axes, i.e. across the flattened DP
domain. ZeRO-2 and ZeRO-3 on the outer axis come for free.

num_distributed_optimizer_instances decides only whether an outer axis exists;
outer_dp_sharding_strategy decides what lives on it. Both strategy fields are now
validated in DistributedDataParallelConfig.__post_init__ rather than deep in the FSDP
code, since the names are config vocabulary.

Validation

On 4 GPUs, TestMcoreAdapterHybrid — 5 passed:

placements are (Replicate, Shard(0)) for HSDP and (Shard(0), Shard(0)) for HFSDP,
on a ("dp_outer", "dp_shard") mesh
both produce the same losses over three steps as a single data-parallel instance,
so splitting the domain leaves the math unchanged
MoE with EP=2 and two instances puts dense parameters on the hybrid mesh and experts
on ("expert_dp",), and nothing else — experts never acquire an outer axis

test_mcore_adapter.py is 11 passed; with test_param_and_grad_buffer.py, 108 passed /
4 skipped, since the new config validation runs on every DistributedDataParallelConfig.

🤖 Generated with Claude Code

copy-pr-bot · 2026-08-26T02:13:41Z

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.

wujingyue · 2026-08-26T04:53:57Z

/ok to test

Autumn1998 · 2026-08-26T05:36:19Z

+        )
+
+    inner_size = inner_group.size()
+    layout = torch.arange(dist.get_world_size()).reshape(outer_group.size(), inner_size).tolist()


I’m a bit confused🤔: why can we use the get_world_size() to compute the expected indexs? Are we assuming that world_size = inner_dp × outer_dp?
It would be helpful to understand how this work when TP, PP, or EP(maybe with #6798) is also present, or whether support for those configurations is outside the current scope.

Yes, the assumption is world_size == outer_dp × inner_dp. This mirrors v1: _get_hsdp_tp_mesh (:864) also reshapes torch.arange(world_size) (:869), just with extra ep/tp factors that reduce away when those are 1.

It's checked rather than assumed — this rank's row is compared against intra_dp_cp and its column against inter_dist_opt (:845, :850), so a mismatch raises instead of silently mis-mapping shards. TP/PP/CP above 1 are rejected in _validate_config (:658), so a config that would break the identity can't reach here.

copy-pr-bot · 2026-08-26T20:54:05Z

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 · 2026-08-26T20:54:36Z

/ok to test

wujingyue · 2026-08-26T20:57:05Z

/ok to test

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

nemo-automation-bot · 2026-08-31T00:52:31Z

🔄 Merge queue validation started!

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

…ware optimizer, HSDP, and optimizer CUDA graphs Megatron-LM landed dedicated MFSDP v2 support for these, each with its own commits and tests, after this validation function was written: - clip_grad: MFSDP v2 gradient clipping across dense and expert meshes (NVIDIA/Megatron-LM#6489) - use_precision_aware_optimizer: mixed-precision optimizers in MFSDP v2 (NVIDIA/Megatron-LM#6506) - num_distributed_optimizer_instances / outer_dp_sharding_strategy (HSDP): expose HSDP and HFSDP through the MFSDP v2 adapter (NVIDIA/Megatron-LM#6873); the adapter itself now enforces the correct paired constraint, so no replacement check is needed here - optimizer_cuda_graph: verified for MFSDP v2 (NVIDIA/Megatron-LM#6498) Verified end-to-end against Megatron-LM main + Megatron-Bridge with the DeepSeek-V3 MFSDP v2 + EP proxy config (clip_grad=1.0): 20/20 iterations, finite decreasing loss, no skipped or NaN iterations.

…ware optimizer, HSDP, and optimizer CUDA graphs Megatron-LM landed dedicated MFSDP v2 support for these, each with its own commits and tests, after this validation function was written: - clip_grad: MFSDP v2 gradient clipping across dense and expert meshes (NVIDIA/Megatron-LM#6489) - use_precision_aware_optimizer: mixed-precision optimizers in MFSDP v2 (NVIDIA/Megatron-LM#6506) - num_distributed_optimizer_instances / outer_dp_sharding_strategy (HSDP): expose HSDP and HFSDP through the MFSDP v2 adapter (NVIDIA/Megatron-LM#6873); the adapter itself now enforces the correct paired constraint, so no replacement check is needed here - optimizer_cuda_graph: verified for MFSDP v2 (NVIDIA/Megatron-LM#6498) Verified end-to-end against Megatron-LM main + Megatron-Bridge with the DeepSeek-V3 MFSDP v2 + EP proxy config (clip_grad=1.0): 20/20 iterations, finite decreasing loss, no skipped or NaN iterations. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

…A#6873) Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

…ware optimizer, HSDP, and optimizer CUDA graphs Megatron-LM landed dedicated MFSDP v2 support for these, each with its own commits and tests, after this validation function was written: - clip_grad: MFSDP v2 gradient clipping across dense and expert meshes (NVIDIA/Megatron-LM#6489) - use_precision_aware_optimizer: mixed-precision optimizers in MFSDP v2 (NVIDIA/Megatron-LM#6506) - num_distributed_optimizer_instances / outer_dp_sharding_strategy (HSDP): expose HSDP and HFSDP through the MFSDP v2 adapter (NVIDIA/Megatron-LM#6873); the adapter itself now enforces the correct paired constraint, so no replacement check is needed here - optimizer_cuda_graph: verified for MFSDP v2 (NVIDIA/Megatron-LM#6498) Verified end-to-end against Megatron-LM main + Megatron-Bridge with the DeepSeek-V3 MFSDP v2 + EP proxy config (clip_grad=1.0): 20/20 iterations, finite decreasing loss, no skipped or NaN iterations. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

wujingyue commented Aug 26, 2026

View reviewed changes

wujingyue marked this pull request as ready for review August 26, 2026 04:53

wujingyue requested review from a team as code owners August 26, 2026 04:53

copy-pr-bot Bot temporarily deployed to test August 26, 2026 04:54 Inactive

copy-pr-bot Bot temporarily deployed to public August 26, 2026 04:57 Inactive

github-actions Bot added the complexity: medium label Aug 26, 2026

Autumn1998 reviewed Aug 26, 2026

View reviewed changes

wujingyue requested a review from Autumn1998 August 26, 2026 17:07

wujingyue added the Run tests label Aug 26, 2026

wujingyue force-pushed the agent/mfsdp-v2-hsdp branch from ce2b895 to 3ed254e Compare August 26, 2026 20:54

copy-pr-bot Bot had a problem deploying to test August 26, 2026 20:55 Error

wujingyue force-pushed the agent/mfsdp-v2-hsdp branch from 3ed254e to ceb56c6 Compare August 26, 2026 20:56

copy-pr-bot Bot temporarily deployed to test August 26, 2026 21:01 Inactive

copy-pr-bot Bot temporarily deployed to public August 26, 2026 21:03 Inactive

Autumn1998 approved these changes Aug 27, 2026

View reviewed changes

wujingyue force-pushed the agent/mfsdp-v2-hsdp branch from ceb56c6 to dbb852b Compare August 27, 2026 22:45

copy-pr-bot Bot temporarily deployed to test August 27, 2026 22:46 Inactive

copy-pr-bot Bot temporarily deployed to public August 27, 2026 22:49 Inactive

copy-pr-bot Bot temporarily deployed to public August 30, 2026 06:51 Inactive

test(fsdp): sort MFSDP adapter test imports

f483003

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

copy-pr-bot Bot temporarily deployed to test August 30, 2026 14:47 Inactive

copy-pr-bot Bot temporarily deployed to public August 30, 2026 14:49 Inactive

wujingyue mentioned this pull request Aug 30, 2026

validate data-parallel sharding strategies #6983

Merged

wujingyue added 2 commits August 30, 2026 14:59

refactor(fsdp): split sharding strategy validation

ee1ba7c

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

style(fsdp): match main config spacing

30dddd1

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

svcnvidia-nemo-ci added the Approved All necessary approvals have been made label Aug 30, 2026

copy-pr-bot Bot temporarily deployed to test August 30, 2026 15:01 Inactive

copy-pr-bot Bot temporarily deployed to public August 30, 2026 15:04 Inactive

wujingyue removed request for a team August 30, 2026 15:16

wujingyue added this pull request to the merge queue Aug 31, 2026

Merged via the queue into NVIDIA:main with commit 713ea6a Aug 31, 2026
201 checks passed

wujingyue deleted the agent/mfsdp-v2-hsdp branch August 31, 2026 02:32

Autumn1998 mentioned this pull request Sep 1, 2026

Support independent HSDP instances for expert parameters #7012

Merged

wujingyue mentioned this pull request Sep 1, 2026

fix(fsdp): drop stale MFSDP v2 guard rails for clip_grad, precision-aware optimizer, HSDP, and optimizer CUDA graphs NVIDIA-NeMo/Megatron-Bridge#5933

Merged

2 of 5 tasks

ksivaman pushed a commit to ksivaman/Megatron-LM that referenced this pull request Sep 1, 2026

feat(fsdp): expose HSDP and HFSDP through the MFSDP v2 adapter (NVIDI…

b5c8abc

…A#6873) Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Navigation Menu

Repository navigation

feat(fsdp): expose HSDP and HFSDP through the MFSDP v2 adapter - #6873

feat(fsdp): expose HSDP and HFSDP through the MFSDP v2 adapter#6873
wujingyue merged 11 commits into
NVIDIA:mainfrom
wujingyue:agent/mfsdp-v2-hsdp

wujingyue commented Aug 26, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Aug 26, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wujingyue commented Aug 26, 2026

Uh oh!

Autumn1998 Aug 26, 2026

Uh oh!

wujingyue Aug 26, 2026

Uh oh!

copy-pr-bot Bot commented Aug 26, 2026

Uh oh!

wujingyue commented Aug 26, 2026

Uh oh!

wujingyue commented Aug 26, 2026

Uh oh!

nemo-automation-bot Bot commented Aug 31, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

wujingyue commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

No HSDP/HFSDP special case

Validation

Uh oh!

copy-pr-bot Bot commented Aug 26, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wujingyue commented Aug 26, 2026

Uh oh!

Autumn1998 Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

wujingyue Aug 26, 2026

Choose a reason for hiding this comment

Uh oh!

copy-pr-bot Bot commented Aug 26, 2026

Uh oh!

wujingyue commented Aug 26, 2026

Uh oh!

wujingyue commented Aug 26, 2026

Uh oh!

nemo-automation-bot Bot commented Aug 31, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

wujingyue commented Aug 26, 2026 •

edited

Loading