fix(recipes): preserve Kimi Muon layout settings

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

Author: cuichenx

docs/training/optimizer-scheduler.md

@@ -53,11 +53,16 @@ By default, `use_layer_wise_param_layout=True` builds shard-aligned DDP buffers
 for the LayerWise optimizer. Set it to `False` to retain compact DDP buffers and
 use LayerWise's variable-size parameter synchronization path. This reduces
 layout padding while preserving the configured parameter-gather behavior.
+Existing Kimi recipes keep the compact layout and parameter-gather overlap
+disabled until their MoE and pipeline-parallel configurations are validated
+with the shard-aligned layout; other recipes can opt in independently.
 
 Resuming Muon training restores Muon optimizer state. When changing optimization
 algorithms between runs, load model weights with fresh optimizer state using
 `checkpoint.load_optim=False` or the appropriate finetuning configuration;
 Bridge does not convert Adam moments into Muon optimizer state.
+Checkpoints created before distributed-optimizer metadata was recorded retain
+the historical model-space sharding fallback when that metadata key is absent.
 
 ## SchedulerConfig
 

src/megatron/bridge/recipes/kimi/h100/kimi_k2.py

@@ -63,16 +63,23 @@ def _apply_kimi_k2_optimizer(cfg: ConfigContainer, optimizer_type: str) -> None:
         grad_reduce_in_fp32 = False
         cfg.ddp.use_distributed_optimizer = True
         cfg.ddp.overlap_param_gather = True
+        if cfg.comm_overlap is not None:
+            cfg.comm_overlap.overlap_param_gather = True
     elif optimizer_type == "muon":
         opt_cfg, scheduler_cfg = distributed_muon_with_cosine_annealing(
             lr_warmup_iters=2000,
             lr_decay_iters=cfg.train.train_iters,
             max_lr=3e-4,
             min_lr=3e-5,
         )
+        # Preserve this recipe's established compact LayerWise layout until the
+        # shard-aligned layout is validated with its MoE/pipeline topology.
+        opt_cfg.use_layer_wise_param_layout = False
         grad_reduce_in_fp32 = True
         cfg.ddp.use_distributed_optimizer = False
         cfg.ddp.overlap_param_gather = False
+        if cfg.comm_overlap is not None:
+            cfg.comm_overlap.overlap_param_gather = False
     else:
         raise ValueError(f"Invalid optimizer type: {optimizer_type}")
 

src/megatron/bridge/recipes/kimi_vl/h100/kimi_k25_vl.py

@@ -121,6 +121,9 @@ def kimi_k25_vl_sft_512gpu_h100_bf16_config() -> ConfigContainer:
         max_lr=3e-4,
         min_lr=3e-5,
     )
+    # Preserve this recipe's established compact LayerWise layout until the
+    # shard-aligned layout is validated with its MoE/pipeline topology.
+    opt_cfg.use_layer_wise_param_layout = False
     cfg.optimizer = opt_cfg
     cfg.scheduler = scheduler_cfg
 
@@ -162,7 +165,7 @@ def kimi_k25_vl_sft_512gpu_h100_bf16_config() -> ConfigContainer:
     cfg.optimizer.exp_avg_sq_dtype = torch.float32
 
     # Communication overlap
-    cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False)
+    cfg.comm_overlap = CommOverlapConfig(tp_comm_overlap=False, overlap_param_gather=False)
     cfg.comm_overlap.delay_wgrad_compute = False
     cfg.comm_overlap.overlap_moe_expert_parallel_comm = False
     cfg.model.moe_shared_expert_overlap = True
@@ -171,7 +174,8 @@ def kimi_k25_vl_sft_512gpu_h100_bf16_config() -> ConfigContainer:
     cfg.checkpoint.save_interval = 2000
     cfg.checkpoint.async_save = False
 
-    # DDP config — Muon requires distributed_optimizer=False and overlap_param_gather=False
+    # LayerWise Muon does not use DDP's ordinary DistributedOptimizer mode.
+    # Keep parameter-gather overlap disabled for this recipe's uneven MIMO inputs.
     cfg.ddp.use_distributed_optimizer = False
     cfg.ddp.overlap_param_gather = False
     cfg.ddp.grad_reduce_in_fp32 = True

tests/unit_tests/recipes/kimi/test_kimi_k2.py

@@ -22,7 +22,7 @@
     _get_kimi_k2_pipeline_layout,
     kimi_k2_pretrain_config,
 )
-from megatron.bridge.training.config import ConfigContainer
+from megatron.bridge.training.config import ConfigContainer, _validate_and_sync_distributed_optimizer_settings
 from megatron.bridge.training.mixed_precision import MixedPrecisionConfig
 from tests.unit_tests.recipes.recipe_test_utils import patch_recipe_module_global
 
@@ -148,18 +148,33 @@ def test_pretrain_config_pipeline_split_settings(self):
         assert cfg.model.num_layers_in_last_pipeline_stage is None
 
     def test_pretrain_config_ddp_settings_for_muon(self):
-        """Test DDP settings configured for Muon optimizer."""
+        """The existing Kimi recipe retains compact, non-overlapped LayerWise Muon."""
         cfg = kimi_k2_pretrain_config()
 
-        # Muon requires these specific DDP settings
         assert cfg.ddp.overlap_grad_reduce is True
-        assert cfg.ddp.overlap_param_gather is False  # Muon needs this to be False
+        assert cfg.ddp.overlap_param_gather is False
+        assert cfg.comm_overlap.overlap_param_gather is False
         assert cfg.ddp.check_for_nan_in_grad is True
-        assert cfg.ddp.use_distributed_optimizer is False  # Muon needs this to be False
+        assert cfg.ddp.use_distributed_optimizer is False
         assert cfg.ddp.use_megatron_fsdp is False
         assert cfg.ddp.grad_reduce_in_fp32 is True
         assert cfg.ddp.average_in_collective is True
         assert cfg.ddp.data_parallel_sharding_strategy == "no_shard"
+        assert cfg.optimizer.optimizer == "dist_muon"
+        assert cfg.optimizer.use_layer_wise_param_layout is False
+
+        _validate_and_sync_distributed_optimizer_settings(cfg)
+
+        assert cfg.optimizer.optimizer == "muon"
+        assert cfg.optimizer.use_layer_wise_distributed_optimizer is True
+        assert cfg.optimizer.use_distributed_optimizer is False
+        assert cfg.ddp.use_distributed_optimizer is False
+        assert cfg.optimizer.use_layer_wise_param_layout is False
+
+        cfg.comm_overlap.finalize()
+        cfg.comm_overlap.data_parallel_size = 16
+        overlap_cfg = cfg.comm_overlap._get_optimizer_overlap_cfgs(cfg.model)
+        assert overlap_cfg.overlap_param_gather is False
 
     def test_pretrain_config_dataset_configuration(self):
         """Test dataset configuration."""

tests/unit_tests/recipes/kimi_vl/test_kimi_k25_vl.py

@@ -203,17 +203,19 @@ def test_sft_config_pipeline_layout_tracks_supported_runner_override(self):
         assert cfg.model.pipeline_model_parallel_layout == _get_kimi_k25_vl_pipeline_layout(4, 1)
 
     def test_sft_config_ddp_settings_for_muon(self):
-        """DDP settings respect Muon's constraints (no dist optimizer, no param overlap)."""
+        """The existing Kimi-VL recipe retains compact, non-overlapped LayerWise Muon."""
         cfg = kimi_k25_vl_sft_config()
 
         assert cfg.ddp.use_distributed_optimizer is False
         assert cfg.ddp.overlap_param_gather is False
+        assert cfg.comm_overlap.overlap_param_gather is False
         assert cfg.ddp.overlap_grad_reduce is True
         assert cfg.ddp.grad_reduce_in_fp32 is True
         assert cfg.ddp.check_for_nan_in_grad is True
         assert cfg.ddp.use_megatron_fsdp is False
         assert cfg.ddp.average_in_collective is True
         assert cfg.ddp.data_parallel_sharding_strategy == "no_shard"
+        assert cfg.optimizer.use_layer_wise_param_layout is False
 
     def test_sft_config_dataset_configuration(self):
         """Dataset uses sequence length 4096 and the model's HF processor path."""