[Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP. by cspades · Pull Request #4427 · NVIDIA/Megatron-LM

cspades · 2026-04-22T15:24:14Z

What does this PR do ?

FusedAdam.use_decoupled_grad not set with FSDP + --use-precision-aware-optimizer, which is a bug introduced in Fix incorrectly set decoupled_grad and DistOpt mechanics for MFSDP. #4133.
- While we have a TransformerEngine MXFP8 + FusedAdam + CPU offloading V2 test for this, param.grad = None is a no-op for FusedAdam.step() which doesn't throw an error or reveal that Megatron-FSDP + FusedAdam would not converge.
- Related to this MBridge PR that mirrors the decoupled gradient logic in MLM: [Megatron-FSDP] Fix segmentation fault when using precision-aware optimizers with the updated MCore. NVIDIA-NeMo/Megatron-Bridge#3510
  - Megatron-Bridge logic also needs to be updated to have use_precision_aware_optimizer activate megatron_fsdp_use_decoupled_grad. This caused seg-faults with clip_grad_norm when MLPerf upgraded MCore + MBridge and param.decoupled_grad = None while clip_grad_norm(use_decoupled_grad=True).

Testing

Added even more assertions and null/zero gradient checks in the unit test to catch this silent bug.
While FP8-Delayed does converge even without this fix, now MXFP8 (which falls outside of the scope of use_precision_aware_optimizer_no_fp8_or_ds_fp8 originally responsible for setting use_decoupled_grad=True) now also converges:

# MXFP8 + FusedAdam(use_decoupled_grad=True)
[2026-04-23 19:40:32.890713] iteration      100/15258789 | consumed samples:        12800 | elapsed time per iteration (ms): 10401.4 | throughput per GPU (TFLOP/s/GPU): 1297.3 | learning rate: 4.915263E-05 | global batch size:   128 | lm loss: 1.168019E-02 | loss scale: 1.0 | grad norm: 0.294 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |

TODO

Update the TransformerEngine E2E test to print out loss for every training step, and update the MCore commit hash to point to this PR's fix.

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code Typing guidelines
I have added relevant documentation
I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

When your PR is ready, click Ready for Review.
An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
- Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

For MRs into `dev` branch

The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

github-actions · 2026-04-22T15:24:24Z

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

Add the oncall reviewer (optional reviewer)
Add required review teams based on your changes

See the contribution guide for more details.

copy-pr-bot · 2026-04-22T15:24:27Z

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.

…-FSDP. Signed-off-by: Cory Ye <cye@nvidia.com>

ko3n1g · 2026-04-22T16:38:47Z

/claude review

claude

LGTM

Signed-off-by: Cory Ye <cye@nvidia.com>

svcnvidia-nemo-ci · 2026-04-27T22:12:24Z

🔄 Merge queue validation started!

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

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com>

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com> Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com>

cspades self-assigned this Apr 22, 2026

cspades requested review from a team as code owners April 22, 2026 15:24

svcnvidia-nemo-ci marked this pull request as draft April 22, 2026 15:24

svcnvidia-nemo-ci added this to the Core 0.16 milestone Apr 22, 2026

copy-pr-bot Bot temporarily deployed to test April 22, 2026 15:25 Inactive

cspades marked this pull request as ready for review April 22, 2026 15:29

svcnvidia-nemo-ci requested a review from a team April 22, 2026 15:29

svcnvidia-nemo-ci added Final Review PR is in the "final review" stage complexity: low labels Apr 22, 2026

Fix segfault caused by not overriding use_decoupled_grad for Megatron…

a1038a6

…-FSDP. Signed-off-by: Cory Ye <cye@nvidia.com>

cspades force-pushed the cye/mfsdp-decgrad-segfault branch from 6de63d2 to a1038a6 Compare April 22, 2026 16:25

copy-pr-bot Bot temporarily deployed to test April 22, 2026 16:26 Inactive

claude Bot approved these changes Apr 22, 2026

View reviewed changes

cspades marked this pull request as draft April 22, 2026 16:52

cspades changed the title ~~[Main] Fix segfault caused by not overriding use_decoupled_grad for Megatron…~~ [Main] Fix invisible issues related to use_decoupled_grad for Megatron-FSDP. Apr 22, 2026

ilml approved these changes Apr 24, 2026

View reviewed changes

Merge branch 'main' into cye/mfsdp-decgrad-segfault

874917b

cspades marked this pull request as ready for review April 24, 2026 02:31

svcnvidia-nemo-ci added complexity: low and removed complexity: low labels Apr 24, 2026

copy-pr-bot Bot temporarily deployed to test April 24, 2026 02:32 Inactive

Add more assert tests that guard against the FusedAdam.step() bug.

1a5fcfa

Signed-off-by: Cory Ye <cye@nvidia.com>

copy-pr-bot Bot temporarily deployed to test April 24, 2026 03:07 Inactive

cspades changed the title ~~[Main] Fix invisible issues related to use_decoupled_grad for Megatron-FSDP.~~ [Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP. Apr 24, 2026

This was referenced Apr 24, 2026

[Megatron-FSDP] Fix segmentation fault when using precision-aware optimizers with the updated MCore. NVIDIA-NeMo/Megatron-Bridge#3510

Merged

[Dev] Fix mis-set decoupled gradient for Megatron-FSDP. #4426

Merged

ericharper approved these changes Apr 24, 2026

View reviewed changes

shjwudp approved these changes Apr 24, 2026

View reviewed changes

cspades added the module: megatron-fsdp label Apr 24, 2026

deepakn94 approved these changes Apr 27, 2026

View reviewed changes

svcnvidia-nemo-ci added Approved All necessary approvals have been made and removed Final Review PR is in the "final review" stage labels Apr 27, 2026

Merge branch 'main' into cye/mfsdp-decgrad-segfault

881fa37

cspades enabled auto-merge April 27, 2026 20:14

copy-pr-bot Bot temporarily deployed to test April 27, 2026 20:14 Inactive

cspades added this pull request to the merge queue Apr 27, 2026

cspades mentioned this pull request Apr 27, 2026

Perf Config for 1 node GB200 DSV3 NVIDIA-NeMo/Megatron-Bridge#3545

Closed

5 tasks

Merged via the queue into NVIDIA:main with commit 0447347 Apr 27, 2026
68 checks passed

cspades deleted the cye/mfsdp-decgrad-segfault branch April 27, 2026 22:46

cspades mentioned this pull request May 18, 2026

[bug] FSDP training does not converge NVIDIA-NeMo/Megatron-Bridge#3826

Open

yangbofun pushed a commit to xlm-research/Megatron-LM that referenced this pull request May 22, 2026

[Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP. (N…

25008c5

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com>

yhgalaxy pushed a commit to yhgalaxy/Megatron-LM that referenced this pull request Jun 17, 2026

[Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP. (N…

0f321a5

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com> Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>

terminator123 pushed a commit to 021ai/Megatron-LM that referenced this pull request Aug 3, 2026

[Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP. (N…

0d51a2c

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com>

Navigation Menu

Uh oh!

[Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP. - #4427

[Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP.#4427
cspades merged 4 commits into
NVIDIA:mainfrom
cspades:cye/mfsdp-decgrad-segfault

cspades commented Apr 22, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Apr 22, 2026

Uh oh!

copy-pr-bot Bot commented Apr 22, 2026

Uh oh!

ko3n1g commented Apr 22, 2026

Uh oh!

claude Bot left a comment

Uh oh!

svcnvidia-nemo-ci commented Apr 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

Uh oh!

Conversation

cspades commented Apr 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Testing

TODO

Contribution process

Pre-checks

Code review

Step 1: Mark PR as "Ready for Review"

Step 2: Final Review

Step 3: Approved

Merge

Uh oh!

github-actions Bot commented Apr 22, 2026

Uh oh!

copy-pr-bot Bot commented Apr 22, 2026

Uh oh!

ko3n1g commented Apr 22, 2026

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Uh oh!

svcnvidia-nemo-ci commented Apr 27, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

cspades commented Apr 22, 2026 •

edited

Loading