From 71e12651903d2737dc1f2a2278bf688b78e2ab7d Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Thu, 14 May 2026 09:38:52 -0700 Subject: [PATCH 1/4] Enhance MimoOptimizer and tests for distributed checkpointing support - Add handling for sharded metadata recovery in MimoOptimizer. - Improve optimizer state restoration logic to accommodate DistributedOptimizer. - Update test cases to validate optimizer step continuity and structural integrity for distributed setups. Signed-off-by: Kamran Jafari --- megatron/core/models/mimo/optimizer.py | 39 +++- .../unit_tests/models/test_mimo_checkpoint.py | 178 ++++++++++++++++-- 2 files changed, 195 insertions(+), 22 deletions(-) diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py index 1a79c1f91ff..988bf0fe558 100644 --- a/megatron/core/models/mimo/optimizer.py +++ b/megatron/core/models/mimo/optimizer.py @@ -6,7 +6,7 @@ from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import torch @@ -48,6 +48,11 @@ def __init__(self, module_infos: Dict[str, ModuleOptimizerInfo], config: Optimiz ] self.is_stub_optimizer = len(self._active_optimizers) == 0 self.optimizer = None # Base class compat + # Stashed by `sharded_state_dict` so `load_state_dict` can recover + # non-sharded scalars (notably `param_state_sharding_type`) that + # dist_checkpointing's common-state path drops on ranks whose active + # module isn't present on rank 0. + self._last_sharded_metadata: Dict[str, Any] = {} @torch.no_grad() def prepare_grads(self) -> bool: @@ -144,6 +149,12 @@ def load_state_dict(self, state_dict: Dict): as ShardedObjects by sharded_state_dict(), then delegates to each per-module optimizer's load_state_dict. """ + # The sharding-type metadata isn't saved to the checkpoint — it's a + # load-time interpretation hint that the caller supplies via + # the metadata kwarg on the most recent sharded_state_dict call. We + # recover it here for ranks whose active module wasn't on rank 0 at + # save time (and so dropped it via the common-state path). + recovered_sharding_type = self._last_sharded_metadata.get('distrib_optim_sharding_type') for name, info in self.module_infos.items(): if not (info.is_active and info.optimizer): continue @@ -154,6 +165,12 @@ 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_grad_scaler(sub_sd) + if ( + recovered_sharding_type is not None + and 'param_state' in sub_sd + and 'param_state_sharding_type' not in sub_sd + ): + sub_sd['param_state_sharding_type'] = recovered_sharding_type info.optimizer.load_state_dict(module_sd) @@ -162,6 +179,8 @@ def sharded_state_dict(self, model_sharded_state_dict, is_loading: bool = False, through distributed save as ShardedObjects (common.pt is rank-0 only, which misses LLM optimizer state in non-colocated mode). """ + # Stash for load_state_dict + self._last_sharded_metadata = dict(kwargs.get('metadata', {}) or {}) sharded_state = {} for name, info in self.module_infos.items(): if info.is_active and info.optimizer: @@ -253,7 +272,13 @@ 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 + # `sub_sd['optimizer']` may be absent on load: when the per-module state_dict + # produced by DistributedOptimizer.state_dict() only contains `param_groups` + # under the 'optimizer' key, `_extract_param_groups` removes it at save time + # and the resulting empty dict can be dropped during dist_checkpointing + # common-state save/load. Use setdefault so the restored param_groups land + # in the right place regardless. + sub_sd.setdefault('optimizer', {})['param_groups'] = loaded_pg def _restore_grad_scaler(sub_sd): @@ -267,17 +292,21 @@ def _restore_grad_scaler(sub_sd): def _get_replica_id(pg_collection: Optional[ProcessGroupCollection]) -> tuple: """Build replica_id tuple for ShardedObject deduplication. - Includes pp_rank so only one PP stage writes the metadata, - and dp_rank so only dp_rank=0 writes (others are replicas). + Returns (tp_rank, pp_rank, dp_rank) so only (0, 0, 0) within each + module's parallelism group is the main replica; all other ranks + in the same module are non-main replicas of the same object. """ assert pg_collection is not None, "pg_collection required for checkpoint replica_id" + assert ( + hasattr(pg_collection, 'tp') and pg_collection.tp is not None + ), "pg_collection.tp must be set for checkpoint deduplication" assert ( hasattr(pg_collection, 'pp') and pg_collection.pp is not None ), "pg_collection.pp must be set for checkpoint deduplication" assert ( hasattr(pg_collection, 'dp') and pg_collection.dp is not None ), "pg_collection.dp must be set for checkpoint deduplication" - return (0, pg_collection.pp.rank(), pg_collection.dp.rank()) + return (pg_collection.tp.rank(), pg_collection.pp.rank(), pg_collection.dp.rank()) def _get_pg_collection_for_optimizer(grid) -> ProcessGroupCollection: diff --git a/tests/unit_tests/models/test_mimo_checkpoint.py b/tests/unit_tests/models/test_mimo_checkpoint.py index 3dc75a05a87..1ccaf33c0e2 100644 --- a/tests/unit_tests/models/test_mimo_checkpoint.py +++ b/tests/unit_tests/models/test_mimo_checkpoint.py @@ -56,8 +56,25 @@ def _randomize_params(model, seed): p.random_() -def _create_model_and_optimizer(encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed): - """Create MIMO model with DDP + optimizer, do a fake step to populate optimizer state. +def _create_model_and_optimizer( + encoder_grid, + llm_grid, + hidden_size, + num_layers, + vocab_size, + seed, + use_distributed_optimizer=False, +): + """Create MIMO model with DDP + optimizer. + + When `use_distributed_optimizer=False` (Float16Optimizer), take a few fake + backward + step iterations so the Adam step counter is non-trivial. The + DistributedOptimizer path requires DDP-driven grad reduce-scatter into the + gradient buffer to take a step, which isn't worth wiring up here — the + parametrization that uses DistributedOptimizer is a structural regression + guard (it checks that the save/load code paths in MimoOptimizer don't crash + when run against a `DistributedOptimizer` + `fully_reshardable` setup), + not a step-continuity test. Caller must call create_all_embedding_groups() before this function. """ @@ -74,27 +91,72 @@ def _create_model_and_optimizer(encoder_grid, llm_grid, hidden_size, num_layers, ) _randomize_params(mimo_model, seed) - # Use Float16Optimizer (not DistributedOptimizer) to exercise the MIMO-specific - # param_groups/grad_scaler extraction in sharded_state_dict. DistributedOptimizer - # handles its own checkpointing internally and our code is transparent to it. opt_config = OptimizerConfig( optimizer='adam', lr=1e-4, weight_decay=0.01, clip_grad=1.0, bf16=True, - use_distributed_optimizer=False, + use_distributed_optimizer=use_distributed_optimizer, ) optimizer = get_mimo_optimizer(mimo_model, opt_config) - # Fake backward + step to populate optimizer state (Adam m/v) - for param in mimo_model.parameters(): - param.grad = torch.randn_like(param) - optimizer.step() + if not use_distributed_optimizer: + # Float16Optimizer path: take several fake backward + step iterations + # so we can later assert Adam step-counter continuity. + for _ in range(3): + for param in mimo_model.parameters(): + param.grad = torch.randn_like(param) + optimizer.step() + else: + # DistributedOptimizer path: a real step requires DDP-driven grad + # reduce-scatter into the gradient buffer. We don't need actual step + # values for this regression test (its purpose is to exercise the + # save/load *structure* — see `test_distributed_optimizer_fully_reshardable`), + # so just initialize the inner torch optimizer's state buffers + # (`exp_avg`, `exp_avg_sq`) with zeros via `init_state_fn`. Without + # this, `get_parameter_state_dp_zero` later fails with + # `KeyError: 'exp_avg'` when it tries to read state for each param. + # + # Each per-module optimizer is wrapped in a `ChainedOptimizer` (which + # holds N>=1 DistributedOptimizers); init_state_fn lives on the inner + # ones, not on the chain wrapper. + from megatron.core.optimizer.optimizer import ChainedOptimizer + + for info in optimizer.module_infos.values(): + if not (info.is_active and info.optimizer is not None): + continue + inner_opts = ( + info.optimizer.chained_optimizers + if isinstance(info.optimizer, ChainedOptimizer) + else [info.optimizer] + ) + for inner in inner_opts: + inner.init_state_fn(inner.optimizer, inner.config) return mimo_model, optimizer +def _find_optimizer_step(state_dict): + """Return the optimizer step counter from a per-module state_dict, regardless of + whether it lives in `state['common_step']`, `param_groups[*]['step']`, or per-param state. + """ + opt = state_dict.get('optimizer', {}) + state = opt.get('state', {}) + if 'common_step' in state: + val = state['common_step'] + return int(val.item()) if torch.is_tensor(val) else int(val) + for g in opt.get('param_groups', []): + if 'step' in g: + val = g['step'] + return int(val.item()) if torch.is_tensor(val) else int(val) + for s in state.values(): + if isinstance(s, dict) and 'step' in s: + val = s['step'] + return int(val.item()) if torch.is_tensor(val) else int(val) + return None + + def run_checkpoint_test( encoder_tp, encoder_pp, @@ -107,8 +169,16 @@ def run_checkpoint_test( hidden_size=256, num_layers=2, vocab_size=1000, + use_distributed_optimizer=False, + distrib_optim_sharding_type='fully_reshardable', ): - """Save model + optimizer checkpoint, load into fresh instances, verify match.""" + """Save model + optimizer checkpoint, load into fresh instances, verify match. + + `use_distributed_optimizer` selects between Float16Optimizer (default; step + continuity is asserted) and DistributedOptimizer (structural regression + guard for the load-side fixes — `_restore_param_groups` setdefault and + `param_state_sharding_type` recovery from stashed metadata). + """ # Clear NVTE env vars that the conftest set_env fixture sets to '0'. # GPTModel (LanguageModule) asserts these are unset or match the attention backend. os.environ.pop('NVTE_FLASH_ATTN', None) @@ -121,9 +191,21 @@ def run_checkpoint_test( llm_grid = create_hypercomm_grid(offset=llm_offset, tp=llm_tp, cp=1, pp=llm_pp, dp=llm_dp) create_all_embedding_groups([encoder_grid, llm_grid]) + optim_metadata = ( + {'distrib_optim_sharding_type': distrib_optim_sharding_type} + if use_distributed_optimizer + else {} + ) + # --- Create model A + optimizer, snapshot state --- model_a, optimizer_a = _create_model_and_optimizer( - encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed=1 + encoder_grid, + llm_grid, + hidden_size, + num_layers, + vocab_size, + seed=1, + use_distributed_optimizer=use_distributed_optimizer, ) params_a = {name: p.clone() for name, p in model_a.named_parameters()} @@ -139,15 +221,26 @@ def run_checkpoint_test( # Save model save(model_a.sharded_state_dict(), model_ckpt) - # Save optimizer (needs fresh model sharded_state_dict since save() consumes tensor refs) - optim_sd_a = optimizer_a.sharded_state_dict(model_a.sharded_state_dict(), is_loading=False) - save(optim_sd_a, optim_ckpt, validate_access_integrity=False) + # Save optimizer (needs fresh model sharded_state_dict since save() consumes tensor refs). + # validate_access_integrity=True is the regression guard for the _get_replica_id fix: + # without including TP rank in replica_id, every TP rank at (pp=0, dp=0) would emit + # the same `_mimo_*` ShardedObject as a main replica, producing duplicate-key errors. + optim_sd_a = optimizer_a.sharded_state_dict( + model_a.sharded_state_dict(), is_loading=False, metadata=optim_metadata + ) + save(optim_sd_a, optim_ckpt, validate_access_integrity=True) dist.barrier() # --- Create model B + optimizer with different weights (reuse same grids) --- model_b, optimizer_b = _create_model_and_optimizer( - encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed=2 + encoder_grid, + llm_grid, + hidden_size, + num_layers, + vocab_size, + seed=2, + use_distributed_optimizer=use_distributed_optimizer, ) # Load model @@ -162,7 +255,9 @@ def run_checkpoint_test( model_b.load_state_dict(loaded_model_sd) # Load optimizer - optim_sd_b = optimizer_b.sharded_state_dict(model_b.sharded_state_dict(), is_loading=True) + optim_sd_b = optimizer_b.sharded_state_dict( + model_b.sharded_state_dict(), is_loading=True, metadata=optim_metadata + ) loaded_optim_sd = load(optim_sd_b, optim_ckpt, validate_access_integrity=False) optimizer_b.load_state_dict(loaded_optim_sd) @@ -201,6 +296,19 @@ def run_checkpoint_test( state_a[param_id][key], state_b[param_id][key] ), f"Optimizer {name} param {param_id} {key} mismatch" + # Verify Adam step counter survives save/load. Without restoring step, + # bias correction would reset and cause an effective LR drop on resume. + # Only meaningful when real optimizer steps were taken (the + # DistributedOptimizer parametrization skips steps — see + # _create_model_and_optimizer). + if not use_distributed_optimizer: + step_a = _find_optimizer_step(sd_a) + step_b = _find_optimizer_step(sd_b) + assert step_a is not None, f"Optimizer {name}: could not locate step in model A" + assert step_a == step_b, ( + f"Optimizer {name}: step mismatch after load " f"(A={step_a}, B={step_b})" + ) + finally: _cleanup_tmpdir(ckpt_dir) @@ -271,3 +379,39 @@ def test_encoder_tp2_pp2_llm_tp2_pp2(self): hidden_size=256, num_layers=2, ) + + def test_distributed_optimizer_fully_reshardable(self): + """Regression guard for the DistributedOptimizer + fully_reshardable load path. + + Exercises the exact code path the hetero LLaVA trainer hits and which + the Float16Optimizer-based parametrizations above don't cover: + + - `DistributedOptimizer.state_dict()` puts only `param_groups` under + the 'optimizer' key (state is sharded separately as `param_state`). + After `_extract_param_groups` deletes `param_groups`, what's left is + an empty dict that is dropped through dist_checkpointing's common- + state round-trip on ranks where the active module isn't on rank 0. + `_restore_param_groups` must use `setdefault` to survive this. + - `param_state_sharding_type` is a plain scalar that also doesn't + survive the common-state round-trip on those ranks. `load_state_dict` + must recover it from the metadata stash set by `sharded_state_dict`. + + Uses an 8-GPU non-colocated layout where rank 0 is NOT in the language + module, so the dropped-on-rank-0 codepath is forced. + """ + if self.world_size != 8: + pytest.skip(f"Requires 8 GPUs, got {self.world_size}") + run_checkpoint_test( + encoder_tp=2, + encoder_pp=1, + encoder_dp=1, + encoder_offset=0, + llm_tp=2, + llm_pp=3, + llm_dp=1, + llm_offset=2, + hidden_size=256, + num_layers=3, + use_distributed_optimizer=True, + distrib_optim_sharding_type='fully_reshardable', + ) From 7edaa8c3a3f0d51379e0368d0c13f174f7adc34d Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Fri, 15 May 2026 13:35:16 -0700 Subject: [PATCH 2/4] Refactor MimoOptimizer to enhance state management for param_state_sharding_type in distributed checkpointing Signed-off-by: Kamran Jafari --- megatron/core/models/mimo/optimizer.py | 63 +++++++++++++++----------- 1 file changed, 37 insertions(+), 26 deletions(-) diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py index 988bf0fe558..6981afe306a 100644 --- a/megatron/core/models/mimo/optimizer.py +++ b/megatron/core/models/mimo/optimizer.py @@ -6,7 +6,7 @@ from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -48,11 +48,6 @@ def __init__(self, module_infos: Dict[str, ModuleOptimizerInfo], config: Optimiz ] self.is_stub_optimizer = len(self._active_optimizers) == 0 self.optimizer = None # Base class compat - # Stashed by `sharded_state_dict` so `load_state_dict` can recover - # non-sharded scalars (notably `param_state_sharding_type`) that - # dist_checkpointing's common-state path drops on ranks whose active - # module isn't present on rank 0. - self._last_sharded_metadata: Dict[str, Any] = {} @torch.no_grad() def prepare_grads(self) -> bool: @@ -145,16 +140,10 @@ def state_dict(self): def load_state_dict(self, state_dict: Dict): """Load per-module optimizer state dicts. - Reassembles param_groups and grad_scaler that were extracted and saved - as ShardedObjects by sharded_state_dict(), then delegates to each - per-module optimizer's load_state_dict. + Reassembles param_groups, grad_scaler, and param_state_sharding_type + that were extracted and saved as ShardedObjects by sharded_state_dict(), + then delegates to each per-module optimizer's load_state_dict. """ - # The sharding-type metadata isn't saved to the checkpoint — it's a - # load-time interpretation hint that the caller supplies via - # the metadata kwarg on the most recent sharded_state_dict call. We - # recover it here for ranks whose active module wasn't on rank 0 at - # save time (and so dropped it via the common-state path). - recovered_sharding_type = self._last_sharded_metadata.get('distrib_optim_sharding_type') for name, info in self.module_infos.items(): if not (info.is_active and info.optimizer): continue @@ -164,23 +153,17 @@ 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) - if ( - recovered_sharding_type is not None - and 'param_state' in sub_sd - and 'param_state_sharding_type' not in sub_sd - ): - sub_sd['param_state_sharding_type'] = recovered_sharding_type info.optimizer.load_state_dict(module_sd) def sharded_state_dict(self, model_sharded_state_dict, is_loading: bool = False, **kwargs): - """Build sharded state dict, routing param_groups and grad_scaler - through distributed save as ShardedObjects (common.pt is rank-0 only, - which misses LLM optimizer state in non-colocated mode). + """Build sharded state dict, routing param_groups, grad_scaler, and + param_state_sharding_type through distributed save as ShardedObjects + (common.pt is rank-0 only, which misses non-colocated LLM optimizer + state). """ - # Stash for load_state_dict - self._last_sharded_metadata = dict(kwargs.get('metadata', {}) or {}) sharded_state = {} for name, info in self.module_infos.items(): if info.is_active and info.optimizer: @@ -194,6 +177,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 @@ -237,6 +221,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): @@ -251,6 +237,23 @@ 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. + + Plain non-tensor scalars at the per-module level otherwise travel through + dist_checkpointing's common-state path (rank 0 only), so for non-colocated + MIMO they are lost on ranks whose module is inactive on rank 0. + """ + 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) @@ -289,6 +292,14 @@ def _restore_grad_scaler(sub_sd): break +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 _get_replica_id(pg_collection: Optional[ProcessGroupCollection]) -> tuple: """Build replica_id tuple for ShardedObject deduplication. From 2f7158a3e67222df58f87a00e5d7d1fa6f5080aa Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Mon, 18 May 2026 19:14:43 -0700 Subject: [PATCH 3/4] revert PR #4791 fixes Signed-off-by: Kamran Jafari --- megatron/core/models/mimo/optimizer.py | 42 +---- .../unit_tests/models/test_mimo_checkpoint.py | 156 +++--------------- 2 files changed, 29 insertions(+), 169 deletions(-) diff --git a/megatron/core/models/mimo/optimizer.py b/megatron/core/models/mimo/optimizer.py index 6981afe306a..09b230e1557 100644 --- a/megatron/core/models/mimo/optimizer.py +++ b/megatron/core/models/mimo/optimizer.py @@ -140,9 +140,9 @@ def state_dict(self): def load_state_dict(self, state_dict: Dict): """Load per-module optimizer state dicts. - Reassembles param_groups, grad_scaler, and param_state_sharding_type - that were extracted and saved as ShardedObjects by sharded_state_dict(), - then delegates to each per-module optimizer's load_state_dict. + Reassembles param_groups and grad_scaler that were extracted and saved + as ShardedObjects by sharded_state_dict(), then delegates to each + per-module optimizer's load_state_dict. """ for name, info in self.module_infos.items(): if not (info.is_active and info.optimizer): @@ -153,16 +153,14 @@ 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) def sharded_state_dict(self, model_sharded_state_dict, is_loading: bool = False, **kwargs): - """Build sharded state dict, routing param_groups, grad_scaler, and - param_state_sharding_type through distributed save as ShardedObjects - (common.pt is rank-0 only, which misses non-colocated LLM optimizer - state). + """Build sharded state dict, routing param_groups and grad_scaler + through distributed save as ShardedObjects (common.pt is rank-0 only, + which misses LLM optimizer state in non-colocated mode). """ sharded_state = {} for name, info in self.module_infos.items(): @@ -177,7 +175,6 @@ 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 @@ -221,8 +218,6 @@ 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): @@ -237,23 +232,6 @@ 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. - - Plain non-tensor scalars at the per-module level otherwise travel through - dist_checkpointing's common-state path (rank 0 only), so for non-colocated - MIMO they are lost on ranks whose module is inactive on rank 0. - """ - 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) @@ -292,14 +270,6 @@ def _restore_grad_scaler(sub_sd): break -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 _get_replica_id(pg_collection: Optional[ProcessGroupCollection]) -> tuple: """Build replica_id tuple for ShardedObject deduplication. diff --git a/tests/unit_tests/models/test_mimo_checkpoint.py b/tests/unit_tests/models/test_mimo_checkpoint.py index 1ccaf33c0e2..1f2079db723 100644 --- a/tests/unit_tests/models/test_mimo_checkpoint.py +++ b/tests/unit_tests/models/test_mimo_checkpoint.py @@ -56,25 +56,8 @@ def _randomize_params(model, seed): p.random_() -def _create_model_and_optimizer( - encoder_grid, - llm_grid, - hidden_size, - num_layers, - vocab_size, - seed, - use_distributed_optimizer=False, -): - """Create MIMO model with DDP + optimizer. - - When `use_distributed_optimizer=False` (Float16Optimizer), take a few fake - backward + step iterations so the Adam step counter is non-trivial. The - DistributedOptimizer path requires DDP-driven grad reduce-scatter into the - gradient buffer to take a step, which isn't worth wiring up here — the - parametrization that uses DistributedOptimizer is a structural regression - guard (it checks that the save/load code paths in MimoOptimizer don't crash - when run against a `DistributedOptimizer` + `fully_reshardable` setup), - not a step-continuity test. +def _create_model_and_optimizer(encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed): + """Create MIMO model with DDP + optimizer, do fake steps to populate optimizer state. Caller must call create_all_embedding_groups() before this function. """ @@ -91,48 +74,25 @@ def _create_model_and_optimizer( ) _randomize_params(mimo_model, seed) + # Use Float16Optimizer (not DistributedOptimizer) to exercise the MIMO-specific + # param_groups/grad_scaler extraction in sharded_state_dict. DistributedOptimizer + # handles its own checkpointing internally and our code is transparent to it. opt_config = OptimizerConfig( optimizer='adam', lr=1e-4, weight_decay=0.01, clip_grad=1.0, bf16=True, - use_distributed_optimizer=use_distributed_optimizer, + use_distributed_optimizer=False, ) optimizer = get_mimo_optimizer(mimo_model, opt_config) - if not use_distributed_optimizer: - # Float16Optimizer path: take several fake backward + step iterations - # so we can later assert Adam step-counter continuity. - for _ in range(3): - for param in mimo_model.parameters(): - param.grad = torch.randn_like(param) - optimizer.step() - else: - # DistributedOptimizer path: a real step requires DDP-driven grad - # reduce-scatter into the gradient buffer. We don't need actual step - # values for this regression test (its purpose is to exercise the - # save/load *structure* — see `test_distributed_optimizer_fully_reshardable`), - # so just initialize the inner torch optimizer's state buffers - # (`exp_avg`, `exp_avg_sq`) with zeros via `init_state_fn`. Without - # this, `get_parameter_state_dp_zero` later fails with - # `KeyError: 'exp_avg'` when it tries to read state for each param. - # - # Each per-module optimizer is wrapped in a `ChainedOptimizer` (which - # holds N>=1 DistributedOptimizers); init_state_fn lives on the inner - # ones, not on the chain wrapper. - from megatron.core.optimizer.optimizer import ChainedOptimizer - - for info in optimizer.module_infos.values(): - if not (info.is_active and info.optimizer is not None): - continue - inner_opts = ( - info.optimizer.chained_optimizers - if isinstance(info.optimizer, ChainedOptimizer) - else [info.optimizer] - ) - for inner in inner_opts: - inner.init_state_fn(inner.optimizer, inner.config) + # Take several fake backward + step iterations so the Adam step counter is + # non-trivial; this lets us verify step-counter continuity across save/load. + for _ in range(3): + for param in mimo_model.parameters(): + param.grad = torch.randn_like(param) + optimizer.step() return mimo_model, optimizer @@ -169,16 +129,8 @@ def run_checkpoint_test( hidden_size=256, num_layers=2, vocab_size=1000, - use_distributed_optimizer=False, - distrib_optim_sharding_type='fully_reshardable', ): - """Save model + optimizer checkpoint, load into fresh instances, verify match. - - `use_distributed_optimizer` selects between Float16Optimizer (default; step - continuity is asserted) and DistributedOptimizer (structural regression - guard for the load-side fixes — `_restore_param_groups` setdefault and - `param_state_sharding_type` recovery from stashed metadata). - """ + """Save model + optimizer checkpoint, load into fresh instances, verify match.""" # Clear NVTE env vars that the conftest set_env fixture sets to '0'. # GPTModel (LanguageModule) asserts these are unset or match the attention backend. os.environ.pop('NVTE_FLASH_ATTN', None) @@ -191,21 +143,9 @@ def run_checkpoint_test( llm_grid = create_hypercomm_grid(offset=llm_offset, tp=llm_tp, cp=1, pp=llm_pp, dp=llm_dp) create_all_embedding_groups([encoder_grid, llm_grid]) - optim_metadata = ( - {'distrib_optim_sharding_type': distrib_optim_sharding_type} - if use_distributed_optimizer - else {} - ) - # --- Create model A + optimizer, snapshot state --- model_a, optimizer_a = _create_model_and_optimizer( - encoder_grid, - llm_grid, - hidden_size, - num_layers, - vocab_size, - seed=1, - use_distributed_optimizer=use_distributed_optimizer, + encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed=1 ) params_a = {name: p.clone() for name, p in model_a.named_parameters()} @@ -225,22 +165,14 @@ def run_checkpoint_test( # validate_access_integrity=True is the regression guard for the _get_replica_id fix: # without including TP rank in replica_id, every TP rank at (pp=0, dp=0) would emit # the same `_mimo_*` ShardedObject as a main replica, producing duplicate-key errors. - optim_sd_a = optimizer_a.sharded_state_dict( - model_a.sharded_state_dict(), is_loading=False, metadata=optim_metadata - ) + optim_sd_a = optimizer_a.sharded_state_dict(model_a.sharded_state_dict(), is_loading=False) save(optim_sd_a, optim_ckpt, validate_access_integrity=True) dist.barrier() # --- Create model B + optimizer with different weights (reuse same grids) --- model_b, optimizer_b = _create_model_and_optimizer( - encoder_grid, - llm_grid, - hidden_size, - num_layers, - vocab_size, - seed=2, - use_distributed_optimizer=use_distributed_optimizer, + encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed=2 ) # Load model @@ -255,9 +187,7 @@ def run_checkpoint_test( model_b.load_state_dict(loaded_model_sd) # Load optimizer - optim_sd_b = optimizer_b.sharded_state_dict( - model_b.sharded_state_dict(), is_loading=True, metadata=optim_metadata - ) + optim_sd_b = optimizer_b.sharded_state_dict(model_b.sharded_state_dict(), is_loading=True) loaded_optim_sd = load(optim_sd_b, optim_ckpt, validate_access_integrity=False) optimizer_b.load_state_dict(loaded_optim_sd) @@ -298,16 +228,12 @@ def run_checkpoint_test( # Verify Adam step counter survives save/load. Without restoring step, # bias correction would reset and cause an effective LR drop on resume. - # Only meaningful when real optimizer steps were taken (the - # DistributedOptimizer parametrization skips steps — see - # _create_model_and_optimizer). - if not use_distributed_optimizer: - step_a = _find_optimizer_step(sd_a) - step_b = _find_optimizer_step(sd_b) - assert step_a is not None, f"Optimizer {name}: could not locate step in model A" - assert step_a == step_b, ( - f"Optimizer {name}: step mismatch after load " f"(A={step_a}, B={step_b})" - ) + step_a = _find_optimizer_step(sd_a) + step_b = _find_optimizer_step(sd_b) + assert step_a is not None, f"Optimizer {name}: could not locate step in model A" + assert ( + step_a == step_b + ), f"Optimizer {name}: step mismatch after load (A={step_a}, B={step_b})" finally: _cleanup_tmpdir(ckpt_dir) @@ -379,39 +305,3 @@ def test_encoder_tp2_pp2_llm_tp2_pp2(self): hidden_size=256, num_layers=2, ) - - def test_distributed_optimizer_fully_reshardable(self): - """Regression guard for the DistributedOptimizer + fully_reshardable load path. - - Exercises the exact code path the hetero LLaVA trainer hits and which - the Float16Optimizer-based parametrizations above don't cover: - - - `DistributedOptimizer.state_dict()` puts only `param_groups` under - the 'optimizer' key (state is sharded separately as `param_state`). - After `_extract_param_groups` deletes `param_groups`, what's left is - an empty dict that is dropped through dist_checkpointing's common- - state round-trip on ranks where the active module isn't on rank 0. - `_restore_param_groups` must use `setdefault` to survive this. - - `param_state_sharding_type` is a plain scalar that also doesn't - survive the common-state round-trip on those ranks. `load_state_dict` - must recover it from the metadata stash set by `sharded_state_dict`. - - Uses an 8-GPU non-colocated layout where rank 0 is NOT in the language - module, so the dropped-on-rank-0 codepath is forced. - """ - if self.world_size != 8: - pytest.skip(f"Requires 8 GPUs, got {self.world_size}") - run_checkpoint_test( - encoder_tp=2, - encoder_pp=1, - encoder_dp=1, - encoder_offset=0, - llm_tp=2, - llm_pp=3, - llm_dp=1, - llm_offset=2, - hidden_size=256, - num_layers=3, - use_distributed_optimizer=True, - distrib_optim_sharding_type='fully_reshardable', - ) From 927d0830cd4a83ce69a2905306c670bf22dd9f5f Mon Sep 17 00:00:00 2001 From: Kamran Jafari Date: Tue, 19 May 2026 05:00:43 -0700 Subject: [PATCH 4/4] revert changes to the tests Signed-off-by: Kamran Jafari --- .../unit_tests/models/test_mimo_checkpoint.py | 41 +++---------------- 1 file changed, 5 insertions(+), 36 deletions(-) diff --git a/tests/unit_tests/models/test_mimo_checkpoint.py b/tests/unit_tests/models/test_mimo_checkpoint.py index 1f2079db723..5cb28313cfd 100644 --- a/tests/unit_tests/models/test_mimo_checkpoint.py +++ b/tests/unit_tests/models/test_mimo_checkpoint.py @@ -57,7 +57,7 @@ def _randomize_params(model, seed): def _create_model_and_optimizer(encoder_grid, llm_grid, hidden_size, num_layers, vocab_size, seed): - """Create MIMO model with DDP + optimizer, do fake steps to populate optimizer state. + """Create MIMO model with DDP + optimizer, do a fake step to populate optimizer state. Caller must call create_all_embedding_groups() before this function. """ @@ -87,36 +87,14 @@ def _create_model_and_optimizer(encoder_grid, llm_grid, hidden_size, num_layers, ) optimizer = get_mimo_optimizer(mimo_model, opt_config) - # Take several fake backward + step iterations so the Adam step counter is - # non-trivial; this lets us verify step-counter continuity across save/load. - for _ in range(3): - for param in mimo_model.parameters(): - param.grad = torch.randn_like(param) - optimizer.step() + # Fake backward + step to populate optimizer state (Adam m/v) + for param in mimo_model.parameters(): + param.grad = torch.randn_like(param) + optimizer.step() return mimo_model, optimizer -def _find_optimizer_step(state_dict): - """Return the optimizer step counter from a per-module state_dict, regardless of - whether it lives in `state['common_step']`, `param_groups[*]['step']`, or per-param state. - """ - opt = state_dict.get('optimizer', {}) - state = opt.get('state', {}) - if 'common_step' in state: - val = state['common_step'] - return int(val.item()) if torch.is_tensor(val) else int(val) - for g in opt.get('param_groups', []): - if 'step' in g: - val = g['step'] - return int(val.item()) if torch.is_tensor(val) else int(val) - for s in state.values(): - if isinstance(s, dict) and 'step' in s: - val = s['step'] - return int(val.item()) if torch.is_tensor(val) else int(val) - return None - - def run_checkpoint_test( encoder_tp, encoder_pp, @@ -226,15 +204,6 @@ def run_checkpoint_test( state_a[param_id][key], state_b[param_id][key] ), f"Optimizer {name} param {param_id} {key} mismatch" - # Verify Adam step counter survives save/load. Without restoring step, - # bias correction would reset and cause an effective LR drop on resume. - step_a = _find_optimizer_step(sd_a) - step_b = _find_optimizer_step(sd_b) - assert step_a is not None, f"Optimizer {name}: could not locate step in model A" - assert ( - step_a == step_b - ), f"Optimizer {name}: step mismatch after load (A={step_a}, B={step_b})" - finally: _cleanup_tmpdir(ckpt_dir)