fix(optimizer): checkpoint native FP32 parameters in mixed precision

megatron/core/optimizer/optimizer.py

@@ -7,7 +7,6 @@
 import math
 import warnings
 from abc import ABC, abstractmethod
-from itertools import chain
 from logging import getLogger
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
@@ -1019,25 +1018,53 @@ def sharded_state_dict(
 
         state_dict = self.state_dict()
 
+        # Optimizer state ids enumerate the inner optimizer params: the fp32 main
+        # copies of float16 params and the native fp32 params, interleaved in the
+        # original param-group order. Yield the model-side param for each inner
+        # param in that order so the ids line up even when both kinds are present.
+        def model_params_in_optimizer_order():
+            for inner_group, float16_group, fp32_group in zip(
+                self.optimizer.param_groups, self.float16_groups, self.fp32_from_fp32_groups
+            ):
+                float16_params = iter(float16_group)
+                fp32_param_ids = {id(param) for param in fp32_group}
+                for param in inner_group['params']:
+                    yield param if id(param) in fp32_param_ids else next(float16_params)
+
         id_to_sharded_param_map = get_param_id_to_sharded_param_map(
-            model_sharded_state_dict, chain.from_iterable(g for g in self.float16_groups)
+            model_sharded_state_dict, model_params_in_optimizer_order()
         )
 
         # Convert fp32_from_fp16_params
         assert len(state_dict['fp32_from_fp16_params']) == len(
             state_dict['optimizer']['param_groups']
         )
+        # State ids of the fp32 main copies only, skipping native fp32 params.
+        float16_param_ids_per_group = []
+        for state_group, inner_group, fp32_group in zip(
+            state_dict['optimizer']['param_groups'],
+            self.optimizer.param_groups,
+            self.fp32_from_fp32_groups,
+        ):
+            fp32_param_ids = {id(param) for param in fp32_group}
+            float16_param_ids_per_group.append(
+                [
+                    param_id
+                    for param_id, param in zip(state_group['params'], inner_group['params'])
+                    if id(param) not in fp32_param_ids
+                ]
+            )
         state_dict['fp32_from_fp16_params'] = [
             [
                 make_sharded_optimizer_tensor(
                     id_to_sharded_param_map[param_id],
                     fp32_param,
                     prefix=f'optimizer.state.fp32_param',
                 )
-                for param_id, fp32_param in zip(state_group['params'], fp32_group)
+                for param_id, fp32_param in zip(param_ids, fp32_group)
             ]
-            for fp32_group, state_group in zip(
-                state_dict['fp32_from_fp16_params'], state_dict['optimizer']['param_groups']
+            for fp32_group, param_ids in zip(
+                state_dict['fp32_from_fp16_params'], float16_param_ids_per_group
             )
         ]
 

tests/unit_tests/dist_checkpointing/test_optimizer.py

@@ -79,6 +79,25 @@ def sharded_state_dict(self):
         return sharded_state_dict
 
 
+class NativeFp32Model(torch.nn.Module):
+    """BF16 model with one parameter kept natively in FP32."""
+
+    def __init__(self):
+        super().__init__()
+        self.pre = torch.nn.Linear(8, 8, bias=False, dtype=torch.bfloat16)
+        self.gate = torch.nn.Parameter(torch.zeros(24, dtype=torch.float32))
+        self.post = torch.nn.Linear(8, 8, bias=False, dtype=torch.bfloat16)
+        self.config = TransformerConfig(
+            hidden_size=8, num_attention_heads=1, num_layers=1, bf16=True
+        )
+
+    def sharded_state_dict(self):
+        return {
+            key: ShardedTensor.from_rank_offsets(key, value)
+            for key, value in self.state_dict(keep_vars=True).items()
+        }
+
+
 class SwigluFactoryModel(torch.nn.Module):
     def __init__(self, pp_separate_model: bool = False):
         super().__init__()
@@ -238,6 +257,51 @@ def test_optimizer_params(self, tmp_path_dist_ckpt):
             ]
         )
 
+    def test_float16_optimizer_with_native_fp32_params(self):
+        """Native FP32 params must retain their optimizer-state ids among BF16 params."""
+        from megatron.core.optimizer import OptimizerConfig
+        from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params
+
+        Utils.initialize_model_parallel(1, 1)
+        model = NativeFp32Model().cuda()
+        assert model.pre.weight.dtype == torch.bfloat16
+        assert model.gate.dtype == torch.float32
+
+        # Force optimizer state initialization.
+        for param in model.parameters():
+            param.grad = torch.zeros_like(param)
+        inner_optim = Adam(model.parameters())
+        inner_optim.step()
+
+        optim = Float16OptimizerWithFloat16Params(
+            inner_optim,
+            OptimizerConfig(optimizer='adam', lr=1e-4, bf16=True),
+            None,
+            lambda opt, cfg: None,
+        )
+        sharded_state_dict = optim.sharded_state_dict(model.sharded_state_dict())
+
+        # FP32 main copies pair with the BF16 params only, in optimizer order.
+        fp32_params = sharded_state_dict['fp32_from_fp16_params'][0]
+        assert [(sharded.key, tuple(sharded.data.shape)) for sharded in fp32_params] == [
+            ('optimizer.state.fp32_param.pre.weight', (8, 8)),
+            ('optimizer.state.fp32_param.post.weight', (8, 8)),
+        ]
+
+        # Per-param state maps every param, including the native FP32 one, to the right key.
+        state = sharded_state_dict['optimizer']['state']
+        # parameters() yields the root module's own params first, then submodules.
+        expected = {0: ('gate', (24,)), 1: ('pre.weight', (8, 8)), 2: ('post.weight', (8, 8))}
+        for param_id, (model_key, shape) in expected.items():
+            for state_key in ('exp_avg', 'exp_avg_sq'):
+                sharded = state[param_id][state_key]
+                assert sharded.key == f'optimizer.state.{state_key}.{model_key}', sharded.key
+                assert tuple(sharded.data.shape) == shape, (
+                    param_id,
+                    sharded.key,
+                    sharded.data.shape,
+                )
+
 
 def initialize_pp_agnostic_model(pre_process=True, post_process=True, seed=0, **config_kwargs):
     torch.manual_seed(seed)