Fix incorrectly set decoupled_grad and DistOpt mechanics for MFSDP.

[cspades]

What does this PR do ?

• Update DistributedOptimizer to correctly use decoupled_grad with Megatron-FSDP. (FSDP will always use decoupled gradients when using FusedAdam, which simplifies the logic a lot.)

TL;DR

use_decoupled_grad=self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8
or (
    # Megatron-FSDP always uses decoupled_grad with FusedAdam.
    self.config.use_precision_aware_optimizer
    and getattr(params[0], "__fsdp_param__", False)  # == ddp_config.use_megatron_fsdp
),

Details

• Megatron-FSDP does not use FusedAdam's master weights, but Megatron-LM hard-codes master_weights=True if OptimizerConfig.use_precision_aware_optimizer_no_fp8_or_ds_fp8 / use_decoupled_grad are True. This PR turns off FusedAdam master weights when using Megatron-FSDP, as FusedAdam should only provide an optimizer.step() to Megatron-FSDP's DTensor(FP32/BF16) main weights.
• Improved Megatron-FSDP argument validation and added NotImplementedError for completely unexpected paths in the distributed optimizer.

Testing

• Adding E2E --use-precision-aware-optimizer unit test to temporarily guarantee functionality. Cross-tested with this TE E2E test: Add Megatron-FSDP E2E integration test to TE CI/CD (L1). TransformerEngine#2845
• Using HFSDP + FP8 delayed scaling, we get a lot of memory back by turning off master_weights:

# No FusedAdam Master Weights (HFSDP + FP8 Delayed Scaling)
[Rank 0] (after 2 iterations) memory (MB) | allocated: 19595.55 | max allocated: 27022.17 | reserved: 22504.00 | max reserved: 30774.00
[2026-04-03 12:37:38.220805] iteration      100/15258789 | consumed samples:        12800 | elapsed time per iteration (ms): 3263.1 | throughput per GPU (TFLOP/s/GPU): 230.1 | learning rate: 4.915198E-07 | global batch size:   128 | lm loss: 5.473558E+00 | loss scale: 1.0 | grad norm: 9.413 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |

# FusedAdam Master Weights (HFSDP + FP8 Delayed Scaling)
[Rank 0] (after 2 iterations) memory (MB) | allocated: 23425.69 | max allocated: 30852.31 | reserved: 26344.00 | max reserved: 34788.00
[2026-04-03 12:44:12.848377] iteration      100/15258789 | consumed samples:        12800 | elapsed time per iteration (ms): 3149.4 | throughput per GPU (TFLOP/s/GPU): 238.4 | learning rate: 4.915198E-07 | global batch size:   128 | lm loss: 5.476564E+00 | loss scale: 1.0 | grad norm: 9.424 | num zeros: 0 | number of skipped iterations:   0 | number of nan iterations:   0 |

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @mcore-oncall.

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.

[shjwudp]

Overall LGTM, just have a little concern about megatron_fsdp_use_decoupled_grad.

This PR avoids FusedAdam from redundantly creating master weights when M-FSDP already maintains FP32 main weights, which is the correct way to use FusedAdam under M-FSDP (in terms of memory usage).

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[cspades]

Random unit test failure in merge queue: https://github.com/NVIDIA/Megatron-LM/actions/runs/24220276331/job/70712019031

🤷🏻 Sending it back in.

[cspades]

Previous unit test passed: https://github.com/NVIDIA/Megatron-LM/actions/runs/24222278675/job/70717988646

Then something else randomly failed: https://github.com/NVIDIA/Megatron-LM/actions/runs/24222278675/job/70717988616

Rerunning this test to be sure: https://github.com/NVIDIA/Megatron-LM/actions/runs/24222278675/job/70801031886 (✅ 🤷🏻‍♂️)

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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