fix(training): realign the distributed optimizer's parameter index with its group order

[kamran-nvidia]

What does this PR do ?

Works around a Megatron-Core DistributedOptimizer bug that makes optimizer checkpoint saves fail when an optimizer group mixes FP32 and 16-bit trainable parameters. setup_optimizer now rebuilds the optimizer's parameter index in the order the group actually stores its shards.

The bug. _build_optimizer_group_ranges records each parameter's position in its optimizer group in gradient-buffer order. _build_model_and_main_param_groups then stores the group as all FP32 shards followed by all 16-bit shards. When a 16-bit buffer comes before an FP32 one, model_param_group_index_map points at other parameters' shards. FP32 parameters typically come from mark_keep_in_fp32, which MCore uses for hyper-connection, CSA, and MTP-head parameters.

• Training is unaffected, because the optimizer step uses param_groups directly.
• Saving optimizer state fails in sharded_param_state_dp_reshardable, for example with AssertionError: (torch.Size([8192]), 8192, 12288).
• Loading could silently restore state onto the wrong parameters wherever sizes happen to agree.

The same code is on Megatron-LM main today. This workaround is meant to stay until an upstream fix lands; it leaves already-aligned maps untouched, so it then becomes a no-op.

Changelog

• training/optim.py: _realign_distributed_optimizer_param_index, called at the end of setup_optimizer, rebuilds model_param_group_index_map for every DistributedOptimizer, including inside a ChainedOptimizer. It logs the number of corrected entries and leaves the map unchanged if MCore's groups and index ever disagree on the parameter set.
• tests/unit_tests/training/test_optim.py:
  • Stand-in optimizer tests: a mixed-dtype group, an already-aligned map, chained and non-distributed optimizers, and unknown parameters.
  • A one-GPU test with a real DDP model and a BF16 buffer ahead of an FP32 one. After an Adam step it checks that every parameter's optimizer state has its own size and values, and it builds the dp_reshardable checkpoint state.

Testing

• On MCore e30d863a (the current pin), GB200: the 5 new tests pass, and the rest of test_optim.py passes.
• With the realignment disabled, the GPU test fails with AssertionError on the param state, so it reproduces the bug.
• On a large MoE model with FP32 router weights and residual gates, the step-2 full-state save that previously failed now succeeds. A run resumed from it matches the uninterrupted run at steps 3 and 4 within 0.0003%, with identical learning rates.
• Ruff lint and format are clean.

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? (internal workaround, no user-facing change)
• Does the PR affect components that are optional to install? (no)

Additional Information

• The proper fix belongs in Megatron-Core (megatron/core/optimizer/distrib_optimizer.py); an upstream PR will follow.

🤖 Generated with Claude Code

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

[kamran-nvidia]

Upstream fix in NVIDIA/Megatron-LM#7336