[Megatron-FSDP] Fix segmentation fault when using precision-aware optimizers with the updated MCore.

[cspades]

What does this PR do ?

• Fixes a segmentation fault caused by passing param.decoupled_grad = None into clip_grad_norm -> multi_tensor_apply, which is fixed by setting megatron_fsdp_use_decoupled_grad = use_precision_aware_optimizer as Megatron-FSDP always uses a decoupled gradient when using TE.FusedAdam.
• Requires these PR in MCore which ensure that FusedAdam(use_decoupled_grad) is set appropriately, otherwise the training will not converge (silently without hard errors): [Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP. NVIDIA/Megatron-LM#4427 / [Dev] Fix mis-set decoupled gradient for Megatron-FSDP. NVIDIA/Megatron-LM#4426

# export PYTHONFAULTHANDLER=1
# export CUDA_LAUNCH_BLOCKING=1
Current thread 0x0000fffff7ff47a0 (most recent call first):
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/pytorch/optimizers/multi_tensor_apply.py", line 21 in __call__
  File "/workspace/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/optimizer/clip_grads.py", line 190 in clip_grad_by_total_norm_fp32
  File "/workspace/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/optimizer/optimizer.py", line 232 in clip_grad_norm
  File "/workspace/Megatron-Bridge/3rdparty/Megatron-LM/megatron/core/optimizer/optimizer.py", line 635 in step

Testing

• Tested in MLPerf container on single-node, and removing the megatron_fsdp_use_decoupled_grad config causes the segfault.
• Add test to make sure FusedAdam(use_decoupled_grad), MegatronFSDP.param_and_grad_buffer.use_decoupled_grad, and clip_grad_norm(use_decoupled_grad) are all consistent and produce non-zero gradients.
• Deleted a redundant checkpointing test in favor of save and resume with loss parity validation.
• Had to fork Megatron-FSDP tests due to OOM on downstream test.
  • Didn't add a GB200 version yet, there is currently zero reason to add it, will add it when I add MXFP8 tests.

GitHub Actions CI

See the CI sectionin the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  • Reviewer: Does the PR have correct import guards for all optional libraries?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[cspades]

Either some corrupt global state / config (test OOM'ing on FusedAdam, which shouldn't be used) or something else. Debugging that test...

[cuichenx]

LGTM!

[pruprakash]

QA RCCA Analysis

1. Fix Reference

• PR: [Megatron-FSDP] Fix segmentation fault when using precision-aware optimizers with the updated MCore. #3510 - [Megatron-FSDP] Fix segmentation fault when using precision-aware optimizers with the updated MCore
• Issue: PR is the fix itself (no linked issue)

2. Root Cause

Segmentation fault caused by passing param.decoupled_grad = None into clip_grad_norm -> multi_tensor_apply. Megatron-FSDP always uses a decoupled gradient when using TE.FusedAdam, but the flag wasn't being set correctly.

3. Trigger Configuration

• Using Megatron-FSDP with precision-aware optimizers (FusedAdam)
• Updated MCore version

4. Nature of the Bug

Classification: CODE BUG - Incorrect optimizer configuration causing segfault

5. Existing Test Coverage

In Fix PR: YES - 2 test files added/modified:

• tests/functional_tests/launch_scripts/h100/active/L0_Launch_megatron_fsdp.sh
• tests/functional_tests/test_groups/megatron_fsdp/test_megatron_fsdp.py

In NMFW Tests: Not specifically covering this precision-aware optimizer scenario

6. Coverage Assessment

Test Type	Exists	Covers Bug
Fix PR functional tests	YES	YES
Fix PR unit tests	YES	YES
NMFW regression tests	NO	N/A

7. New Regression Test

NOT NEEDED - Fix PR includes functional tests for Megatron-FSDP with precision-aware optimizers.

8. Conclusion

Verdict: ADEQUATE COVERAGE - Fix PR includes comprehensive functional tests.