test(training): align optimizer fixtures with Muon setup

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

tests/functional_tests/training/muon_distributed_checkpoint.py

@@ -1,3 +1,5 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
 """Two-GPU numerical and save/resume regression for distributed optimizers.
 
 Run with torch.distributed.run. Each case uses a fresh process; --compare checks

tests/unit_tests/data/megatron_mimo/test_loaders.py

@@ -299,9 +299,15 @@ def test_mimo_runtime_update_normalizes_zero_worker_provider_before_loader_const
     runtime_cfg.validation.eval_micro_batch_size = None
     runtime_cfg.profiling = None
     runtime_cfg.ddp.use_distributed_optimizer = False
+    runtime_cfg.ddp.num_distributed_optimizer_instances = 1
     runtime_cfg.optimizer.use_distributed_optimizer = False
+    runtime_cfg.optimizer.use_layer_wise_distributed_optimizer = False
+    runtime_cfg.optimizer.optimizer = "adam"
+    runtime_cfg.optimizer.overlap_param_gather_with_optimizer_step = False
     runtime_cfg.ddp.overlap_param_gather = False
     runtime_cfg.optimizer.overlap_param_gather = False
+    runtime_cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = False
+    runtime_cfg.optimizer.reuse_grad_buf_for_mxfp8_param_ag = False
 
     megatron_mimo_runtime_config_update(runtime_cfg)
 

tests/unit_tests/training/megatron_mimo/test_megatron_mimo_checkpointing.py

@@ -54,6 +54,16 @@ def test_megatron_mimo_runtime_config_update_resolves_eval_batch_size_defaults()
     cfg.validation.eval_global_batch_size = None
     cfg.validation.eval_micro_batch_size = None
     cfg.profiling = None
+    cfg.ddp.use_distributed_optimizer = False
+    cfg.ddp.num_distributed_optimizer_instances = 1
+    cfg.ddp.overlap_param_gather = False
+    cfg.ddp.reuse_grad_buf_for_mxfp8_param_ag = False
+    cfg.optimizer.optimizer = "adam"
+    cfg.optimizer.use_distributed_optimizer = False
+    cfg.optimizer.use_layer_wise_distributed_optimizer = False
+    cfg.optimizer.overlap_param_gather = False
+    cfg.optimizer.overlap_param_gather_with_optimizer_step = False
+    cfg.optimizer.reuse_grad_buf_for_mxfp8_param_ag = False
 
     megatron_mimo_runtime_config_update(cfg)
 

tests/unit_tests/training/test_decentralized_pg.py

@@ -687,12 +687,12 @@ def test_setup_optimizer_passes_none_pg_collection_when_not_provided(
         assert call_kwargs["pg_collection"] is None
 
     @patch("megatron.bridge.training.optim.get_model_config")
-    @patch("megatron.bridge.training.optim.get_megatron_muon_optimizer")
+    @patch("megatron.bridge.training.optim.get_megatron_optimizer")
     @patch("megatron.bridge.training.optim.OptimizerParamScheduler")
     def test_setup_optimizer_passes_pg_collection_to_muon_optimizer(
-        self, mock_scheduler, mock_get_muon_optimizer, mock_get_model_config
+        self, mock_scheduler, mock_get_optimizer, mock_get_model_config
     ):
-        """Test that setup_optimizer passes pg_collection to muon optimizer."""
+        """Test that setup_optimizer passes pg_collection for Muon construction."""
         from megatron.core.optimizer import OptimizerConfig
 
         from megatron.bridge.training.config import SchedulerConfig
@@ -703,7 +703,7 @@ def test_setup_optimizer_passes_pg_collection_to_muon_optimizer(
         mock_model = MagicMock()
         mock_pg_collection = MagicMock()
         mock_optimizer = MagicMock()
-        mock_get_muon_optimizer.return_value = mock_optimizer
+        mock_get_optimizer.return_value = mock_optimizer
 
         optimizer_config = OptimizerConfig(optimizer="muon", lr=1e-3)
         scheduler_config = SchedulerConfig()
@@ -717,9 +717,10 @@ def test_setup_optimizer_passes_pg_collection_to_muon_optimizer(
             pg_collection=mock_pg_collection,
         )
 
-        # Verify pg_collection was passed to muon optimizer
-        mock_get_muon_optimizer.assert_called_once()
-        call_kwargs = mock_get_muon_optimizer.call_args[1]
+        # Verify pg_collection was passed through the unified optimizer factory.
+        mock_get_optimizer.assert_called_once()
+        call_kwargs = mock_get_optimizer.call_args[1]
+        assert call_kwargs["config"] is optimizer_config
         assert call_kwargs["pg_collection"] == mock_pg_collection
 
     @patch("megatron.bridge.training.optim.get_model_config")

tests/unit_tests/training/test_gtp.py

@@ -324,7 +324,11 @@ def test_distributed_model_build_obeys_gtp_lifecycle(mock_configure, mock_classi
     cfg = SimpleNamespace(
         model=model_config,
         ddp=object(),
-        optimizer=SimpleNamespace(overlap_param_gather_with_optimizer_step=False),
+        optimizer=SimpleNamespace(
+            optimizer="adam",
+            overlap_param_gather_with_optimizer_step=False,
+            use_layer_wise_distributed_optimizer=False,
+        ),
         dist=DistributedInitConfig(
             use_megatron_fsdp=False,
             use_torch_fsdp2=False,

tests/unit_tests/training/test_optimizer_distribution.py

@@ -1,3 +1,5 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
 """Regression tests for optimizer selection, DDP layout and checkpoint state."""
 
 from types import SimpleNamespace

tests/unit_tests/training/test_pg_collection_wiring.py

@@ -116,7 +116,9 @@ def test_train_stops_nsys_profiler_when_skipped_iteration_reaches_profile_end(mo
         straggler=None,
         ddp=SimpleNamespace(use_megatron_fsdp=False, overlap_param_gather=False),
         optimizer=SimpleNamespace(
+            optimizer="adam",
             use_distributed_optimizer=False,
+            use_layer_wise_distributed_optimizer=False,
             optimizer_cuda_graph=False,
         ),
         model=SimpleNamespace(
@@ -224,7 +226,9 @@ def test_train_stops_nsys_profiler_when_rerun_requests_exit(monkeypatch):
         straggler=None,
         ddp=SimpleNamespace(use_megatron_fsdp=False, overlap_param_gather=False),
         optimizer=SimpleNamespace(
+            optimizer="adam",
             use_distributed_optimizer=False,
+            use_layer_wise_distributed_optimizer=False,
             optimizer_cuda_graph=False,
         ),
         model=SimpleNamespace(

tests/unit_tests/training/test_runtime_callbacks.py

@@ -89,7 +89,11 @@ def _run_setup(
             set_level_for_all_loggers=False,
         ),
         model=provider,
-        optimizer=SimpleNamespace(overlap_param_gather_with_optimizer_step=False),
+        optimizer=SimpleNamespace(
+            optimizer="adam",
+            overlap_param_gather_with_optimizer_step=False,
+            use_layer_wise_distributed_optimizer=False,
+        ),
         optimizer_config_override_provider=None,
         peft=None,
         profiling=SimpleNamespace(),