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

[cspades]

What does this PR do ?

• FusedAdam.use_decoupled_grad not set with FSDP, 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

Main Branch: #4427

Related Consequences

• 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

• 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"

1. When your PR is ready, click Ready for Review.
2. 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.

[copy-pr-bot]

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.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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