fix(training): preserve compact LayerWise workflows

Propagate the LayerWise parameter-layout choice through standard and MIMO model wrapping, and keep the ordinary distributed-optimizer DDP flag disabled for compact layouts. Detect nested distributed optimizer leaves when choosing checkpoint metadata and cover DP-resize resume behavior.

Co-authored-by: ykarnati <ykarnati@nvidia.com>

Co-authored-by: yaoyu-33 <yaoyu.094@gmail.com>
Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

docs/training/optimizer-scheduler.md

@@ -49,6 +49,11 @@ layer-wise Muon. The current Core implementation requires one distributed
 optimizer instance and does not support `overlap_param_gather_with_optimizer_step`
 or either FSDP implementation with Muon.
 
+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.
+
 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;
@@ -145,4 +150,4 @@ These scheduling options control how the weight decay coefficient changes over t
 |---------------|-----------------------------------------------------------------------------|
 | **Constant**  | Fixed weight decay throughout training.                                     |
 | **Linear**    | Linear progression from start to end weight decay.                          |
-| **Cosine**    | Cosine progression from start to end weight decay.                          |
+| **Cosine**    | Cosine progression from start to end weight decay.                          |

src/megatron/bridge/models/megatron_mimo/build_model.py

@@ -41,6 +41,7 @@ def build_megatron_mimo_model(
     wrap_with_ddp: bool = True,
     data_parallel_random_init: bool = True,
     use_layer_wise_distributed_optimizer: bool = False,
+    use_layer_wise_param_layout: bool = True,
 ) -> tuple["MimoModel", "MegatronMIMOInfra"]:
     """Build a distributed MegatronMIMO model and return ``(model, infra)``.
 
@@ -78,6 +79,7 @@ def build_megatron_mimo_model(
         wrap_with_ddp=wrap_with_ddp,
         data_parallel_random_init=data_parallel_random_init,
         use_layer_wise_distributed_optimizer=use_layer_wise_distributed_optimizer,
+        use_layer_wise_param_layout=use_layer_wise_param_layout,
     )
     mimo_model = model_list[0]
 

src/megatron/bridge/models/megatron_mimo/megatron_mimo_ddp.py

@@ -14,6 +14,8 @@
 from megatron.core.models.mimo.config.role import MIMO_LANGUAGE_MODULE_KEY
 from torch import nn
 
+from megatron.bridge.utils.common_utils import warn_rank_0
+
 
 if TYPE_CHECKING:
     from megatron.core.distributed import (
@@ -36,6 +38,7 @@ def wrap_megatron_mimo_model_distributed(
     pg_collections: Dict[str, Optional["ProcessGroupCollection"]],
     *,
     use_layer_wise_distributed_optimizer: bool = False,
+    use_layer_wise_param_layout: bool = True,
 ) -> "MimoModel":
     """Wrap MegatronMIMO model's submodules with DDP.
 
@@ -48,6 +51,8 @@ def wrap_megatron_mimo_model_distributed(
         grids: Module name to HyperCommGrid mapping.
         pg_collections: Module name to ProcessGroupCollection mapping.
         use_layer_wise_distributed_optimizer: Build shard-aligned layouts for layer-wise optimizers.
+        use_layer_wise_param_layout: Use shard-aligned DDP buffers for LayerWise optimizers.
+            False retains compact DDP buffers and variable-size parameter synchronization.
 
     Returns:
         The same megatron_mimo_model with wrapped submodules.
@@ -61,14 +66,26 @@ def wrap_megatron_mimo_model_distributed(
     # Lazy import to avoid circular dependency (models layer loads before training layer)
     from megatron.bridge.training.megatron_mimo_parallel_utils import is_current_rank_in_grid
 
+    disable_encoder_overlap = ddp_config.overlap_grad_reduce is True or ddp_config.overlap_param_gather is True
+    has_modality_grid = any(
+        module_name != MIMO_LANGUAGE_MODULE_KEY and grid is not None for module_name, grid in grids.items()
+    )
+    if disable_encoder_overlap and has_modality_grid:
+        # Emit before per-grid rank filtering so global rank zero reports the
+        # override even when it does not participate in an encoder grid.
+        warn_rank_0(
+            "Disabling gradient-reduce and parameter-gather overlap for modality encoders because "
+            "not every rank is guaranteed to execute an encoder backward pass."
+        )
+
     def wrap_module(
         module: nn.Module,
         config: "TransformerConfig",
         pg: "ProcessGroupCollection",
         module_ddp_config: "DistributedDataParallelConfig",
     ) -> "DistributedDataParallel":
         layout = None
-        if use_layer_wise_distributed_optimizer:
+        if use_layer_wise_distributed_optimizer and use_layer_wise_param_layout:
             module_ddp_config.use_distributed_optimizer = True
             tag_params_for_buffer_routing([module])
             groups = pg.setup_process_groups_for_ddp(pg, config, module_ddp_config)
@@ -144,10 +161,12 @@ def wrap_module(
                         f"a 'config' attribute. Encoders must be MegatronModule subclasses."
                     )
 
-                # An encoder rank that gets no media in a step runs no encoder forward or backward,
-                # so the DDP hooks would not issue its overlapped param all-gather or grad reduce.
-                encoder_ddp_config = dataclasses.replace(
-                    ddp_config, overlap_grad_reduce=False, overlap_param_gather=False
+                # Uneven input participation can leave encoder replicas
+                # unsynchronized when communication is hook-driven.
+                encoder_ddp_config = (
+                    dataclasses.replace(ddp_config, overlap_grad_reduce=False, overlap_param_gather=False)
+                    if disable_encoder_overlap
+                    else ddp_config
                 )
                 wrapped = wrap_module(
                     submodule,

src/megatron/bridge/models/megatron_mimo/megatron_mimo_provider.py

@@ -511,6 +511,7 @@ def provide_distributed_model(
         post_wrap_hook: Optional[Callable[[List[MegatronModule]], List[MegatronModule]]] = None,
         *,
         use_layer_wise_distributed_optimizer: bool = False,
+        use_layer_wise_param_layout: bool = True,
     ) -> List[MegatronModule]:
         """Build MegatronMIMO model with heterogeneous parallelism and DDP wrapping.
 
@@ -545,6 +546,8 @@ def provide_distributed_model(
             pre_wrap_hook: Callable(s) to modify model before wrapping.
             post_wrap_hook: Callable to modify model after wrapping.
             use_layer_wise_distributed_optimizer: Build shard-aligned layouts for layer-wise optimizers.
+            use_layer_wise_param_layout: Use shard-aligned DDP buffers for LayerWise optimizers.
+                False retains compact DDP buffers and variable-size parameter synchronization.
 
         Returns:
             List containing the wrapped MimoModel.
@@ -626,6 +629,7 @@ def provide_distributed_model(
                     grids=infra.module_to_grid_map,
                     pg_collections=infra.pg_collections,
                     use_layer_wise_distributed_optimizer=use_layer_wise_distributed_optimizer,
+                    use_layer_wise_param_layout=use_layer_wise_param_layout,
                 )
                 for m in model_list
             ]

src/megatron/bridge/models/model_provider.py

@@ -215,6 +215,7 @@ def provide_distributed_model(
         pg_collection: ProcessGroupCollection | None = None,
         *,
         use_layer_wise_distributed_optimizer: bool = False,
+        use_layer_wise_param_layout: bool = True,
     ) -> list[ModelT]:
         """Instantiate and wrap the model for distributed training.
 
@@ -225,6 +226,8 @@ def provide_distributed_model(
 
         Args:
             use_layer_wise_distributed_optimizer: Build DDP layouts for layer-wise optimizer state.
+            use_layer_wise_param_layout: Use shard-aligned DDP buffers for LayerWise optimizers.
+                False retains compact DDP buffers and variable-size parameter synchronization.
             ddp_config: Configuration for distributed data parallel.
             model_type: Type of model (encoder, decoder, or both).
             overlap_param_gather_with_optimizer_step: Whether to overlap param gathering.
@@ -307,6 +310,7 @@ def composed_pre_wrap_hook(model: list[MegatronModule]) -> list[MegatronModule]:
             mixed_precision_wrapper=mixed_precision_wrapper,
             pg_collection=pg_collection,
             use_layer_wise_distributed_optimizer=use_layer_wise_distributed_optimizer,
+            use_layer_wise_param_layout=use_layer_wise_param_layout,
         )
 
         if final_post_wrap_hook:
@@ -612,6 +616,7 @@ def get_model(
     *,
     pg_collection: ProcessGroupCollection,
     use_layer_wise_distributed_optimizer: bool = False,
+    use_layer_wise_param_layout: bool = True,
 ) -> list[MegatronModule]:
     """Create and configure a model for distributed training.
 
@@ -635,6 +640,8 @@ def get_model(
         use_megatron_fsdp: Use Megatron's Fully Sharded Data Parallel
         use_torch_fsdp2: Use PyTorch's Fully Sharded Data Parallel v2
         use_layer_wise_distributed_optimizer: Build shard-aligned DDP layouts for layer-wise optimizers.
+        use_layer_wise_param_layout: Use shard-aligned DDP buffers for LayerWise optimizers.
+            False retains compact DDP buffers and variable-size parameter synchronization.
         wrap_with_ddp: Whether to wrap the model with DDP
         data_parallel_random_init: Whether to use random initialization for
             data parallel ranks (vs broadcasting from rank 0)
@@ -740,6 +747,7 @@ def get_model(
             use_megatron_fsdp=use_megatron_fsdp,
             use_torch_fsdp2=use_torch_fsdp2,
             use_layer_wise_distributed_optimizer=use_layer_wise_distributed_optimizer,
+            use_layer_wise_param_layout=use_layer_wise_param_layout,
             pg_collection=pg_collection,
         )
 

src/megatron/bridge/training/checkpointing.py

@@ -1394,9 +1394,10 @@ def save_checkpoint(
     rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
 
     # Collect cfg, model, RNG.
-    sharded_sd_metadata = _build_sharded_state_dict_metadata(uses_distributed_optimizer(cfg.optimizer), ckpt_cfg)
+    use_distributed_optimizer = _uses_distributed_optimizer(cfg.optimizer, optimizer)
+    sharded_sd_metadata = _build_sharded_state_dict_metadata(use_distributed_optimizer, ckpt_cfg)
     sharded_sd_metadata["dp_cp_group"] = _checkpoint_dp_cp_group(pg_collection)
-    if uses_distributed_optimizer(cfg.optimizer):
+    if use_distributed_optimizer:
         print_rank_0(
             f"Storing distributed optimizer sharded state of type {sharded_sd_metadata['distrib_optim_sharding_type']}"
         )
@@ -2913,6 +2914,7 @@ def _load_checkpoint_from_path(
     model = unwrap_model(model)
     pg_collection = pg_collection or get_pg_collection(model)
     ckpt_format = cfg.checkpoint.ckpt_format
+    use_distributed_optimizer = _uses_distributed_optimizer(cfg.optimizer, optimizer)
 
     # HuggingFace-format directories are only honored for pretrained model
     # initialization.  Resume training via ``checkpoint.load`` should use a
@@ -3048,9 +3050,7 @@ def _load_checkpoint_from_path(
 
         if ckpt_type == CheckpointType.LOCAL:
             # Local checkpoints don't store content metadata in common.pt.
-            sharded_sd_metadata = _build_sharded_state_dict_metadata(
-                uses_distributed_optimizer(cfg.optimizer), cfg.checkpoint
-            )
+            sharded_sd_metadata = _build_sharded_state_dict_metadata(use_distributed_optimizer, cfg.checkpoint)
         else:
             sharded_sd_metadata = dist_checkpointing.load_content_metadata(preloaded_state_dict=state_dict)
         print_rank_0(f"sharded_state_dict metadata loaded from the checkpoint: {sharded_sd_metadata}")
@@ -3065,23 +3065,23 @@ def _load_checkpoint_from_path(
             gen_sd_optim = optimizer
             gen_sd_opt_param_scheduler = opt_param_scheduler
 
-            if uses_distributed_optimizer(cfg.optimizer):
+            if use_distributed_optimizer:
                 if sharded_sd_metadata is None:
                     sharded_sd_metadata = {
                         "distrib_optim_sharding_type": (
                             "fully_sharded_model_space"
-                            if run_config["checkpoint"]["fully_parallel_save"]
+                            if not cfg.optimizer.use_distributed_optimizer
+                            or run_config["checkpoint"]["fully_parallel_save"]
                             else "dp_zero_gather_scatter"
                         ),
                     }
-                if (
-                    not tp_pp_match
-                    and sharded_sd_metadata["distrib_optim_sharding_type"]
-                    not in DistributedOptimizer.checkpoint_fully_reshardable_formats
-                ):
+                # Older mixed-optimizer checkpoints omitted this key and used
+                # DistributedOptimizer's model-space default.
+                sharding_type = sharded_sd_metadata.get("distrib_optim_sharding_type", "fully_sharded_model_space")
+                if not tp_pp_match and sharding_type not in DistributedOptimizer.checkpoint_fully_reshardable_formats:
                     raise RuntimeError(
                         f"{mismatch_msg}: not supported for DistributedOptimizer with sharding type"
-                        f" {sharded_sd_metadata['distrib_optim_sharding_type']}."
+                        f" {sharding_type}."
                         f" Please use `checkpoint_config.fully_parallel_save=True` for checkpoint saving."
                     )
         else:
@@ -3199,7 +3199,7 @@ def _load_checkpoint_from_path(
                 print_rank_0("Optimizer state was not saved in the FSDP checkpoint; using a fresh optimizer")
 
         optim_sd_kwargs = dict(
-            metadata=_build_sharded_state_dict_metadata(uses_distributed_optimizer(cfg.optimizer), cfg.checkpoint),
+            metadata=_build_sharded_state_dict_metadata(use_distributed_optimizer, cfg.checkpoint),
             is_loading=True,
         )
         # DTensor optimizer state does not contain dtype-keyed FQNs.
@@ -4149,6 +4149,20 @@ def load_fsdp_dtensor_checkpoint(
     return state_dict, checkpoint_name, release, CheckpointType.FSDP_DTENSOR
 
 
+def _has_distributed_optimizer(optimizer: MegatronOptimizer | None) -> bool:
+    """Find tensor-sharded DistributedOptimizer leaves inside optimizer wrappers."""
+    if isinstance(optimizer, DistributedOptimizer):
+        return True
+    if not isinstance(optimizer, MegatronOptimizer):
+        return False
+    return any(_has_distributed_optimizer(child) for child in getattr(optimizer, "chained_optimizers", ()))
+
+
+def _uses_distributed_optimizer(config: Any, optimizer: MegatronOptimizer | None) -> bool:
+    """Combine declared LayerWise mode with the constructed optimizer tree."""
+    return uses_distributed_optimizer(config) or _has_distributed_optimizer(optimizer)
+
+
 def _build_sharded_state_dict_metadata(use_distributed_optimizer: bool, cfg: CheckpointConfig) -> dict:
     """Builds metadata used for sharded_state_dict versioning.
 

src/megatron/bridge/training/config.py

@@ -134,6 +134,14 @@ class OptimizerConfig(MCoreOptimizerConfig):
     for field modifications after construction but before computed fields are calculated.
     """
 
+    use_layer_wise_param_layout: bool = True
+    """Use the precomputed shard-aligned layout for the LayerWise optimizer.
+
+    Set to ``False`` to retain compact DDP buffers and use LayerWise's
+    variable-size parameter synchronization path. This setting has no effect unless
+    ``use_layer_wise_distributed_optimizer`` is enabled.
+    """
+
     def __post_init__(self) -> None:
         """Skip MCore post_init during initial construction.
 
@@ -2359,9 +2367,10 @@ def _validate_and_sync_distributed_optimizer_settings(config: ConfigContainer) -
     if uses_layer_wise_optimizer(config.optimizer):
         if config.ddp.num_distributed_optimizer_instances != 1:
             raise ValueError("Layer-wise optimizers require num_distributed_optimizer_instances=1.")
-        # Both distributed implementations use reduce-scatter. Model setup must
-        # additionally pass layer-wise intent so Core chooses the correct layout.
-        config.ddp.use_distributed_optimizer = True
+        # This flag selects the ordinary DistributedOptimizer DDP layout. The
+        # model wrapper enables it only when a shard-aligned LayerWise layout is
+        # requested; compact LayerWise mode intentionally leaves it disabled.
+        config.ddp.use_distributed_optimizer = False
     elif ddp_setting or optimizer_setting:
         if ddp_setting != optimizer_setting:
             warn_rank_0(

src/megatron/bridge/training/optim.py

@@ -35,6 +35,7 @@
     OptimizerConfigOverrideProviderContext,
     SchedulerConfig,
 )
+from megatron.bridge.training.utils.optimizer_utils import normalize_optimizer_config
 
 
 G_LOGGER = logging.getLogger(__name__)
@@ -60,6 +61,11 @@ def setup_optimizer(
     Returns:
         tuple containing the optimizer and scheduler
     """
+    if not hasattr(optimizer_config, "provide"):
+        # setup_optimizer is also used by callers that construct Core's config
+        # directly, without ConfigContainer validation.
+        normalize_optimizer_config(optimizer_config)
+
     if optimizer_config_override_provider is None:
         optimizer_config_override_provider = OptimizerConfigOverrideProvider()
 

src/megatron/bridge/training/setup.py

@@ -633,6 +633,7 @@ def _build_distributed_model(cfg: ConfigContainer, pg_collection: ProcessGroupCo
             use_megatron_fsdp=cfg.dist.use_megatron_fsdp,
             use_torch_fsdp2=cfg.dist.use_torch_fsdp2,
             data_parallel_random_init=cfg.rng.data_parallel_random_init,
+            use_layer_wise_param_layout=cfg.optimizer.use_layer_wise_param_layout,
         )
     else:
         model = model_config.provide_distributed_model(
@@ -641,6 +642,7 @@ def _build_distributed_model(cfg: ConfigContainer, pg_collection: ProcessGroupCo
             use_torch_fsdp2=cfg.dist.use_torch_fsdp2,
             overlap_param_gather_with_optimizer_step=cfg.optimizer.overlap_param_gather_with_optimizer_step,
             use_layer_wise_distributed_optimizer=uses_layer_wise_optimizer(cfg.optimizer),
+            use_layer_wise_param_layout=cfg.optimizer.use_layer_wise_param_layout,
             data_parallel_random_init=cfg.rng.data_parallel_random_init,
             pg_collection=pg_collection,
         )

src/megatron/bridge/training/setup_megatron_mimo.py

@@ -138,6 +138,8 @@ def _update_megatron_mimo_model_config_funcs(
 def setup_megatron_mimo(
     state: GlobalState,
     build_data_iterators_fn: Optional[Callable] = None,
+    *,
+    use_layer_wise_param_layout: Optional[bool] = None,
 ) -> MegatronMIMOSetupOutput:
     """MegatronMIMO-specific setup helper.
 
@@ -156,6 +158,8 @@ def setup_megatron_mimo(
             create the optimizer.
         build_data_iterators_fn: Optional function to build data iterators.
             Should have signature: (cfg, megatron_mimo_infra) -> (train_iter, valid_iter)
+        use_layer_wise_param_layout: Override the optimizer config's LayerWise layout choice.
+            False retains compact DDP buffers and variable-size parameter synchronization.
 
     Returns:
         MegatronMIMOSetupOutput containing all components for training.
@@ -187,13 +191,16 @@ def setup_megatron_mimo(
     from megatron.bridge.models.megatron_mimo import build_megatron_mimo_model
 
     megatron_mimo_provider = cfg.model
+    if use_layer_wise_param_layout is None:
+        use_layer_wise_param_layout = cfg.optimizer.use_layer_wise_param_layout
     model, megatron_mimo_infra = build_megatron_mimo_model(
         megatron_mimo_provider,
         ddp_config=cfg.ddp,
         use_layer_wise_distributed_optimizer=uses_layer_wise_optimizer(cfg.optimizer),
         fp16=getattr(cfg.model, "fp16", False),
         bf16=getattr(cfg.model, "bf16", True),
         seed=cfg.rng.seed,
+        use_layer_wise_param_layout=use_layer_wise_param_layout,
     )
 
     logger.info(f"Rank {dist.get_rank()}: Creating multimodule communicator")