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

wujingyue · 2026-09-01T19:52:54Z

Summary

Remove stale MFSDP v2 validation guards now that Megatron-LM supports gradient clipping, precision-aware optimizers, HSDP/HFSDP, optimizer CUDA graphs, and full-iteration CUDA graphs.

Setting	Megatron-LM support
Gradient clipping	NVIDIA/Megatron-LM#6489
Precision-aware optimizer	NVIDIA/Megatron-LM#6506; clipping compatibility fix in NVIDIA/Megatron-LM#7074
HSDP/HFSDP	NVIDIA/Megatron-LM#6873
Optimizer CUDA graphs	NVIDIA/Megatron-LM#6498
Full-iteration CUDA graphs	NVIDIA/Megatron-LM#7075

Current Bridge main's MCore pin includes the required support. The remaining unsupported-configuration guards are unchanged.

Also wire MFSDP v2's no_sync into the training schedule even when overlap_grad_reduce=False, so gradient accumulation defers synchronization correctly.

Test coverage

Extend the dense HybridModel v2 smoke test to enable precision-aware optimization, full-iteration graphs, and optimizer graphs with FP32 main gradients.
Disable check_for_nan_in_loss only in that graph test: the host-side loss check is incompatible with capture and rejected by Bridge's configuration validator.
Add two-instance HSDP and HFSDP smoke cases.
Retain the existing MoE EP=2 smoke coverage.

Validation

Previously verified the DeepSeek-V3 MFSDP v2 + EP=2 proxy on 8 H100 GPUs against Megatron-LM main with optimizer.clip_grad=1.0: 20/20 iterations, finite and decreasing loss, no skipped/NaN iterations. This is prior proxy validation, not a run of the updated functional tests.
Ruff lint, formatting, and git diff --check for the graph-test fix.
Full pre-commit checks: attempted, but uv is unavailable on the login node.
Run the updated targeted FSDP functional tests against current main's MCore pin.
Current-head CI.

copy-pr-bot · 2026-09-01T19:52:57Z

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-09-01T19:53:00Z

/ok to test 6764976

wujingyue · 2026-09-01T20:02:39Z

/ok to test 65b7a90

copy-pr-bot · 2026-09-01T20:05:08Z

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-09-01T20:05:10Z

/ok to test 4455bfc

wujingyue · 2026-09-01T20:10:27Z

/ok to test f2289d2

wujingyue · 2026-09-04T05:52:07Z

/ok to test c2028ab

wujingyue · 2026-09-04T06:49:13Z

/ok to test 3c7c146

wujingyue · 2026-09-04T21:10:19Z

/ok to test 3c7c146

…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>

… smoke test The override existed only to dodge the guard rail removed in the previous commit; the base default (clip_grad=1.0) now exercises real gradient clipping under MFSDP v2. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Asserts ConfigContainer.validate() accepts clip_grad, use_precision_aware_optimizer, optimizer_cuda_graph and HSDP under megatron_fsdp_version=2. Config validation only, so it runs on CPU; skipped when the installed MCore has no MFSDP v2. Restoring any of the removed guards fails these cases. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Parameterizes the dense HybridModel smoke test over use_precision_aware_optimizer, optimizer_cuda_graph and HSDP, so each one actually trains rather than only passing config validation. clip_grad is covered throughout by the base optimizer default of 1.0. optimizer_cuda_graph pins the main-gradient dtype to the main-weight dtype, which capturable TE FusedAdam requires. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

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

MCore's v2 hybrid path builds its device mesh from pg_collection's inter_dist_opt and intra_dp_cp groups, and Bridge already supplies both via use_mpu_process_groups(), where intra_dp_cp is created with partial_data_parallel=True -- exactly the inner group MCore expects. That plumbing is version-agnostic, which is why v1 has supported HSDP here all along, so the two v2-only guard rails were stale rather than load-bearing. Cover both strategies with a training test: no_shard outer over an optim_grads_params inner axis (HSDP) and optim outer (HFSDP). Also drop the MCORE_HAS_MEGATRON_FSDP_V2 skipif from the v2 tests, since the pinned MCore ships v2 unconditionally. Signed-off-by: Jingyue Wu <jingyuew@nvidia.com> Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Megatron-LM #7075 added full-iteration CUDA graph support to MFSDP v2, so the model-side guard is stale. The dense HybridModel smoke test now runs with cuda_graph_impl="full_iteration"; MoE stays out because its dispatch shapes are not capturable. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

DDP and MFSDP v1 only reduce during backward when overlap_grad_reduce is on, so gating no_sync_func on that flag is right for them. MFSDP v2 always reduces in backward, so no_sync is how the schedule marks a non-final microbatch -- a correctness requirement, not an overlap optimization. Without it, with an outer DP axis and more than one microbatch, every backward finalizes the DP-outer axis and only the last microbatch's gradient reaches the optimizer. Training still converges, just to the wrong curve. Measured on 8xH100 with a DeepSeek-V3 proxy (EP=2, HFSDP, 2 microbatches), max relative lm-loss deviation from a no-FSDP reference over 20 iterations: v1 4.1e-03, v2 before 1.8e-01, v2 after 3.2e-04. Requires Megatron-LM #7186, which adds FullyShardedDataParallelV2.no_sync; without it v2 inherits the no-op from _BaseDataParallel. Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

wujingyue · 2026-09-16T16:54:39Z

/ok to test a710c53

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

wujingyue · 2026-09-16T19:56:58Z

/ok to test aa4d4db

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

…for Megatron-Bridge Upstream Megatron-Core split Megatron-FSDP into FullyShardedDataParallelV1 (the implementation miles-main ships as FullyShardedDataParallel) and FullyShardedDataParallelV2, and turned FullyShardedDataParallel into a factory (NVIDIA#5865). Megatron-Bridge main imports both names from mcore_fsdp_adapter since NVIDIA-NeMo/Megatron-Bridge#5933 (conversion/utils.py::unwrap_model, training/setup.py, training/train.py), so on miles-main every bridge-mode weight load / export raised ImportError with bridge-new @ 2e09c234a. Add FullyShardedDataParallelV1 as an alias of the existing class and a FullyShardedDataParallelV2 placeholder whose constructor raises NotImplementedError. The v2 implementation is not ported and FullyShardedDataParallel stays the v1 class, so the existing isinstance checks in megatron/core/utils.py and megatron/training/training.py are unchanged. Signed-off-by: Yusheng Su <yushengsu.thu@gmail.com>

copy-pr-bot Bot temporarily deployed to test September 1, 2026 19:54 Inactive

wujingyue force-pushed the agent/mfsdp-v2-guard-rails branch from 6764976 to 65b7a90 Compare September 1, 2026 20:02

copy-pr-bot Bot had a problem deploying to test September 1, 2026 20:04 Error

wujingyue force-pushed the agent/mfsdp-v2-guard-rails branch from 65b7a90 to 4455bfc Compare September 1, 2026 20:05

copy-pr-bot Bot had a problem deploying to test September 1, 2026 20:06 Error

copy-pr-bot Bot temporarily deployed to test September 1, 2026 20:11 Inactive

wujingyue marked this pull request as ready for review September 1, 2026 21:27

yaoyu-33 added area:training Training loop, callbacks, and runtime integration bug Something isn't working full-test-suite needs-more-tests Requires additional L0 and L1 test coverage before merge needs-review PR is ready for code review and waiting on a reviewer labels Sep 1, 2026

copy-pr-bot Bot temporarily deployed to test September 4, 2026 06:50 Inactive

wujingyue marked this pull request as draft September 8, 2026 05:30

wujingyue force-pushed the agent/mfsdp-v2-guard-rails branch from c1d2475 to 37f8f91 Compare September 8, 2026 05:38

wujingyue and others added 6 commits September 11, 2026 13:53

fix(fsdp): retain unsupported MFSDP v2 guard rails

467458c

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

fix(fsdp): support precision-aware optimizer in MFSDP v2

bb3d810

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

wujingyue and others added 3 commits September 11, 2026 13:53

fix(fsdp): enable optimizer CUDA graphs in MFSDP v2

d3778ab

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

wujingyue force-pushed the agent/mfsdp-v2-guard-rails branch from 37f8f91 to f075e16 Compare September 11, 2026 20:53

wujingyue mentioned this pull request Sep 11, 2026

[BUG] CPU latency makes MFSDP v2 slower than v1 on a DeepSeek-V3 proxy NVIDIA/Megatron-LM#7264

Open

wujingyue linked an issue Sep 11, 2026 that may be closed by this pull request

Integrate MFSDP v2 with Megatron Bridge NVIDIA/Megatron-LM#6325

Closed

test(fsdp): disable loss NaN check in full-iteration graph smoke test

9d5a652

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

wujingyue mentioned this pull request Sep 16, 2026

[BUG] MFSDP v2 with HybridEP: intermittent bad gradients in eager training NVIDIA/Megatron-LM#7400

Closed

wujingyue added 3 commits September 16, 2026 09:47

refactor(fsdp): simplify v2 compatibility and separate graph smoke test

2ef5fb1

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

refactor(fsdp): inline wrapper types and remove compatibility module

a419b68

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

refactor(fsdp): simplify no-sync wrapper condition

a710c53

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

wujingyue marked this pull request as ready for review September 16, 2026 16:54

wujingyue requested a review from yaoyu-33 September 16, 2026 16:55

copy-pr-bot Bot deployed to test September 16, 2026 16:55 Active

wujingyue requested a review from cuichenx September 16, 2026 16:55

test(fsdp): use TE RNG tracker for CUDA graph smoke test

aa4d4db

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

copy-pr-bot Bot deployed to test September 16, 2026 19:59 Active

test(fsdp): reset MCore RNG tracker before graph smoke test

73c5e63

Signed-off-by: Jingyue Wu <wujingyue@gmail.com>

copy-pr-bot Bot deployed to test September 16, 2026 22:28 Active

yaoyu-33 approved these changes Sep 17, 2026

View reviewed changes

yaoyu-33 merged commit 809f0e6 into NVIDIA-NeMo:main Sep 17, 2026
143 checks passed

yushengsu-thu mentioned this pull request Sep 18, 2026

[port] Megatron-Core APIs needed by Megatron-Bridge main (#4452 MLA resolver, gdn alias, #5865 FSDP names) radixark/Megatron-LM#102

Merged

malay-nagda mentioned this pull request Sep 26, 2026

1f1b cuTE DSL fusion full-iter CG | rubin- NT3 ultra, DSv4 flash #6046

Merged

Navigation Menu

Repository navigation

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

fix(fsdp): drop stale MFSDP v2 guard rails for clip_grad, precision-aware optimizer, HSDP, and optimizer CUDA graphs#5933
yaoyu-33 merged 16 commits into
NVIDIA-NeMo:mainfrom
wujingyue:agent/mfsdp-v2-guard-rails

wujingyue commented Sep 1, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 1, 2026

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 4, 2026

Uh oh!

wujingyue commented Sep 4, 2026

Uh oh!

wujingyue commented Sep 4, 2026

Uh oh!

wujingyue commented Sep 16, 2026

Uh oh!

wujingyue commented Sep 16, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

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

Uh oh!

Summary

Test coverage

Validation

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 1, 2026

Uh oh!

copy-pr-bot Bot commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 1, 2026

Uh oh!

wujingyue commented Sep 4, 2026

Uh oh!

wujingyue commented Sep 4, 2026

Uh oh!

wujingyue commented Sep 4, 2026

Uh oh!

wujingyue commented Sep 16, 2026

Uh oh!

wujingyue commented Sep 16, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

wujingyue commented Sep 1, 2026 •

edited

Loading