Add fully_shard_optimizer for mixed-precision FSDP

[wujingyue]

Summary

• Add fully_shard_optimizer(optimizer) and export it for the experimental fully_shard path.
• Adapt optimizer steps in place with pre/post hooks that temporarily cast FSDP-managed grads while the optimizer consumes them, then restore the original grad tensors.
• Reject optimizer step closures because the hook-only adapter cannot intercept grads materialized inside the closure.
• Add focused MFSDP v2 optimizer coverage for mixed param/grad precision:
  • raw torch.optim.Adam fails on the mismatch without adaptation
  • TE FusedAdam accepts mismatched grads and updates parameters
  • fully_shard_optimizer covers Adam mixed-precision stepping through the adapter

Notes

• DistributedOptimizer is intentionally not used for the MFSDP v2 test. With use_megatron_fsdp=True, it delegates to the wrapped optimizer for this path; with False, it expects the older _ParamAndGradBuffer setup rather than v2 FsdpParameterGroup/DBuffer state.
• A follow-up can extend the adapter to sync main_weight to model_weight from a post-step hook. This PR keeps the current per-microbatch sync behavior.

Testing

• BASE_REF=main CHECK_ONLY=true SKIP_DOCS=false bash tools/autoformat.sh
• uv run --no-sync python -m torch.distributed.run --nproc-per-node 1 -m pytest -q tests/unit_tests/distributed/megatron_fsdp/test_optimizer.py --capture=fd --tb=short --disable-warnings -rN (2 passed)
• uv run --no-sync python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/megatron_fsdp/test_optimizer.py --capture=fd --tb=short --disable-warnings -rN (2 passed on each rank)
• uv run --no-sync python -m torch.distributed.run --nproc-per-node 2 -m pytest -q tests/unit_tests/distributed/megatron_fsdp/test_experimental_fully_shard.py -k "adam or optimizer" --capture=fd --tb=short --disable-warnings -rN (2 passed, 8 deselected on each rank)

[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.

[copy-pr-bot]

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.

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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