NMFW-464: Distributed checkpoint save/load for the hetero MIMO training loop

[yashaswikarnati]

Summary

Adds distributed-checkpoint save/load to the standalone hetero MIMO training loop (examples/mimo/train_hetero.py) so MimoModel + MimoOptimizer (Chained DistributedOptimizer in the MoE recipe) + LR/WD scheduler + Python/Torch/CUDA RNG state round-trip through megatron.core.dist_checkpointing without depending on parallel_state. Layout is intentionally compatible with megatron/training/checkpointing.py output.

Validated end-to-end on cw-dfw 8-GPU 20L mock (stage2): save at iter 3 → reload → resume at iter 4 with cosine LR continuation (1.59e-4 → 1.32e-4 → 1.01e-4) and matching loss trajectory.

What's in this PR

New files

• examples/mimo/training/hetero/checkpointing.py — save_checkpoint / load_checkpoint wrapping dist_checkpointing.save/load with a tracker file, per-branch RNG ShardedObjects, content_metadata persistence, and atomic tracker writes.

CLI surface (in examples/mimo/training/hetero/args.py)

• --save, --load, --save-interval, --no-save-optim, --no-load-optim, --no-load-scheduler, --no-save-rng, --no-load-rng, --finetune, --dist-ckpt-optim-fully-reshardable (off by default; help text warns about save-time gather OOM on <80GB GPUs).

Loop wiring (examples/mimo/training/hetero/loop.py)

• Load before train iters, periodic save on --save-interval, always save final iter when --save is set.

Three pre-existing bug fixes uncovered by hetero usage

• megatron/core/ssm/mamba_mixer.py: MambaMixer.sharded_state_dict now passes tp_group=self.tp_group and dp_cp_group=metadata['dp_cp_group'] to make_sharded_tensors_for_checkpoint (was falling back to parallel_state; gated_delta_net already does this correctly).
• megatron/core/models/mimo/optimizer.py:MimoOptimizer.sharded_state_dict: now applies add_prefix_for_sharding(module_sd, f'mimo.{name}.') so the two branches' optimizers don't collide on identical internal ShardedObject keys like chained_0.optimizer.distributed.dp_group_idx_0.optimizer/shard_0_1.
• megatron/core/models/mimo/optimizer.py:_get_replica_id: now returns (pp_rank, tp_rank, dp_rank) instead of (0, pp_rank, dp_rank) — without tp_rank, two TP ranks within the same DP rank both claimed primary writer and dist_checkpointing.save errored with 'Duplicate ShardedObject keys'.

New _extract_param_state_sharding_type helpers in MimoOptimizer route DistributedOptimizer's top-level param_state_sharding_type config string through a per-module ShardedObject (with a unique key per branch), so the non-rank-0 module owner doesn't lose it when only rank 0's common.pt is authoritative. Without this, DistributedOptimizer.load_state_dict raised AssertionError: 'param_state_sharding_type' in state_dict.

TP group propagation (examples/mimo/training/hetero/runtime.py): a new _propagate_tp_groups_for_checkpoint walker stamps self.tp_group on every descendant that doesn't already have one (e.g. ExtendedRMSNorm in mamba_mixer, RADIO encoder submodules) so the default MegatronModule.sharded_state_dict path doesn't fall through to parallel_state.get_tensor_model_parallel_group. The walker uses hasattr so a module that intentionally sets tp_group = None is left alone.

Megatron-parity additions in common state

• args (dict(vars(args))) so the saved layout is inspectable / future-checkable.
• checkpoint_version = 3.0.
• content_metadata = _clean_metadata_for_serialization(_build_optim_metadata(args)) passed to dist_checkpointing.save, persisting distrib_optim_sharding_type in metadata.json.
• Cross-rank max-reduce in _read_tracker (mirrors megatron.training.checkpointing.read_metadata).
• Atomic tracker write (.tmp + os.replace).

Test coverage

• tests/unit_tests/models/test_mimo_checkpoint.py parameterized on use_distributed_optimizer. New test_encoder_tp2_llm_tp2_pp3_distributed_optimizer exercises the MimoOptimizer prefix walk + param_state_sharding_type extract that the existing Float16Optimizer variant doesn't reach.

Test plan

• cw-dfw 8-GPU 20L mock smoke: save iter 3 → reload → train iter 4-5 (DistributedOptimizer + EP=4 + TP=2 + 2-branch Chained), exit 0 on both runs.
• Verified resume reads optimizer=yes, scheduler=yes, rng=yes, finetune=False and LR continues cosine schedule across the reload boundary.
• tools/autoformat.sh (black, isort) clean.
• tests/unit_tests/models/test_mimo_checkpoint.py::TestMimoCheckpoint (existing Float16Optimizer + new DistributedOptimizer variant) — run on the next CI sweep.
• Cross-TP/EP reload with --dist-ckpt-optim-fully-reshardable — opt-in, deferred to a follow-up due to save-time memory cost.

Notes

• This is a fork-side PR against ykarnati/nmfw-464-nemotron-vlm-with-hetero-parallel (the NMFW-464 integration branch), not against NVIDIA/Megatron-LM:main. Cluster session nmfw-464-ckpt-iter2 retains the latest validation artifacts if you want to inspect the saved common.pt.
• Independent reviewer findings (B1 docstring on --save-interval, B4 mkdir-on-all-ranks, R1 tp_group propagation comment, R2 OOM warning in --dist-ckpt-optim-fully-reshardable help) are all addressed in this commit.

🤖 Generated with Claude Code

[yashaswikarnati]

should tp rank be first or pp rank? how do we know ? can we check existing megatron code for optimizer or other classes? whats the expectation here ?

[yashaswikarnati]

just for my own understanding why do we need this ?

[yashaswikarnati]

for my own understanding, why do we need this ? this walks over all modules and sets tp group for every module in self.tp_group? which modules are missing ? trying to think if this indirect way is better or we directly fix the module that does not have tp group set, depends on what modules are missing. now

[yashaswikarnati]

can we reuse functions and functionality from megatron/training where ever possible ?