Enable Pipeline Parallelism Support for M-FSDP v2 by shjwudp · Pull Request #6486 · NVIDIA/Megatron-LM

shjwudp · 2026-08-12T16:26:49Z

Summary

Under virtual pipeline parallelism, wrap_model_chunks_with_ddp wraps each chunk in its own FullyShardedDataParallelV2, and each adapter __init__ previously opened a fresh fully_shard_context. Every chunk therefore got independent communication streams, microbatch state, and prefetch orders, so the combined-1F1B schedule could not coordinate all-gathers or reduce-scatters across chunks.

Changes

fully_shard_context gains reuse_existing: when a context is already active on the same device, it is yielded instead of creating a new one, and only the outermost scope finalizes it. Nesting rejection remains the default.
The MFSDP v2 adapter opens its construction scope with reuse_existing, so it joins an ambient training-loop context when present.
wrap_model_chunks_with_ddp wraps the per-chunk wrapping loop in one fully_shard_context when DP is the MFSDP v2 adapter and there is more than one chunk, so all VPP chunks share one context (streams, orders, is_last_microbatch) and finalize is deferred to the outer scope.
ModelBuilder path (prepare_existing_model_chunks_for_distributed_training) mirrors the same ambient-context fix.
at do not follow the static forward/backward orders.

Tests

Unit coverage for trace, replay wrap-around, orientation-aware prefetch, divergence re-trace, stale release, and nested reuse scopes sharing the outer context.

Validation (GPU test results)

Tested on EOS H100 with deepseek_v3_proxy + real DeepSeek-V3 SlimPajama, seed 1234, PP=2/VPP=2, TP1/EP1/CP1 (DP=4), MBS=1/GBS=8, seq 4096, 1000 iterations, tensorboard + W&B.

1000-step real-data convergence — M-FSDP v2 ≡ DistOpt. Both PP=2/VPP=2 configs converge to the same val-loss at every 100-iter eval; max diff ~0.02 (~0.3% relative), both 0 NaN, exit 0. Final val-loss @iter1000: M-FSDP v2 6.023 vs DistOpt 6.019 (lm-loss @1000: 6.490 vs 6.485; both start ~11.93 → ~6.0). W&B runs: nidznep9 (M-FSDP v2) / 6d21q4g7 (DistOpt), project jianbinc-mfsdp-v2-vpp-benchmark.

copy-pr-bot · 2026-08-12T16:26:53Z

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.

PR NVIDIA#6486 shares a single FsdpContext across VPP model chunks, but FullyShardedOptimizer._validate_config still rejects len(model_chunks) != 1, so MFSDP v2 + VPP aborts at optimizer build with "MFSDP v2 currently supports exactly one model chunk." FullyShardedOptimizer already iterates every entry in self.model_chunks (zero_grad) and get_parameters() aggregates across all chunks, so the lone non-empty guard is the only blocker. Mirror the author's later dev/vpp-shared-fsdp-context branch: reject only an empty model_chunks list. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Minimal port of NVIDIA#6486's core idea: all VPP model chunks constructed by one wrap call share a single MFSDP v2 fully_shard_context, so their FSDP modules join one context and it is finalized once after the chunk loop. - fully_shard_context gains reuse_existing=True; reused scopes join the outermost scope's context on the same device (no finalize) and raise if asked to share across devices. Default (reuse_existing=False) keeps the existing no-nesting RuntimeError. - FsdpContext now records its device so the same-device reuse check has a value to compare. - FullyShardedDataParallelV2 opens its construction scope with reuse_existing=True so per-chunk wrappers join an ambient context. - wrap_model_chunks_with_ddp and _ddp_wrap wrap the per-chunk DP construction loop in one ambient fully_shard_context when wrapping MFSDP v2 and more than one chunk. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

PR NVIDIA#6486 (split/pr-d-shared-context) adds FsdpContext sharing across VPP model chunks but forgets to relax FullyShardedDataParallelV2._validate_config, which still unconditionally requires TP=PP=CP=1. Because virtual pipeline parallelism requires pipeline_model_parallel_size > 1, the exact config this feature targets (MFSDP v2 + VPP) is rejected at model init with "MFSDP v2 does not currently support: pipeline_model_parallel_size=N". Also, wrap_model_chunks_with_ddp / _ddp_wrap set disable_bucketing=True on non-first VPP chunks, which _validate_config rejects ("MFSDP v2 does not support disabling bucketing.") since MFSDP v2 keeps bucketing enabled and shards every chunk independently. Changes: - Remove pipeline_model_parallel_size from the unsupported-parallelisms list and drop pp from the process-group size-1 check in _validate_config. The DP mesh is built solely from dp_cp, so DP sharding is unaffected by PP/VPP. - Force disable_bucketing=False for MFSDP v2 chunks in both wrap functions so a multi-chunk (VPP) wrap no longer trips the disable-bucketing guard. Base: split/pr-d-shared-context @ 2cd3d9f (PR NVIDIA#6486 head). Signed-off-by: Jack Chang <jianbinc@nvidia.com>

PR NVIDIA#6486 shares a single FsdpContext across VPP model chunks, but FullyShardedOptimizer._validate_config still rejects len(model_chunks) != 1, so MFSDP v2 + VPP aborts at optimizer build with "MFSDP v2 currently supports exactly one model chunk." FullyShardedOptimizer already iterates every entry in self.model_chunks (zero_grad) and get_parameters() aggregates across all chunks, so the lone non-empty guard is the only blocker. Mirror the author's later dev/vpp-shared-fsdp-context branch: reject only an empty model_chunks list. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

wujingyue · 2026-09-10T00:41:28Z

    *,
    use_symmetric_memory: bool = False,
    unify_communication_stream: bool = False,
+    reuse_existing: bool = False,


This looks fishy... Why (why do we need this) and why not (why not other approaches)?

The goal here is to share a single FsdpContext across VPP chunks, thereby avoiding an excessive number of communication streams, potential CCA fragmentation, and other issues introduced by having multiple FsdpContext instances.

share a single FsdpContext across VPP chunks

Why can't we share them by wrapping everything in the same FsdpContext?

Good point—thanks! Removed reuse_existing.

guihong-nv · 2026-09-10T12:50:46Z

+        # wrap_model_chunks_with_ddp / dist_utils._ddp_wrap). The data-parallel mesh is
+        # built solely from dp_cp, so DP sharding is unaffected by how layers are split
+        # across pipeline stages or virtual stages.
        unsupported_parallelisms = [


Enabling PP introduces duplicate embedding/output weights, but FullyShardedOptimizer.get_grad_norm() counts both copies in its WORLD reduction. Its replication correction only covers DP. A four-rank reproduction returned 7.071 instead of 5, causing excessive gradient clipping. Preserve the shared-parameter marker and exclude duplicate copies from gradient statistics.

Thanks! When share_embeddings_and_output_weights=False, counting the embedding and output-weight gradient norms twice is definitely incorrect.

Thanks for pointing this out—I’ll address it.

Are you saying there's a bug in FullyShardedOptimizer.get_grad_norm() when handling tied weights? Can you file a bug with a minimal, non-mcore repro so I can take a look? It doesn't seem to be PP-specific and can affect everything.

In MCore’s DP path, get_grads_for_grad_norm selects the gradients eligible for gradient-norm computation. For shared weights, it uses param_is_not_shared to identify shared parameters; gradients associated with shared parameters are excluded from the calculation.

PR NVIDIA#6486 shares a single FsdpContext across VPP model chunks, but FullyShardedOptimizer._validate_config still rejects len(model_chunks) != 1, so MFSDP v2 + VPP aborts at optimizer build with "MFSDP v2 currently supports exactly one model chunk." FullyShardedOptimizer already iterates every entry in self.model_chunks (zero_grad) and get_parameters() aggregates across all chunks, so the lone non-empty guard is the only blocker. Mirror the author's later dev/vpp-shared-fsdp-context branch: reject only an empty model_chunks list. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

…llelism MFSDP v2 shards embeddings with uneven (flat-buffer) offsets, and different PP stages can use different flattened layouts. Two ranks owning the same logical slice may hold it at different flat-buffer offsets with different local shard shapes, so reducing the raw local shards directly would combine mismatched elements. Reconstruct the full logical embedding gradient via uneven_dtensor_to_full_tensor before the PP all-reduce, then copy back the slice this rank owns into its shard's local tensor. Without this, MFSDP v2 + PP/VPP fails in finalize_model_grads._allreduce_embedding_grad with "attempting to assign a gradient of size '[N,1024]' to a tensor of size '[V,1024]'". This is NVIDIA PR NVIDIA#6500's change (author: Jianbin Chang), ported so shjwudp:split/pr-d-shared-context (PR NVIDIA#6486) actually runs MFSDP v2 + VPP. Excluded the separate MFSDP-v2+MoE empty-shard TE workaround (PR NVIDIA#7002). Validated on EOS H100 (deepseek_v3_proxy, real DeepSeek-V3 SlimPajama, PP=2/VPP=2, seed 1234, 1000 iters): M-FSDP v2 and DistOpt converge to the same val-loss (diff <= ~0.02, 0 NaN, exit 0). Signed-off-by: Jianbin Chang <jianbinc@nvidia.com> Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Extract the inline scatter-back code in _allreduce_embedding_grad into two named helpers (no behavior change): - _full_embedding_grad(orig_grad): expand an MFSDP v2 uneven (flat-buffer) sharded embedding grad to the full logical gradient via uneven_dtensor_to_full_tensor (or None when the grad is None). - _copy_embedding_grad_back(orig_grad, full_grad): copy this rank's logical slice (from chunk.offsets/sizes) back into its shard's local tensor. This keeps the cross-stage all-reduce operating on a canonical (uniform-row) logical gradient, which is correct because different PP stages can hold the same logical embedding at different flat-buffer offsets with different local shard shapes (a direct all-reduce of raw local shards would combine mismatched elements). Signed-off-by: Jianbin Chang <jianbinc@nvidia.com> Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

…parallel.py Signed-off-by: Jack Chang <jianbinc@nvidia.com>

The linting CI job failed in the black step on this change. With the repository's black flags (--skip-magic-trailing-comma, --skip-string-normalization) three bracketed expressions fit within the 100-column limit and must collapse onto single lines: - tests/unit_tests/distributed/mfsdp_v2/test_pipeline_parallel.py: the @pytest.mark.skipif decorator and the chunks = [...] assignment - megatron/core/distributed/fsdp/mcore_fsdp_adapter.py: the unsupported_parallelisms = [...] assignment Formatting only; no behavior change. Verified with the repository's linting toolchain (black 26.3.0, isort 5.13.2, pylint 3.2.6, ruff 0.9.10): all four steps exit 0. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Restore the view guard in deallocate_output_tensor() and turn the pipeline-output pseudo-free off for Megatron-FSDP v2 instead of letting the guard abort training. Megatron-FSDP v2's full-backward hook makes the pipeline stage output an autograd view of the module's own output, and deallocate_output_tensor() cannot pseudo-free a view: the base tensor owns the storage and keeps it alive, so the swap reclaims no memory. Keep the guard meaningful and disable the optimization for the configuration that cannot use it, in both the Python-args and YAML config paths. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

The YAML config path merges the language_model and model_parallel sections, and the fixture already carries every TransformerConfig field - including overlap_p2p_comm - in language_model, so supplying it again through model_parallel raised "got multiple values for keyword argument". Stop setting init_method to the "xavier_uniform" string as well: that branch emits a scaled_init_method kwarg that TransformerConfig does not accept. Signed-off-by: Jack Chang <jianbinc@nvidia.com>

shjwudp · 2026-09-22T14:12:45Z

/ok to test 16ee180

nemo-automation-bot · 2026-09-22T16:07:17Z

🔄 Merge queue validation started!

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

nemo-automation-bot · 2026-09-23T03:38:07Z

🔄 Merge queue validation started!

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

nemo-automation-bot · 2026-09-23T04:54:19Z

🔄 Merge queue validation started!

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

shjwudp changed the title ~~Share one FsdpContext across VPP model chunks (#123)~~ [Prototype] Share one FsdpContext across VPP model chunks Aug 12, 2026

shjwudp mentioned this pull request Aug 12, 2026

Expose MFSDP v2 custom-schedule lifecycle #6484

Merged

This was referenced Aug 27, 2026

Megatron-LM Monthly Activity Report — August 2026 (Final, through August 31, 2026) #6921

Open

Megatron-LM Monthly Activity Report — August 2026 #7006

Open

shjwudp force-pushed the split/pr-d-shared-context branch 2 times, most recently from b182624 to ea04810 Compare September 2, 2026 14:51

lhb8125 mentioned this pull request Sep 7, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

shjwudp force-pushed the split/pr-d-shared-context branch from ea04810 to de0529a Compare September 8, 2026 07:43

shjwudp mentioned this pull request Sep 9, 2026

Allow MFSDP v2 with pipeline parallelism (VPP shared context) shjwudp/Megatron-LM#158

Merged

shjwudp changed the title ~~[Prototype] Share one FsdpContext across VPP model chunks~~ Enable Pipeline Parallelism Support for M-FSDP v2 Sep 9, 2026

shjwudp mentioned this pull request Sep 9, 2026

Fix MFSDP v2 shared-embedding gradient all-reduce under pipeline parallelism shjwudp/Megatron-LM#160

Merged

shjwudp marked this pull request as ready for review September 9, 2026 23:57

shjwudp requested a review from a team as a code owner September 9, 2026 23:57

shjwudp requested a review from a team September 9, 2026 23:57

shjwudp requested review from a team as code owners September 9, 2026 23:57

github-actions Bot added the complexity: medium label Sep 9, 2026

shjwudp force-pushed the split/pr-d-shared-context branch from 59fce5c to d20eb99 Compare September 10, 2026 00:03

wujingyue reviewed Sep 10, 2026

View reviewed changes

shjwudp requested a review from a team as a code owner September 10, 2026 01:51

shjwudp mentioned this pull request Sep 10, 2026

Prototype: Enable MXFP8, VPP, Combined 1F1B, and Distributed Muon with MFSDP v2 #6197

Draft

7 of 9 tasks

guihong-nv reviewed Sep 10, 2026

View reviewed changes

shjwudp added 10 commits September 22, 2026 21:58

remove reuse_existing

2333f32

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

remove if out._base is not None:

f6c4653

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

Polish unit test tests/unit_tests/distributed/mfsdp_v2/test_pipeline_…

e52b2e3

…parallel.py Signed-off-by: Jack Chang <jianbinc@nvidia.com>

revert unnecessary change

e661729

Signed-off-by: Jack Chang <jianbinc@nvidia.com>

shjwudp force-pushed the split/pr-d-shared-context branch from 2676443 to 16ee180 Compare September 22, 2026 14:09

copy-pr-bot Bot temporarily deployed to test September 22, 2026 14:15 Inactive

copy-pr-bot Bot temporarily deployed to public September 22, 2026 14:16 Inactive

shjwudp enabled auto-merge September 22, 2026 14:42

shjwudp added this pull request to the merge queue Sep 22, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Sep 22, 2026

Connor-XY mentioned this pull request Sep 22, 2026

Silent numerical-correctness bugs in Megatron Core parallelism composition #7452

Open

shjwudp added this pull request to the merge queue Sep 23, 2026

Merged via the queue into NVIDIA:main with commit 2539bc3 Sep 23, 2026
102 of 103 checks passed

lhb8125 mentioned this pull request Oct 1, 2026

Megatron-LM Monthly Activity Report — September 2026 #7773

Open

This was referenced Oct 12, 2026

fix(distributed): restore FSDP v1 pipeline training and checkpoints #8159

Draft

fix(checkpoint): use global FSDP v1 pipeline layer keys NVIDIA-NeMo/Megatron-Bridge#6424

Draft

sbhavani mentioned this pull request Oct 15, 2026

[ROADMAP][2026 Q3] Megatron Core Roadmap #5676

Open

Conversation

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

Uh oh!

Summary

Changes

Tests

Validation (GPU test results)

Uh oh!

copy-pr-bot Bot commented Aug 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

shjwudp commented Sep 22, 2026

Uh oh!

nemo-automation-bot Bot commented Sep 22, 2026

Uh oh!

Uh oh!

nemo-automation-bot Bot commented Sep 23, 2026

Uh oh!

nemo-automation-bot Bot commented Sep 23, 2026

Uh oh!

Uh oh!

This branch was previously deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

shjwudp commented Aug 12, 2026 •

edited

Loading