fix mimo optimizer checkpoint metadata restore

megatron/core/models/mimo/optimizer.py

@@ -153,6 +153,7 @@ def load_state_dict(self, state_dict: Dict):
 
             for sub_sd, inner_opt in _iter_optimizer_sub_dicts(module_sd, info.optimizer):
                 _restore_param_groups(sub_sd, inner_opt, name)
+                _restore_param_state_sharding_type(sub_sd)
                 _restore_grad_scaler(sub_sd)
 
             info.optimizer.load_state_dict(module_sd)
@@ -175,6 +176,7 @@ def sharded_state_dict(self, model_sharded_state_dict, is_loading: bool = False,
                 ):
                     suffix = f'.{idx}' if idx > 0 else ''
                     _extract_param_groups(sub_sd, name, suffix, replica_id)
+                    _extract_param_state_sharding_type(sub_sd, name, suffix, replica_id)
                     _extract_grad_scaler(sub_sd, name, suffix, replica_id)
 
                 sharded_state[name] = module_sd
@@ -218,6 +220,8 @@ def _extract_param_groups(sub_sd, module_name, suffix, replica_id):
             replica_id=replica_id,
         )
         del opt_sub['param_groups']
+        if not opt_sub:
+            del sub_sd['optimizer']
 
 
 def _extract_grad_scaler(sub_sd, module_name, suffix, replica_id):
@@ -232,6 +236,18 @@ def _extract_grad_scaler(sub_sd, module_name, suffix, replica_id):
         )
 
 
+def _extract_param_state_sharding_type(sub_sd, module_name, suffix, replica_id):
+    """Save: extract param_state_sharding_type into a ShardedObject."""
+    if 'param_state_sharding_type' in sub_sd:
+        sub_sd[f'_mimo_param_state_sharding_type{suffix}'] = ShardedObject(
+            f'optimizer.mimo.{module_name}{suffix}.param_state_sharding_type',
+            sub_sd.pop('param_state_sharding_type'),
+            (1,),
+            (0,),
+            replica_id=replica_id,
+        )
+
+
 def _restore_param_groups(sub_sd, inner_optimizer, module_name):
     """Load: restore param_groups with current param IDs from the inner optimizer."""
     # Find the _mimo_param_groups key (may have a suffix for chained optimizers)
@@ -262,6 +278,14 @@ def _restore_param_groups(sub_sd, inner_optimizer, module_name):
     sub_sd.setdefault('optimizer', {})['param_groups'] = loaded_pg
 
 
+def _restore_param_state_sharding_type(sub_sd):
+    """Load: restore param_state_sharding_type from ShardedObject key."""
+    for k in list(sub_sd.keys()):
+        if k.startswith('_mimo_param_state_sharding_type'):
+            sub_sd['param_state_sharding_type'] = sub_sd.pop(k)
+            break
+
+
 def _restore_grad_scaler(sub_sd):
     """Load: restore grad_scaler from ShardedObject key."""
     for k in list(sub_sd.keys()):

tests/unit_tests/models/test_mimo_checkpoint.py

@@ -274,3 +274,98 @@ def test_encoder_tp2_pp2_llm_tp2_pp2(self):
             hidden_size=256,
             num_layers=2,
         )
+
+
+class TestOptimizerCheckpointHelpers:
+    """CPU-only coverage for the dist-checkpoint extract/restore helpers."""
+
+    @staticmethod
+    def _extract_param_state_sharding_type(*args, **kwargs):
+        from megatron.core.models.mimo.optimizer import _extract_param_state_sharding_type
+
+        return _extract_param_state_sharding_type(*args, **kwargs)
+
+    @staticmethod
+    def _restore_param_state_sharding_type(*args, **kwargs):
+        from megatron.core.models.mimo.optimizer import _restore_param_state_sharding_type
+
+        return _restore_param_state_sharding_type(*args, **kwargs)
+
+    @staticmethod
+    def _extract_param_groups(*args, **kwargs):
+        from megatron.core.models.mimo.optimizer import _extract_param_groups
+
+        return _extract_param_groups(*args, **kwargs)
+
+    @staticmethod
+    def _restore_param_groups(*args, **kwargs):
+        from megatron.core.models.mimo.optimizer import _restore_param_groups
+
+        return _restore_param_groups(*args, **kwargs)
+
+    def test_extract_param_state_sharding_type_wraps_value_into_sharded_object(self):
+        from megatron.core.dist_checkpointing.mapping import ShardedObject
+
+        sub_sd = {'param_state_sharding_type': 'fully_sharded_bucket_space'}
+
+        self._extract_param_state_sharding_type(sub_sd, 'images', '.1', replica_id=(0, 0, 0))
+
+        assert 'param_state_sharding_type' not in sub_sd
+        wrapped = sub_sd['_mimo_param_state_sharding_type.1']
+        assert isinstance(wrapped, ShardedObject)
+        assert wrapped.key == 'optimizer.mimo.images.1.param_state_sharding_type'
+        assert wrapped.data == 'fully_sharded_bucket_space'
+        assert wrapped.replica_id == (0, 0, 0)
+
+    def test_extract_param_state_sharding_type_noop_when_missing(self):
+        sub_sd = {'unrelated': 1}
+
+        self._extract_param_state_sharding_type(sub_sd, 'images', '', replica_id=0)
+
+        assert sub_sd == {'unrelated': 1}
+
+    def test_restore_param_state_sharding_type_renames_suffixed_key(self):
+        sub_sd = {'_mimo_param_state_sharding_type.0': 'fully_sharded_bucket_space'}
+
+        self._restore_param_state_sharding_type(sub_sd)
+
+        assert sub_sd == {'param_state_sharding_type': 'fully_sharded_bucket_space'}
+
+    def test_restore_param_state_sharding_type_noop_when_missing(self):
+        sub_sd = {'unrelated': 1}
+
+        self._restore_param_state_sharding_type(sub_sd)
+
+        assert sub_sd == {'unrelated': 1}
+
+    def test_extract_param_groups_deletes_empty_optimizer_dict(self):
+        sub_sd = {'optimizer': {'param_groups': [{'lr': 0.1, 'params': [0]}]}}
+
+        self._extract_param_groups(sub_sd, 'images', '', replica_id=0)
+
+        assert 'optimizer' not in sub_sd
+        assert '_mimo_param_groups' in sub_sd
+
+    def test_extract_param_groups_keeps_optimizer_when_other_keys_remain(self):
+        sub_sd = {
+            'optimizer': {'param_groups': [{'lr': 0.1, 'params': [0]}], 'state': {0: {'step': 5}}}
+        }
+
+        self._extract_param_groups(sub_sd, 'images', '', replica_id=0)
+
+        assert sub_sd['optimizer'] == {'state': {0: {'step': 5}}}
+        assert '_mimo_param_groups' in sub_sd
+
+    def test_restore_param_groups_recreates_missing_optimizer_wrapper(self):
+        from unittest.mock import MagicMock
+
+        inner_optimizer = MagicMock()
+        inner_optimizer.optimizer.state_dict.return_value = {
+            'param_groups': [{'lr': 0.1, 'params': [42, 43]}]
+        }
+        sub_sd = {'_mimo_param_groups': [{'lr': 0.1, 'params': []}]}
+
+        self._restore_param_groups(sub_sd, inner_optimizer, 'images')
+
+        assert sub_sd['optimizer']['param_groups'][0]['params'] == [42, 43]
+        assert '_mimo_param_groups' not in sub_sd