Thread MIMO support through the stock training loop (schedule + optimizer)

megatron/core/optimizer/__init__.py

@@ -1002,6 +1002,17 @@ def get_megatron_optimizer(
         Instance of MegatronOptimizer.
     """
 
+    # A MimoModel routes to the heterogeneous per-module optimizer builder.
+    from megatron.core.models.mimo.model.base import MimoModel
+
+    if isinstance(model_chunks[0], MimoModel):
+        from megatron.core.models.mimo.optimizer import get_mimo_optimizer
+
+        assert (
+            len(model_chunks) == 1
+        ), "MimoModel does not support virtual pipeline parallelism (multiple model chunks)"
+        return get_mimo_optimizer(model_chunks[0], config)
+
     # None → apply standard defaults. To extend defaults with custom overrides,
     # start from get_standard_config_overrides(config) and merge yours in.
     if config_overrides is None:

megatron/training/training.py

@@ -156,7 +156,11 @@ def set_startup_timestamps(program_start=None, main_entry=None):
     is_vp_first_stage,
     is_vp_last_stage,
 )
-from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.pipeline_parallel.p2p_communication import P2PCommunicator
+from megatron.core.process_groups_config import (
+    MultiModuleProcessGroupCollection,
+    ProcessGroupCollection,
+)
 from megatron.core.transformer.cuda_graphs import TECudaGraphHelper
 from megatron.core.transformer.module import Float16Module
 from megatron.core.transformer.moe.paged_stash import PagedStashRunner
@@ -1957,6 +1961,7 @@ def setup_model_and_optimizer(
     model_provider_func,
     model_type,
     checkpointing_context=None,
+    pg_collection=None,
 ):
     """Setup model and optimizer."""
     args = get_args()
@@ -1969,7 +1974,9 @@ def setup_model_and_optimizer(
     has_rl_optimizer = args.perform_rl_step and not args.no_load_optim
     skip_optimizer = not (has_normal_optimizer or has_rl_optimizer)
     wrap_with_ddp = not skip_optimizer
-    model = get_model(model_provider_func, model_type, wrap_with_ddp=wrap_with_ddp)
+    model = get_model(
+        model_provider_func, model_type, wrap_with_ddp=wrap_with_ddp, pg_collection=pg_collection
+    )
     unwrapped_model = unwrap_model(model)
 
     if args.logits_save_dir is not None:
@@ -2195,11 +2202,15 @@ def dummy_train_step(data_iterator):
             )
 
 
-def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None, pg_collection: Optional[ProcessGroupCollection] = None):
+def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=None, pg_collection: Optional[ProcessGroupCollection] = None, p2p_communicator: Optional[P2PCommunicator] = None, schedule_pg_collection: Optional[MultiModuleProcessGroupCollection] = None):
     """Single training step.
 
     pg_collection: optional per-module :class:`ProcessGroupCollection`; None uses the mpu globals,
         otherwise it must define mp, pp, and dp_cp.
+    p2p_communicator: optional communicator forwarded to the schedule for cross-grid P2P; None
+        preserves the default behavior.
+    schedule_pg_collection: optional per-module groups forwarded to the schedule for the
+        cross-grid case; None preserves the default behavior.
     """
     args = get_args()
     timers = get_timers()
@@ -2275,6 +2286,8 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch
             forward_only=False,
             adjust_tensor_shapes_fn=adjust_tensor_shapes_fn,
             force_all_reduce=save_wgrads_in_this_iteration,
+            p2p_communicator=p2p_communicator,
+            pg_collection=schedule_pg_collection,
         )
         if save_activations_in_this_iteration:
             save_activations(iteration + 1)
@@ -3116,8 +3129,16 @@ def train(
     checkpointing_context,
     non_loss_data_func,
     inference_model=None,
+    p2p_communicator: Optional[P2PCommunicator] = None,
+    schedule_pg_collection: Optional[MultiModuleProcessGroupCollection] = None,
 ):
-    """Training function: run train_step desired number of times, run validation, checkpoint."""
+    """Training function: run train_step desired number of times, run validation, checkpoint.
+
+    p2p_communicator: optional communicator forwarded to the schedule for cross-grid P2P; None
+        preserves the default behavior.
+    schedule_pg_collection: optional per-module groups forwarded to the schedule for the
+        cross-grid case; None preserves the default behavior.
+    """
     args = get_args()
     timers = get_timers()
     fault_injector_kwargs = {}
@@ -3572,7 +3593,9 @@ def trace_handler(p):
                 num_zeros_in_grad,
                 max_attention_logit,
             ) = train_step(
-                forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration
+                forward_step_func, train_data_iterator, model, optimizer, opt_param_scheduler, config, forward_backward_func, iteration=iteration,
+                pg_collection=model_pg_collection,
+                p2p_communicator=p2p_communicator, schedule_pg_collection=schedule_pg_collection
             )
             ft_integration.on_training_step_end()
             if _maybe_raise_workload_exception is not None and iteration != start_iteration:

tests/unit_tests/training/test_setup_model_and_optimizer_mimo.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+
+"""get_megatron_optimizer dispatches a MimoModel to get_mimo_optimizer."""
+
+from types import SimpleNamespace
+from unittest import mock
+
+from megatron.core.models.mimo.model.base import MimoModel
+from megatron.core.optimizer import get_megatron_optimizer
+
+
+def test_get_megatron_optimizer_dispatches_mimo_model():
+    # __new__ gives an isinstance-true MimoModel without running its heavy __init__.
+    fake_mimo = MimoModel.__new__(MimoModel)
+    sentinel = object()
+    with mock.patch(
+        "megatron.core.models.mimo.optimizer.get_mimo_optimizer", return_value=sentinel
+    ) as get_mimo:
+        result = get_megatron_optimizer(SimpleNamespace(optimizer="adam"), [fake_mimo])
+    get_mimo.assert_called_once()
+    assert result is sentinel

tests/unit_tests/training/test_train_step_schedule_plumbing.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""train_step forwards p2p_communicator and schedule pg_collection to forward_backward_func."""
+
+from types import SimpleNamespace
+from unittest import mock
+
+from megatron.training import training as training_mod
+
+
+class _Rerun:
+    """Run the forward/backward body once, then ask train_step to exit before optimizer.step."""
+
+    _ran = False
+
+    def should_run_forward_backward(self, data_iterator):
+        run, self._ran = not self._ran, True
+        return run
+
+    def should_checkpoint_and_exit(self):
+        return False, True, 0  # (checkpoint, exit, code)
+
+
+def _run(**kwargs):
+    args = SimpleNamespace(
+        save_params_interval=None,
+        save_activations_interval=None,
+        save_tokens_per_expert_interval=None,
+        save_wgrads_interval=None,
+        save_dgrads_interval=None,
+        reuse_grad_buf_for_mxfp8_param_ag=False,
+        overlap_param_gather=False,
+        seq_length=8,
+        micro_batch_size=1,
+        decoder_seq_length=None,
+        empty_unused_memory_level=0,
+    )
+    captured = {}
+    model = [SimpleNamespace(force_all_reduce=False, zero_grad_buffer=lambda: None)]
+    with (
+        mock.patch.object(training_mod, "get_args", return_value=args),
+        mock.patch.object(training_mod, "get_timers", return_value=mock.MagicMock()),
+        mock.patch.object(training_mod, "get_rerun_state_machine", return_value=_Rerun()),
+        mock.patch.object(training_mod, "get_num_microbatches", return_value=1),
+        mock.patch.object(training_mod, "has_nvidia_modelopt", False),
+    ):
+        training_mod.train_step(
+            forward_step_func=lambda *a, **k: None,
+            data_iterator=iter([]),
+            model=model,
+            optimizer=SimpleNamespace(zero_grad=lambda: None),
+            opt_param_scheduler=None,
+            config=SimpleNamespace(),
+            forward_backward_func=lambda **kw: captured.update(kw) or [],
+            iteration=0,
+            **kwargs,
+        )
+    return captured
+
+
+def test_train_step_forwards_schedule_plumbing():
+    p2p, pg = object(), object()
+    captured = _run(p2p_communicator=p2p, schedule_pg_collection=pg)
+    assert captured["p2p_communicator"] is p2p and captured["pg_collection"] is pg
+
+
+def test_train_step_defaults_to_none():
+    captured = _run()
+    assert captured["p2p_communicator"] is None and captured["pg_collection"] is None