[BUG] DistributedOptimizer issues in Mixed-Precision: Stale model_param_group_index_map and .copy_ failure during get/load parameter_state

[HelloWorld686]

Describe the bug

During mixed-precision training (BF16 & FP32, FP8 is optional), a RuntimeError (size mismatch) is triggered when saving a checkpoint via get_parameter_state_dp_zero.

Based on our analysis, the hardcoded parameter reordering in _build_model_and_main_param_groups within DistributedOptimizer causes the self.model_param_group_index_map to become out of sync with the actual optimizer.param_groups.

Steps/Code to reproduce bug

1. Before building the DDP model and optimizer，use Float16Module to wrap the model for BF16 training, but manually promote certain modules (both params and inputs) to FP32.
2. Train for several steps.
3. Call get_parameter_state_dp_zero in DistributedOptimizer to collect optimizer states, which triggers the size mismatch error.

Root Cause Analysis

1. Initial Map Construction: In __init__, self.model_param_group_index_map is first constructed via _build_optimizer_group_ranges. This map records the position (group_index, group_order) of param in param_groups

2. Hardcoded Reordering: Subsequently, _build_model_and_main_param_groups reorders the parameters within each group (placing native FP32 shards at the front and main parameter shards converted from FP16/BF16 at the back) and updates optimizer.param_groups accordingly:

3. Index Invalidation: The model_param_group_index_map is not updated after this reordering. Consequently, downstream functions like _get_main_param_and_optimizer_states retrieve the wrong Tensors using stale group_order indices, leading to shape mismatches during buffer copy operations.

Additional question

1. What is the design motivation behind this specific reordering (grouping by DType)?
2. What is the recommended best practice to fix this: disabling the reordering to maintain discovery order consistency, or explicitly updating the model_param_group_index_map after the reordering is performed?

[HelloWorld686]

Regarding the indexing inconsistency discussed above, we have implemented and verified two potential fix approaches internally (Modified the _build_model_and_main_param_groups method of the megatron/core/optimizer/distrib_optimizer.py file). We would like to share them here for feedback:

• Option A: Maintain Original Order (Remove Reordering) Keep the parameter order in param_groups consistent with the initial discovery order.

• Option B: Update Index Map After Reordering Keep the current DType-based grouping logic but rebuild or update model_param_group_index_map immediately after the reordering is finalized.

[HelloWorld686]

Another bug found during checkpoint loading with mixed-precision training

During further testing of the mixed-precision training checkpoint flow, we identified another bug related to state restoration. When resuming from a checkpoint, the call to _set_main_param_and_optimizer_states triggers the following error:

Root Cause Analysis: The error originates may from parameters in shard_fp32_groups. In the _build_model_and_main_param_groups method, the code assigns shard_model_param directly without calling .detach(). Since these tensors remain part of the autograd graph with requires_grad=True, calling .copy_ directly on them during state loading is prohibited by PyTorch's safety checks.

Proposed Fix: We have temporarily resolved this by using .data.copy_() instead of .copy_() to bypass the autograd check.

[ahmadki]

Thanks for the detailed analysis @HelloWorld686

The deterministic dtype reordering is important for checkpoint consistency and correctness, without it, parameters could appear in different positions between checkpoint save and load.

This PR should fix the first issue, I'm working on a fix for the second.

[dylanzhj]

The order of parameter positions between checkpoint saving and loading depends on their sequence in the param_group. Currently, Megatron updates the param_group, while the torch optimizer resets the param IDs. Therefore, there should be no issue with saving/loading. Retrieving the master weight/averages and other states of parameters relies on the gbuf_range, which is determined during model initialization and should remain unchanged. checkpoint consistency and correctness？