From 8bb9bcd953083fb405f7fbc95be5d51cb5a21eb0 Mon Sep 17 00:00:00 2001 From: Li Ding Date: Thu, 30 Apr 2026 16:55:18 -0700 Subject: [PATCH 1/2] fix mimo optimizer checkpoint metadata restore Signed-off-by: Li Ding --- megatron/core/models/mimo/optimizer.py | 32 +++++++++++++++++++++++++- 1 file changed, 31 insertions(+), 1 deletion(-) diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py index 1a79c1f91ff..b6977818556 100644 --- a/megatron/core/models/mimo/optimizer.py +++ b/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) @@ -253,7 +269,21 @@ def _restore_param_groups(sub_sd, inner_optimizer, module_name): ) for loaded_g, current_g in zip(loaded_pg, current_pg): loaded_g['params'] = current_g['params'] - sub_sd['optimizer']['param_groups'] = loaded_pg + # In MIMO, rank-local module optimizer metadata is not common across ranks + # (for example, non-colocated rank 0 may own language while rank 1 owns + # vision). Distributed checkpoint load can therefore return the sharded + # tensor state plus the extracted MIMO param groups, but without the + # original nested "optimizer" common-state wrapper. Recreate it here; the + # inner optimizer load path only needs param_groups from this wrapper. + 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): From 0aab8aa265f09f59430dd99fef8b6dd88d137f8d Mon Sep 17 00:00:00 2001 From: Li Ding Date: Fri, 29 May 2026 08:48:41 -0700 Subject: [PATCH 2/2] tests Signed-off-by: Li Ding --- .../unit_tests/models/test_mimo_checkpoint.py | 95 +++++++++++++++++++ 1 file changed, 95 insertions(+) diff --git a/tests/unit_tests/models/test_mimo_checkpoint.py b/tests/unit_tests/models/test_mimo_checkpoint.py index 3dc75a05a87..06cfdc07871 100644 --- a/tests/unit_tests/models/test_mimo_checkpoint.py +++ b/tests/unit_tests/models/test_mimo_checkpoint.py @@ -271,3 +271,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