diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py
index 2cb454a2027..de47c4ba4ac 100644
--- a/megatron/core/distributed/distributed_data_parallel.py
+++ b/megatron/core/distributed/distributed_data_parallel.py
@@ -473,6 +473,57 @@ def no_sync(self):
             for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
                 bucket_group.is_last_microbatch = True
 
+    def _start_bucket_group_param_sync(
+        self, bucket_group: '_ParamAndGradBucketGroup', force_sync: bool
+    ) -> None:
+        """Dispatch one bucket group's param all-gather + run the FP8 / MXFP8
+        post-all-gather work the synchronous path needs.
+
+        Factored out of :meth:`start_param_sync` so callers that own a subset
+        of bucket groups (e.g. a chained ``LayerWiseDistributedOptimizer`` +
+        ``DistributedOptimizer`` pair) can sync only their own buckets without
+        losing the FP8 post-processing that follows the collective.
+        """
+        bucket_group.start_param_sync(force_sync=force_sync)
+
+        if self.ddp_config.overlap_param_gather:
+            return
+
+        # For MXFP8 params, we need to copy the all-gathered param data from the buffer to
+        # the param.data, since param buffer is not mapped to model params for MXFP8 case.
+        # The paramaters are cast from bf16 to MXFP8 during copy.
+        # In the case of "overlap_param_gather=True", the param copy is done
+        # in "finish_param_sync" stage after zeroing the shared gardient buffers.
+        if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
+            for bucket in bucket_group.buckets:
+                is_bf16_weight_bucket = False
+                for param in bucket.params:
+                    # Skip copying since bf16 weights in the mxfp8 model
+                    # are already mapped to param.data.
+                    if not is_float8tensor(param):
+                        is_bf16_weight_bucket = True
+                        break
+                    param_start, param_end = bucket.param_to_index[param]
+                    param_slice = bucket.param_data.view(-1)[param_start:param_end]
+                    param.data.copy_(param_slice.view(param.data.shape))
+                if is_bf16_weight_bucket:
+                    continue
+                # All-gathered params are not needed after being copied to param.data.
+                # Zero out the param buffer (shared with grad buffer) for gradient
+                # accumulation. We cannot zero out the entire grad buffer because one grad
+                # buffer may correspond to multiple param buffers. If we zero out the entire
+                # grad buffer, it would clear the data of those param buffers that have not
+                # yet completed AG.
+                bucket.param_data.zero_()
+        else:
+            fp8_params = []
+            for bucket in bucket_group.buckets:
+                for param in bucket.params:
+                    if is_float8tensor(param):
+                        fp8_params.append(param)
+            if len(fp8_params) > 0:
+                post_all_gather_processing(fp8_params)
+
     def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bool = False):
         """
         Initiates param sync (all-gather) communication operations for all model parameters.
@@ -493,43 +544,7 @@ def start_param_sync(self, *unused, force_sync: bool = False, force_dispatch: bo
                 return
 
         for bucket_group in self.bucket_groups + self.expert_parallel_bucket_groups:
-            bucket_group.start_param_sync(force_sync=force_sync)
-
-            if not self.ddp_config.overlap_param_gather:
-                # For MXFP8 params, we need to copy the all-gathered param data from the buffer to
-                # the param.data, since param buffer is not mapped to model params for MXFP8 case.
-                # The paramaters are cast from bf16 to MXFP8 during copy.
-                # In the case of "overlap_param_gather=True", the param copy is done
-                # in "finish_param_sync" stage after zeroing the shared gardient buffers.
-                if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag:
-                    for bucket in bucket_group.buckets:
-                        is_bf16_weight_bucket = False
-                        for param in bucket.params:
-                            # Skip copying since bf16 weights in the mxfp8 model
-                            # are already mapped to param.data.
-                            if not is_float8tensor(param):
-                                is_bf16_weight_bucket = True
-                                break
-                            param_start, param_end = bucket.param_to_index[param]
-                            param_slice = bucket.param_data.view(-1)[param_start:param_end]
-                            param.data.copy_(param_slice.view(param.data.shape))
-                        if is_bf16_weight_bucket:
-                            continue
-                        # All-gathered params are not needed after being copied to param.data.
-                        # Zero out the param buffer (shared with grad buffer) for gradient
-                        # accumulation. We cannot zero out the entire grad buffer because one grad
-                        # buffer may correspond to multiple param buffers. If we zero out the entire
-                        # grad buffer, it would clear the data of those param buffers that have not
-                        # yet completed AG.
-                        bucket.param_data.zero_()
-                else:
-                    fp8_params = []
-                    for bucket in bucket_group.buckets:
-                        for param in bucket.params:
-                            if is_float8tensor(param):
-                                fp8_params.append(param)
-                    if len(fp8_params) > 0:
-                        post_all_gather_processing(fp8_params)
+            self._start_bucket_group_param_sync(bucket_group, force_sync=force_sync)
 
     def start_grad_sync(self, *unused):
         """
diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py
index 29f8de1d2d9..9b196a2be57 100644
--- a/megatron/core/distributed/param_and_grad_buffer.py
+++ b/megatron/core/distributed/param_and_grad_buffer.py
@@ -811,15 +811,22 @@ def group_params_for_buffers(
             param_dtype = torch.uint8
         grad_dtype = torch.float if grad_reduce_in_fp32 else param.dtype
         is_expert_parallel = not getattr(param, 'allreduce', True)
+        is_managed_by_layer_wise_optimizer = getattr(
+            param, 'is_managed_by_layer_wise_optimizer', False
+        )
 
-        key = BufferKey(param_dtype, grad_dtype, is_expert_parallel)
+        key = BufferKey(
+            param_dtype, grad_dtype, is_expert_parallel, is_managed_by_layer_wise_optimizer
+        )
         param_list = key_to_params.get(key, [])
         param_list.append(param)
         key_to_params[key] = param_list
 
         # Use param.dtype (not param_dtype) so FP8/NVFP4 params share offsets with their
         # logical high-precision dtype, needed for checkpoint compatibility.
-        offset_key = BufferKey(param.dtype, grad_dtype, is_expert_parallel)
+        offset_key = BufferKey(
+            param.dtype, grad_dtype, is_expert_parallel, is_managed_by_layer_wise_optimizer
+        )
         offset = dtype_to_offsets.get(offset_key, 0)
         dtype_to_offsets[offset_key] = offset + 1
         indices = key_to_indices.get(key, [])
@@ -1498,12 +1505,16 @@ def partition_buckets(
     if len(buffers) == 0:
         return []
 
-    dtype_to_buffer_map = {}
+    # At most one fp8 (uint8) buffer is allowed; Cases 2 and 3 below branch on
+    # whether one is present. Non-uint8 dtypes can legitimately appear in
+    # multiple buffers (e.g. LayerWise-managed bf16 weights + Adam-managed bf16
+    # biases share the bf16 ``param_dtype`` but live in separate buffers), so
+    # the uniqueness check is restricted to uint8.
+    fp8_buffer = None
     for buffer in buffers:
-        dtype = buffer.param_dtype
-        # Make sure that the param_dtype of any two buffers is different.
-        assert dtype not in dtype_to_buffer_map
-        dtype_to_buffer_map[dtype] = buffer
+        if buffer.param_dtype == torch.uint8:
+            assert fp8_buffer is None
+            fp8_buffer = buffer
 
     # Case 1: Put all buckets into a single bucket group if force_single_bucket_group is True.
     if force_single_bucket_group:
@@ -1522,7 +1533,7 @@ def partition_buckets(
         )
         return [bucket_group]
 
-    if torch.uint8 not in dtype_to_buffer_map:
+    if fp8_buffer is None:
         # Case 2: When there is no fp8 buffer in the input buffers, let each bucket group have
         #         only one bucket.
         bucket_groups = []
@@ -1546,7 +1557,6 @@ def partition_buckets(
                     non_fp8_buckets.append(bucket)
 
         bucket_groups = []
-        fp8_buffer = dtype_to_buffer_map[torch.uint8]
         for bucket in fp8_buffer.buckets:
             if len(bucket_groups) == len(fp8_buffer.buckets) - 1:
                 # reduce_scatter_with_fp32_accumulation requires exactly one bucket
diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index ebdd42effe2..ef90bce40ae 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -67,7 +67,7 @@
     _create_emerging_optimizer,
 )
 from .grad_scaler import ConstantGradScaler, DynamicGradScaler
-from .layer_wise_optimizer import LayerWiseDistributedOptimizer
+from .layer_wise_optimizer import LayerWiseDistributedOptimizer, is_managed_by_layer_wise_optimizer
 from .optimizer import (
     ChainedOptimizer,
     Float16OptimizerWithFloat16Params,
@@ -217,7 +217,7 @@ def is_vector_like_parameter(param: torch.nn.Parameter, param_name: str) -> bool
     def is_muon_managed_matrix_parameter(param: torch.nn.Parameter, _: str) -> bool:
         if not is_muon_optimizer:
             return False
-        return param.dim() == 2 and not getattr(param, 'is_embedding_or_output_parameter', False)
+        return is_managed_by_layer_wise_optimizer(param)
 
     def should_scale_lr_with_mup(param: torch.nn.Parameter, param_name: str) -> bool:
         if decoupled_lr_enabled and getattr(param, 'is_embedding_or_output_parameter', False):
@@ -757,6 +757,12 @@ def _get_megatron_emerging_optimizer(
             f"emerging-optimizers package is required for optimizer='{eopt_name}'. "
             "Install it with: pip install emerging-optimizers"
         )
+    assert not (use_layer_wise and config.overlap_param_gather_with_optimizer_step), (
+        "overlap_param_gather_with_optimizer_step is not supported with "
+        "use_layer_wise_distributed_optimizer: the emerging-optimizer path does not "
+        "split model_chunks into (first, rest) groups, so the per-chunk param-gather "
+        "dispatch never fires. Disable one of the two flags."
+    )
     if eopt_name not in _EMERGING_OPTIMIZERS:
         raise ValueError(f"Unsupported emerging optimizer: {eopt_name}")
     if config.fp16:
@@ -790,8 +796,61 @@ def _get_megatron_emerging_optimizer(
         is_expert = group['is_expert_parallel'] and not use_layer_wise
         grouped_param_groups[(opt_name, is_expert)].append(group)
 
+    # Set up DistOpt process groups + filtered buffers once, only if we'll
+    # construct a DistributedOptimizer for non-Muon groups in layer-wise mode.
+    # The DistOpt-vs-LayerWise buffer split only happens when DDP was wrapped
+    # with ``use_distributed_optimizer=True`` (i.e. the layout-based path); in
+    # legacy ping-pong mode all params share one unpadded DDP buffer that
+    # DistOpt cannot manage, so we keep non-Muon params inside LayerWise.
+    ddp_uses_distributed_optimizer = (
+        bool(getattr(model_chunks[0], 'ddp_config', None))
+        and model_chunks[0].ddp_config.use_distributed_optimizer
+    )
+    distopt_process_groups = None
+    distopt_per_model_buffers = None
+    use_separate_distributed_optimizer = ddp_uses_distributed_optimizer and use_layer_wise
+    if use_separate_distributed_optimizer:
+        ddp_config = model_chunks[0].ddp_config
+        assert ddp_config.num_distributed_optimizer_instances == 1, (
+            "Layer-wise + DistributedOptimizer split path does not yet support "
+            "num_distributed_optimizer_instances > 1: distributed_optimizer_instance_id "
+            "is hardcoded to 0 in this path. Disable use_layer_wise_param_layout to "
+            "fall back to the legacy LayerWise ping-pong path."
+        )
+    if use_separate_distributed_optimizer and any(
+        opt_name not in _EMERGING_OPTIMIZERS
+        for (opt_name, _), groups in grouped_param_groups.items()
+        if groups
+    ):
+        # ``setup_process_groups_for_optimizer`` rejects Gloo groups whenever
+        # an explicit ``pg_collection`` is supplied, so the only legal value
+        # here is False.
+        distopt_process_groups = ProcessGroupCollection.setup_process_groups_for_optimizer(
+            pg_collection, model_chunks, use_gloo_process_groups=False
+        )
+        # DistOpt should only manage non-LayerWise buffers (those holding
+        # embeddings, biases, layernorm, etc.). Filter out the LayerWise
+        # shard-aligned buffers that the LayerWiseDistributedOptimizer owns.
+        distopt_per_model_buffers = {}
+        for model_chunk_idx, model_chunk in enumerate(model_chunks):
+            if not hasattr(model_chunk, 'buffers'):
+                continue
+            non_layer_wise_buffers = [
+                buffer
+                for buffer in model_chunk.buffers
+                if buffer.params
+                and not getattr(buffer.params[0], 'is_managed_by_layer_wise_optimizer', False)
+            ]
+            if non_layer_wise_buffers:
+                distopt_per_model_buffers[model_chunk_idx] = non_layer_wise_buffers
+
     # Build an optimizer for each (optimizer_name, is_expert) bucket and combine.
+    # In layer-wise mode, emerging-optimizer (Muon) groups feed into LayerWise,
+    # while non-emerging (Adam) groups are managed by a separate DistributedOptimizer
+    # — that is, the LayerWise optimizer only owns Muon-managed matrix parameters,
+    # and the rest go through DistOpt's standard byte-level shard machinery.
     results = []
+    layer_wise_base_results = []  # (raw_optimizer, init_state_fn) feeding LayerWise.
     for (opt_name, is_expert), groups in grouped_param_groups.items():
         if not groups:
             continue
@@ -803,55 +862,99 @@ def _get_megatron_emerging_optimizer(
                 config, groups, eopt_name, model_chunks, pg_collection
             )
             if use_layer_wise:
-                result = (optimizer, init_state_fn)
+                layer_wise_base_results.append((optimizer, init_state_fn))
+                continue
+            if config.bf16:
+                optimizer = Float16OptimizerWithFloat16Params(
+                    optimizer, config, None, init_state_fn
+                )
             else:
-                if config.bf16:
-                    optimizer = Float16OptimizerWithFloat16Params(
-                        optimizer, config, None, init_state_fn
-                    )
-                else:
-                    optimizer = FP32Optimizer(optimizer, config, init_state_fn)
-                setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group)
-                if pg_collection is None or not hasattr(pg_collection, 'tp'):
-                    tp_group = parallel_state.get_tensor_model_parallel_group()
-                else:
-                    tp_group = pg_collection.tp
-                setattr(optimizer, 'tp_group', tp_group)
-                result = optimizer
+                optimizer = FP32Optimizer(optimizer, config, init_state_fn)
+            setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group)
+            if pg_collection is None or not hasattr(pg_collection, 'tp'):
+                tp_group = parallel_state.get_tensor_model_parallel_group()
+            else:
+                tp_group = pg_collection.tp
+            setattr(optimizer, 'tp_group', tp_group)
+            results.append(optimizer)
+            continue
         else:
             fallback_config = copy.copy(config)
             fallback_config.optimizer = opt_name
-            fallback_config.use_distributed_optimizer = False
-            result = _get_megatron_optimizer_based_on_param_groups(
-                config=fallback_config,
-                model_chunks=model_chunks,
-                param_groups=groups,
-                model_parallel_group=model_parallel_group,
-                pg_collection=pg_collection,
-                skip_megatron_wrapping=use_layer_wise,
-            )
-            # TODO(deyuf): ChainedOptimizer currently asserts all sub-optimizers
-            # share the same config. Revisit this design now that emerging
-            # optimizers mix different optimizer types (e.g. Muon + Adam).
-            # For now, reset to the top-level config so the assertion holds.
-            if not use_layer_wise and hasattr(result, 'config'):
-                result.config = config
-        results.append(result)
+            if use_separate_distributed_optimizer:
+                # Route non-emerging params through a real DistributedOptimizer
+                # (byte-level sharding) instead of stuffing them inside LayerWise.
+                for group in groups:
+                    assert not group['is_expert_parallel'], (
+                        "Non-emerging expert-parallel param groups are not yet "
+                        "supported on the layer-wise + DistributedOptimizer "
+                        "path: they need a separate DistOpt instance with the "
+                        "expert-DP process group, which is not wired up yet. "
+                        "Disable use_layer_wise_param_layout to fall back to "
+                        "the legacy LayerWise ping-pong path for MoE models."
+                    )
+                fallback_config.use_distributed_optimizer = True
+                result = _get_megatron_optimizer_based_on_param_groups(
+                    config=fallback_config,
+                    model_chunks=model_chunks,
+                    param_groups=groups,
+                    per_model_buffers=distopt_per_model_buffers,
+                    model_parallel_group=distopt_process_groups['mp_group'],
+                    data_parallel_group=distopt_process_groups['intra_dp_cp_group'],
+                    data_parallel_group_gloo=distopt_process_groups['intra_dp_cp_group_gloo'],
+                    data_parallel_group_idx=get_pg_rank(distopt_process_groups['mp_group']),
+                    intra_dist_opt_group=distopt_process_groups['intra_dist_opt_group'],
+                    distributed_optimizer_instance_id=0,
+                    pg_collection=pg_collection,
+                    skip_megatron_wrapping=False,
+                )
+                # TODO(deyuf): ChainedOptimizer currently asserts all sub-optimizers
+                # share the same config. Reset to the top-level config so the
+                # assertion holds when DistOpt+LayerWise are chained.
+                if hasattr(result, 'config'):
+                    result.config = config
+                results.append(result)
+            else:
+                # Legacy ping-pong layer-wise path (use_layer_wise=True) or the
+                # non-layer-wise standard chain: keep ``use_distributed_optimizer``
+                # off; in layer-wise mode the raw torch optimizer (returned as a
+                # ``(optimizer, init_state_fn)`` tuple via ``skip_megatron_wrapping``)
+                # feeds into ``LayerWiseDistributedOptimizer``.
+                fallback_config.use_distributed_optimizer = False
+                result = _get_megatron_optimizer_based_on_param_groups(
+                    config=fallback_config,
+                    model_chunks=model_chunks,
+                    param_groups=groups,
+                    model_parallel_group=model_parallel_group,
+                    pg_collection=pg_collection,
+                    skip_megatron_wrapping=use_layer_wise,
+                )
+                if use_layer_wise:
+                    layer_wise_base_results.append(result)
+                else:
+                    if hasattr(result, 'config'):
+                        result.config = config
+                    results.append(result)
 
     if use_layer_wise:
-        base_optimizers, init_fns = (), ()
-        if results:
-            base_optimizers, init_fns = zip(*results)
         log_single_rank(
             logger, logging.INFO, f'Using LayerWiseDistributedOptimizer for {eopt_name}'
         )
-        return LayerWiseDistributedOptimizer(
+        base_optimizers, init_fns = (), ()
+        if layer_wise_base_results:
+            base_optimizers, init_fns = zip(*layer_wise_base_results)
+        layer_wise_optimizer = LayerWiseDistributedOptimizer(
             list(base_optimizers),
             config,
             pg_collection,
             init_state_fn_list=list(init_fns),
             model_chunks=model_chunks,
         )
+        # LayerWise owns Muon-managed params; DistOpt instances in ``results``
+        # own the rest. Chain them so the training loop sees one optimizer.
+        if results:
+            return ChainedOptimizer([layer_wise_optimizer] + results)
+        return layer_wise_optimizer
 
     return ChainedOptimizer(results)
 
diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
index e0d3c2a54ac..e66fa9dfe55 100644
--- a/megatron/core/optimizer/distrib_optimizer.py
+++ b/megatron/core/optimizer/distrib_optimizer.py
@@ -2981,6 +2981,34 @@ def copy_group_params(model_groups, shard_main_groups):
         copy_group_params(self.model_float16_groups, self.shard_fp32_from_float16_groups)
         copy_group_params(self.model_fp32_groups, self.shard_fp32_groups)
 
+    def start_param_sync_for_bucket_group_subset(self) -> None:
+        """Trigger ``start_param_sync`` on DistOpt-managed bucket groups only.
+
+        Walks each model chunk's DDP bucket groups and skips those tagged
+        ``is_managed_by_layer_wise_optimizer=True`` (so a sibling
+        :class:`LayerWiseDistributedOptimizer` does not double-sync the same
+        buckets). When no LayerWise tagging is present every bucket group is
+        included — matching the previous ``model_chunk.start_param_sync()``
+        behaviour. Uses :meth:`DistributedDataParallel._start_bucket_group_param_sync`
+        so FP8 post-all-gather processing (and MXFP8 copy) still runs.
+        """
+        # Deferred import: layer_wise_optimizer's compute_full_param_layout
+        # lazily imports DistributedOptimizer, so importing the helper at
+        # module load here would create a cycle.
+        from .layer_wise_optimizer import _bucket_is_managed_by_layer_wise_optimizer
+
+        for model_chunk in self.model_chunks:
+            for bucket_group in (
+                model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups
+            ):
+                if not bucket_group.buckets:
+                    continue
+                if _bucket_is_managed_by_layer_wise_optimizer(
+                    bucket_group.buckets[0], default_for_untagged=False
+                ):
+                    continue
+                model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False)
+
     @torch.no_grad()
     def step_with_ready_grads(self) -> bool:
         """Step the optimizer with ready gradients, return successful.
@@ -3004,8 +3032,10 @@ def step_with_ready_grads(self) -> bool:
             # the first all-gather is launched asynchronously in the next optimizer.zero_grad()
             # call and subsequent all-gathers are launched in the forward pre-hook.
             if not self.ddp_config.overlap_param_gather:
-                for model_chunk in self.model_chunks:
-                    model_chunk.start_param_sync()
+                # Only sync DistOpt-managed bucket groups so a sibling
+                # LayerWiseDistributedOptimizer's own ``start_param_sync`` call
+                # is not duplicated for the same buckets.
+                self.start_param_sync_for_bucket_group_subset()
         if timers is not None:
             timers('params-all-gather').stop()
 
diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py
index f60934cee26..e8e173ffe06 100644
--- a/megatron/core/optimizer/layer_wise_optimizer.py
+++ b/megatron/core/optimizer/layer_wise_optimizer.py
@@ -2,7 +2,6 @@
 
 import logging
 import math
-from collections import defaultdict
 from typing import Callable, Dict, List, Optional, Tuple
 
 import torch
@@ -33,6 +32,58 @@
 logger = logging.getLogger(__name__)
 
 
+def is_managed_by_layer_wise_optimizer(param: torch.nn.Parameter) -> bool:
+    """Whether a parameter is managed by :class:`LayerWiseDistributedOptimizer`.
+
+    Returns True for the 2D matrix-like weight parameters that Muon orthogonalizes
+    via Newton-Schulz, and False for embeddings, biases, LayerNorm weights, and
+    any other non-matrix parameter (which are handled by Adam through a separate
+    :class:`DistributedOptimizer`).
+
+    Mirrors the routing rule applied by ``_get_param_groups`` /
+    ``default_param_overrides`` for Muon.
+    """
+    if not param.dim() == 2:
+        return False
+    if getattr(param, 'is_embedding_or_output_parameter', False):
+        return False
+    return True
+
+
+def _bucket_is_managed_by_layer_wise_optimizer(bucket, default_for_untagged: bool = True) -> bool:
+    """Whether a DDP bucket belongs to a LayerWise-managed buffer.
+
+    Buckets are built from params that share a :class:`BufferKey`, so checking
+    the first param's tag is sufficient. ``default_for_untagged`` controls the
+    legacy (no-tagging) case: callers asking "is this mine?" from the LayerWise
+    side pass ``True`` (legacy LayerWise owns everything); callers asking from
+    the DistOpt side pass ``False`` (legacy DistOpt also owns everything, so
+    untagged buckets are *not* LayerWise-managed).
+    """
+    if not bucket.params_list:
+        return False
+    param = bucket.params_list[0]
+    if not hasattr(param, 'is_managed_by_layer_wise_optimizer'):
+        return default_for_untagged
+    return param.is_managed_by_layer_wise_optimizer
+
+
+def tag_params_for_buffer_routing(model_chunks) -> None:
+    """Tag every requires-grad param with ``is_managed_by_layer_wise_optimizer``.
+
+    Run this once on the un-DDP-wrapped model chunks before
+    :class:`DistributedDataParallel` constructs its grad/param buffers — the
+    grouping function ``group_params_for_buffers`` reads this attribute to
+    decide which buffer each param lands in (LayerWise shard-aligned buffer vs
+    DistOpt-style byte-level buffer).
+    """
+    for model_chunk in model_chunks:
+        for param in model_chunk.parameters():
+            if not param.requires_grad:
+                continue
+            param.is_managed_by_layer_wise_optimizer = is_managed_by_layer_wise_optimizer(param)
+
+
 class LayerWiseDistributedOptimizer(ChainedOptimizer):
     """Layer-wise distributed optimizer for Megatron-core models.
 
@@ -69,29 +120,30 @@ def _compute_per_buffer_param_layout(
         ddp_config,
         param_indices: Optional[List[int]] = None,
     ) -> 'PerBufferParamLayout':
-        """Compute parameter layout with shard-aligned buckets via size-matching.
+        """Compute parameter layout with shard-aligned buckets via LPT bin-packing.
 
-        Assigns parameters to ``dp_size`` equal-sized shards within each bucket
-        so that no parameter is ever split across a shard boundary.
+        Assigns parameters to ``dp_size`` shards within each bucket so that no
+        parameter is split across a shard boundary, while keeping each bucket
+        confined to a contiguous range in backprop order.
 
         **Algorithm** (operates in reverse model / backprop order):
 
-        1. Separate shared-embedding parameters (isolated buckets, emitted first).
-        2. Pool the remaining parameters in backprop order, indexed by numel.
-        3. Pop the next unassigned parameter and assign it to shard 0.
-        4. For shards 1 … ``dp_size - 1``, assign the next unassigned parameter
-           of the same numel (also in backprop order).  If none is available,
-           insert padding of that numel.  Every shard grows by the same amount,
-           so all shards stay the same size.
-        5. When the bucket total reaches *bucket_size*, finalise the bucket
-           (pad shard size to :meth:`_shard_divisor`) and start a new one.
-        6. Repeat from 3 until all parameters are assigned.
-
-        Because repeated layers produce many parameters of the same shape,
-        size-matching naturally keeps whole parameters together without any
-        name-parsing heuristic.  Padding overhead is low (depending on number
-        of layers and number of shards) — zero when every shape group has a
-        count divisible by ``dp_size``.
+        1. Walk parameters in backprop order, accumulating them into a chunk.
+           A shared (tied) embedding triggers an immediate finalisation
+           followed by an isolated bucket for that embedding alone.
+        2. When the chunk's total numel reaches ``bucket_size`` (or all
+           params have been consumed), bin-pack the chunk into ``dp_size``
+           shards via greedy LPT — sort by numel descending and assign each
+           param to the shard with the smallest current load.
+        3. Pad each shard to ``max(shard_cursors)`` aligned to
+           :meth:`_shard_divisor`, then emit the bucket.
+
+        Each bucket therefore spans a contiguous backprop range so that
+        ``overlap_grad_reduce`` can dispatch the bucket's reduce-scatter as
+        soon as the bucket's backward segment finishes — preserving the
+        original DDP overlap semantics.  LPT bin-packing keeps shards close
+        to balanced; for uniform transformer blocks where ``params_per_layer
+        * num_layers`` is a multiple of ``dp_size`` the packing is perfect.
 
         Args:
             params: Parameters in model-definition (forward) order.
@@ -106,184 +158,116 @@ def _compute_per_buffer_param_layout(
         dp_size = data_parallel_world_size
         shard_divisor = LayerWiseDistributedOptimizer._shard_divisor(dp_size, ddp_config)
 
-        # -- 0. Separate shared-embedding params. -------------------------
-        shared_embedding_params: List[torch.nn.Parameter] = []
-        regular_params: List[torch.nn.Parameter] = []
-        total_param_numel = 0
-        for param in params:
-            total_param_numel += param.data.nelement()
-            if getattr(param, 'shared_embedding', False):
-                shared_embedding_params.append(param)
-            else:
-                regular_params.append(param)
-
-        # -- 1. Build backprop-order pool & per-size index. ---------------
-        pool = list(reversed(regular_params))
-        assigned_param_ids: set[int] = set()  # id(param) of assigned params
-
-        size_groups: Dict[int, List[torch.nn.Parameter]] = defaultdict(list)
-        for param in pool:
-            size_groups[param.data.nelement()].append(param)
-        size_cursors: Dict[int, int] = defaultdict(int)
-
-        overall_cursor = 0
-
-        def _next_unassigned() -> Optional[torch.nn.Parameter]:
-            nonlocal overall_cursor
-            while overall_cursor < len(pool):
-                if id(pool[overall_cursor]) not in assigned_param_ids:
-                    return pool[overall_cursor]
-                overall_cursor += 1
-            return None
-
-        def _next_with_size(param_numel: int) -> Optional[torch.nn.Parameter]:
-            """Next unassigned param of size *param_numel* in backprop order."""
-            group = size_groups[param_numel]
-            cursor = size_cursors[param_numel]
-            while cursor < len(group):
-                if id(group[cursor]) not in assigned_param_ids:
-                    size_cursors[param_numel] = cursor
-                    return group[cursor]
-                cursor += 1
-            size_cursors[param_numel] = cursor
-            return None
-
-        # -- 2. Output accumulators and per-bucket shard state. ----------
+        total_param_numel = sum(p.data.nelement() for p in params)
+
         param_index_map: Dict[torch.nn.Parameter, Tuple[int, int, int]] = {}
         bucket_indices: List[Tuple[int, int]] = []
         per_bucket_numel_unpadded: List[int] = []
-        buffer_cursor = 0  # write position in the contiguous buffer
+        buffer_cursor = 0
         bucket_id = 0
-
-        # Per-shard state for the bucket currently being built.
-        # `shard_assignments[i]` holds an ordered list of (param | None, numel)
-        # entries to be written into shard i; a `None` entry is empty padding
-        # that keeps every shard the same size.
-        shard_assignments: List[List[Tuple[Optional[torch.nn.Parameter], int]]] = [
-            [] for _ in range(dp_size)
-        ]
-        shard_cursor = 0  # position within each shard (identical for all shards)
-        bucket_numel_unpadded = 0
-        size_match_padding_numel = 0  # elements used for empty-shard-slot padding
-
-        def _finalize_bucket() -> None:
-            nonlocal buffer_cursor, bucket_id, shard_assignments
-            nonlocal shard_cursor, bucket_numel_unpadded
-            if shard_cursor == 0:
+        shard_imbalance_padding_numel = 0
+
+        def _emit_bucket(
+            chunk_params: List[torch.nn.Parameter], shared_embedding: bool = False
+        ) -> None:
+            """Bin-pack *chunk_params* into ``dp_size`` shards and emit a bucket.
+
+            With ``shared_embedding=True``, the chunk must contain a single
+            parameter; it goes into shard 0 with same-size padding in
+            shards 1..dp_size-1 so the embedding fits entirely within one
+            shard (needed for the cross-stage tied-embedding all-reduce).
+            """
+            nonlocal buffer_cursor, bucket_id, shard_imbalance_padding_numel
+            if not chunk_params:
                 return
-            padded_shard_size = pad_to_divisor(shard_cursor, shard_divisor)
-            bucket_start_index = buffer_cursor
 
+            shard_assignments: List[List[Tuple[Optional[torch.nn.Parameter], int]]] = [
+                [] for _ in range(dp_size)
+            ]
+            shard_cursors = [0] * dp_size
+
+            if shared_embedding:
+                assert len(chunk_params) == 1
+                param = chunk_params[0]
+                numel = param.data.nelement()
+                shard_assignments[0].append((param, numel))
+                shard_cursors[0] = numel
+                for shard_id in range(1, dp_size):
+                    shard_assignments[shard_id].append((None, numel))
+                    shard_cursors[shard_id] = numel
+            else:
+                # Greedy LPT: largest first, assign to the least-loaded shard.
+                # The within-shard order is sorted-by-numel, not backprop —
+                # that is fine because all params in the chunk share the same
+                # bucket_id, so DDP's backprop-order iteration still sees
+                # monotonic bucket_ids across the chunk boundary.
+                for param in sorted(chunk_params, key=lambda p: -p.data.nelement()):
+                    numel = param.data.nelement()
+                    min_shard = min(range(dp_size), key=lambda s: shard_cursors[s])
+                    placement = pad_param_start(shard_cursors[min_shard])
+                    shard_assignments[min_shard].append((param, numel))
+                    shard_cursors[min_shard] = placement + numel
+
+            padded_shard_size = pad_to_divisor(max(shard_cursors), shard_divisor)
+            bucket_start_index = buffer_cursor
             for shard_id in range(dp_size):
                 shard_start_index = bucket_start_index + shard_id * padded_shard_size
                 cursor = shard_start_index
-                for param, numel in shard_assignments[shard_id]:
+                for p, numel in shard_assignments[shard_id]:
                     cursor = pad_param_start(cursor)
-                    if param is not None:
-                        param_index_map[param] = (cursor, cursor + numel, bucket_id)
+                    if p is not None:
+                        param_index_map[p] = (cursor, cursor + numel, bucket_id)
                     cursor += numel
-
+                shard_imbalance_padding_numel += padded_shard_size - shard_cursors[shard_id]
             bucket_end_index = bucket_start_index + dp_size * padded_shard_size
             bucket_indices.append((bucket_start_index, bucket_end_index))
-            per_bucket_numel_unpadded.append(bucket_numel_unpadded)
+            per_bucket_numel_unpadded.append(sum(p.data.nelement() for p in chunk_params))
             buffer_cursor = bucket_end_index
             bucket_id += 1
 
-            shard_assignments = [[] for _ in range(dp_size)]
-            shard_cursor = 0
-            bucket_numel_unpadded = 0
-
-        # -- 3. Emit one isolated bucket per shared-embedding param. -----
-        # Shared (tied) embeddings need their own bucket — typically because
-        # input and output embeddings are tied across pipeline-parallel
-        # stages and need a cross-stage all-reduce. Each shared embedding
-        # occupies shard 0 of its bucket alone; shards 1..dp_size-1 are
-        # filled with empty (padding) slots of the same numel so the bucket
-        # is shard-aligned and the embedding fits entirely within shard 0.
+        # Each chunk spans a contiguous backprop range. Bucket ids therefore
+        # increase monotonically when ``_ParamAndGradBuffer.__init__`` iterates
+        # params in backprop order, satisfying its ``bucket_id == cur + 1``
+        # invariant.
         #
-        # NOTE: This is expensive. Padding cost per shared embedding is
-        # (dp_size - 1) * pad_to_divisor(numel, shard_divisor) elements,
-        # which for a vocab x hidden embedding (e.g. 128k x 8192) at dp_size
-        # = 8 is roughly 7 * (vocab * hidden) elements — many GBs of the
-        # param buffer (and again of the grad buffer) per shared embedding.
-        # The cost is unavoidable while preserving the "no parameter crosses
-        # a shard boundary" invariant the layerwise scheme depends on for
-        # correct reduce-scatter + local optimizer step.
-        for param in reversed(shared_embedding_params):
+        # Padding floor: the on-buffer bucket size is ``dp_size *
+        # max_shard_cursor``, which is at least ``dp_size * chunk_max_param``
+        # because some shard must hold that param whole. If a single param
+        # dominates the chunk, finalising on ``chunk_numel >= bucket_size``
+        # alone would emit a bucket with most of its shards near-empty
+        # padding. Instead extend the chunk so its raw numel approaches the
+        # padded buffer size, capping per-bucket overhead at ``1 /
+        # PADDING_FLOOR - 1`` (~11% at 0.9). Falls back to ``bucket_size``
+        # when no single param dominates.
+        PADDING_FLOOR = 0.9
+        chunk_params: List[torch.nn.Parameter] = []
+        chunk_numel = 0
+        chunk_max_param = 0
+        for param in reversed(params):
             param_numel = param.data.nelement()
-            assigned_param_ids.add(id(param))
-            shard_assignments[0].append((param, param_numel))
-            bucket_numel_unpadded += param_numel
-            # No size-matching: each shared embedding must be alone in its
-            # bucket. Pad shards 1..dp_size-1 with same-size empty slots.
-            for shard_id in range(1, dp_size):
-                shard_assignments[shard_id].append((None, param_numel))
-                size_match_padding_numel += param_numel
-            shard_cursor = pad_param_start(shard_cursor) + param_numel
-            _finalize_bucket()
-
-        # -- 4. Size-matching loop for regular params. --------------------
-        while True:
-            param = _next_unassigned()
-            if param is None:
-                break
-
-            param_numel = param.data.nelement()
-            assigned_param_ids.add(id(param))
-            shard_assignments[0].append((param, param_numel))
-            bucket_numel_unpadded += param_numel
-
-            for shard_id in range(1, dp_size):
-                # Prefer an exact-numel peer; this gives the cleanest layout
-                # (no inner-shard padding).
-                matched_param = _next_with_size(param_numel)
-                if matched_param is not None:
-                    assigned_param_ids.add(id(matched_param))
-                    shard_assignments[shard_id].append((matched_param, param_numel))
-                    bucket_numel_unpadded += param_numel
-                    continue
-
-                # No exact peer. Greedily pack as many smaller params from the
-                # queue as fit within this shard slot (sized to ``param_numel``).
-                # Cuts overhead from unique-large seeds (e.g. an embedding)
-                # that would otherwise force ``(dp_size - 1) * param_numel`` of
-                # empty padding.
-                useful_in_slot = 0
-                slot_cursor = 0
-                while True:
-                    candidate_param = _next_unassigned()
-                    if candidate_param is None:
-                        break
-                    candidate_numel = candidate_param.data.nelement()
-                    candidate_start = pad_param_start(slot_cursor)
-                    if candidate_start + candidate_numel > param_numel:
-                        break
-                    assigned_param_ids.add(id(candidate_param))
-                    shard_assignments[shard_id].append((candidate_param, candidate_numel))
-                    bucket_numel_unpadded += candidate_numel
-                    slot_cursor = candidate_start + candidate_numel
-                    useful_in_slot += candidate_numel
-
-                # Pad the remainder of the slot up to ``param_numel``.
-                padding_start = pad_param_start(slot_cursor)
-                padding_size = param_numel - padding_start
-                if padding_size > 0:
-                    shard_assignments[shard_id].append((None, padding_size))
-                size_match_padding_numel += param_numel - useful_in_slot
-
-            shard_cursor = pad_param_start(shard_cursor) + param_numel
-
+            if getattr(param, 'shared_embedding', False):
+                # Finalize any in-progress chunk so the shared-embedding
+                # bucket comes after it in backprop order.
+                _emit_bucket(chunk_params)
+                chunk_params = []
+                chunk_numel = 0
+                chunk_max_param = 0
+                _emit_bucket([param], shared_embedding=True)
+                continue
+            chunk_params.append(param)
+            chunk_numel += param_numel
+            chunk_max_param = max(chunk_max_param, param_numel)
             if bucket_size is not None:
-                bucket_total = dp_size * pad_to_divisor(shard_cursor, shard_divisor)
-                if bucket_total >= bucket_size:
-                    _finalize_bucket()
-
-        _finalize_bucket()
+                threshold = max(bucket_size, int(dp_size * chunk_max_param * PADDING_FLOOR))
+                if chunk_numel >= threshold:
+                    _emit_bucket(chunk_params)
+                    chunk_params = []
+                    chunk_numel = 0
+                    chunk_max_param = 0
+        _emit_bucket(chunk_params)
 
-        # -- 5. Log padding overhead. ------------------------------------
         total_buffer_numel = bucket_indices[-1][1] if bucket_indices else 0
         total_padding = total_buffer_numel - total_param_numel
-        alignment_and_shard_end_padding = total_padding - size_match_padding_numel
         log_single_rank(
             logger,
             logging.INFO,
@@ -293,8 +277,7 @@ def _finalize_bucket() -> None:
             f"total_param_numel={total_param_numel}, "
             f"total_buffer_numel={total_buffer_numel}, "
             f"total_padding={total_padding} "
-            f"(size_match={size_match_padding_numel}, "
-            f"alignment+shard_end={alignment_and_shard_end_padding}), "
+            f"(shard_imbalance={shard_imbalance_padding_numel}), "
             f"overhead={total_padding / max(total_param_numel, 1) * 100:.1f}%",
         )
 
@@ -331,6 +314,9 @@ def compute_full_param_layout(
         Returns:
             :class:`FullParamLayout` with a :class:`PerBufferParamLayout` per buffer group.
         """
+        # Avoid a circular import: DistributedOptimizer imports LayerWise indirectly.
+        from .distrib_optimizer import DistributedOptimizer
+
         buffer_groups = group_params_for_buffers(params, ddp_config.grad_reduce_in_fp32)
         layouts = {}
         for buffer_key, (group_params, param_indices) in buffer_groups.items():
@@ -343,7 +329,16 @@ def compute_full_param_layout(
             else:
                 dp_world_size = data_parallel_world_size
 
-            layouts[buffer_key] = LayerWiseDistributedOptimizer._compute_per_buffer_param_layout(
+            # Dispatch per buffer: LayerWise (Muon) params get the shard-aligned
+            # layout; non-LayerWise params (e.g. Adam-managed embeddings, biases)
+            # get DistOpt's byte-level layout.
+            if buffer_key.is_managed_by_layer_wise_optimizer:
+                compute_per_buffer_layout = (
+                    LayerWiseDistributedOptimizer._compute_per_buffer_param_layout
+                )
+            else:
+                compute_per_buffer_layout = DistributedOptimizer._compute_per_buffer_param_layout
+            layouts[buffer_key] = compute_per_buffer_layout(
                 group_params, bucket_size, dp_world_size, ddp_config, param_indices
             )
         return FullParamLayout(layouts=layouts)
@@ -476,6 +471,11 @@ def _shard_params_from_layout(self, optimizers, full_param_layouts, dp_cp_size,
         param_to_shard: Dict[torch.nn.Parameter, int] = {}
         for full_layout in full_param_layouts:
             for buffer_key, layout in full_layout.layouts.items():
+                # Non-LayerWise buffers (e.g. Adam-managed embeddings, biases,
+                # layernorms with a DistOpt byte-level layout) are managed by a
+                # separate DistributedOptimizer; LayerWise does not own them.
+                if not buffer_key.is_managed_by_layer_wise_optimizer:
+                    continue
                 dp_size = expt_dp_size if buffer_key.is_expert_parallel else dp_cp_size
                 for param, (
                     param_start_index,
@@ -595,6 +595,8 @@ def set_bucket_layerwise_params_list(self, model_chunks):
         for model_chunk in model_chunks:
             for group in model_chunk.bucket_groups:
                 for bucket in group.buckets:
+                    if not _bucket_is_managed_by_layer_wise_optimizer(bucket):
+                        continue
                     bucket_params_list = [[] for _ in range(get_pg_size(self.pg_collection.dp_cp))]
                     for bucket_list, full_params_list in zip(
                         bucket_params_list, self.dp_cp_params_list
@@ -606,6 +608,8 @@ def set_bucket_layerwise_params_list(self, model_chunks):
             # Do the same for expert parallel bucket groups.
             for group in model_chunk.expert_parallel_bucket_groups:
                 for bucket in group.buckets:
+                    if not _bucket_is_managed_by_layer_wise_optimizer(bucket):
+                        continue
                     if self.expt_dp_params_list is not None:
                         bucket_params_list = [
                             [] for _ in range(get_pg_size(self.pg_collection.expt_dp))
@@ -713,31 +717,52 @@ def count_zeros(self):
             use_decoupled_grad=self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8,
         )
 
+    def start_param_sync_for_bucket_group_subset(self) -> None:
+        """Trigger ``start_param_sync`` on LayerWise-managed bucket groups only.
+
+        Walks each model chunk's dense + expert-parallel bucket groups and
+        skips any group not managed by LayerWise, so a sibling
+        :class:`DistributedOptimizer`'s own ``start_param_sync`` call does not
+        double-sync the same buckets. Uses
+        :meth:`DistributedDataParallel._start_bucket_group_param_sync` so FP8
+        post-all-gather processing (and MXFP8 copy) still runs.
+        """
+        for model_chunk in self.model_chunks:
+            for bucket_group in (
+                model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups
+            ):
+                if bucket_group.buckets and _bucket_is_managed_by_layer_wise_optimizer(
+                    bucket_group.buckets[0]
+                ):
+                    model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False)
+
     @torch.no_grad()
-    def step(self):  # type: ignore[no-untyped-def]
-        """step function for layer-wise optimizer.
+    def step_with_ready_grads(self) -> bool:
+        """Step then all-gather LayerWise-managed param buffers.
 
-        NOTE: bypassed when this optimizer is a child of an outer
-        ChainedOptimizer; in that case the sibling DistributedOptimizer's
-        step_with_ready_grads handles the param sync.
+        Placed on ``step_with_ready_grads`` (not ``step``) so the param sync also
+        runs when this optimizer is a child of an outer ``ChainedOptimizer``,
+        which calls ``step_with_ready_grads`` directly on each child and bypasses
+        ``step``.
         """
-        update_successful, grad_norm, num_zeros_in_grad = super().step()
+        success = super().step_with_ready_grads()
 
         # All-gather updated params. If overlap_param_gather is True, the all-gather
         # is deferred to the forward pre-hooks via DDP bucket infrastructure.
         if not self.overlap_param_gather:
             if self.use_buffer_param_sync:
                 # Model params are views into the DDP param buffer
-                # (ddp_config.use_distributed_optimizer=True).  The optimizer step
+                # (ddp_config.use_distributed_optimizer=True). The optimizer step
                 # already copied updated fp32 main params → bf16 model params (=
-                # buffer views), so the buffer is up-to-date.  Trigger the standard
-                # buffer all-gather (matches DistributedOptimizer's call site).
-                for model_chunk in self.model_chunks:
-                    model_chunk.start_param_sync()
+                # buffer views), so the buffer is up-to-date. Trigger the standard
+                # buffer all-gather, but only for LayerWise-managed bucket groups
+                # so a sibling DistributedOptimizer's own ``start_param_sync`` call
+                # is not duplicated for the same buckets.
+                self.start_param_sync_for_bucket_group_subset()
             else:
                 self.allgather_params()
 
-        return update_successful, grad_norm, num_zeros_in_grad
+        return success
 
     # TODO(deyuf): need to improve dist checkpointing design to properly handle this
     # fp32_from_fp16_params is list, each sub list could be empty if group is empty
diff --git a/megatron/core/optimizer/param_layout.py b/megatron/core/optimizer/param_layout.py
index 543af88f325..2ee511c6126 100644
--- a/megatron/core/optimizer/param_layout.py
+++ b/megatron/core/optimizer/param_layout.py
@@ -54,11 +54,16 @@ class BufferKey:
         grad_dtype: Gradient reduction dtype.
         is_expert_parallel: Whether the buffer holds expert-parallel parameters,
             which use a separate data-parallel group.
+        is_managed_by_layer_wise_optimizer: Whether parameters in this buffer are
+            managed by :class:`LayerWiseDistributedOptimizer` (shard-aligned layout
+            so each whole param lives in one shard). Non-LayerWise params get
+            :class:`DistributedOptimizer`'s byte-level layout in a separate buffer.
     """
 
     param_dtype: torch.dtype
     grad_dtype: torch.dtype
     is_expert_parallel: bool
+    is_managed_by_layer_wise_optimizer: bool = False
 
 
 @dataclass
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index f3c2ded6907..4d524e53aae 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2765,6 +2765,17 @@ def _add_distributed_args(parser):
                        dest='align_param_gather')
     group.add_argument('--use-distributed-optimizer', action='store_true',
                        help='Use distributed optimizer.')
+    group.add_argument('--no-use-layer-wise-param-layout',
+                       action='store_false',
+                       dest='use_layer_wise_param_layout',
+                       help='Opt out of the precomputed LayerWise param layout. When set, '
+                       'falls back to the legacy LayerWise ping-pong path: all params '
+                       '(including non-Muon embeddings, biases, layernorm) live in a single '
+                       'LayerWise buffer and the optimizer uses the allgather_params() codepath. '
+                       'The default (precomputed layout) routes non-Muon params through a '
+                       'separate DistributedOptimizer with byte-level sharding, which is faster '
+                       'and uses less padding but produces different bf16 reduction ordering '
+                       'and so will not match legacy-path loss curves bit-for-bit.')
     group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub',
                        help='Use the userbuffer registration for DP/FSDP communication buffers.'
                        'This option will reduce GPU SM usage for the DP/FSDP communication,'
diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py
index fac717557a7..119c872aa71 100644
--- a/megatron/training/checkpointing.py
+++ b/megatron/training/checkpointing.py
@@ -431,10 +431,18 @@ def _build_sharded_state_dict_metadata(args: Namespace, dp_cp_group: Optional[to
     """
     metadata = {}
 
-    if args.use_distributed_optimizer and args.ckpt_format == "fsdp_dtensor":
+    # ``use_layer_wise_distributed_optimizer`` shares the DistOpt sub-optimizer
+    # for non-Muon params (embeddings, biases, layernorm, ...), so it needs the
+    # same sharding-type metadata even though the parser flips
+    # ``use_distributed_optimizer`` off in that mode.
+    has_distributed_optimizer = args.use_distributed_optimizer or getattr(
+        args, 'use_layer_wise_distributed_optimizer', False
+    )
+
+    if has_distributed_optimizer and args.ckpt_format == "fsdp_dtensor":
         metadata['distrib_optim_sharding_type'] = 'fsdp_dtensor'
 
-    if args.use_distributed_optimizer and args.ckpt_format != "fsdp_dtensor":
+    if has_distributed_optimizer and args.ckpt_format != "fsdp_dtensor":
         if args.dist_ckpt_optim_fully_reshardable:
             metadata['distrib_optim_sharding_type'] = 'fully_reshardable'
             metadata['distrib_optim_fully_reshardable_mem_efficient'] = args.distrib_optim_fully_reshardable_mem_efficient
diff --git a/megatron/training/training.py b/megatron/training/training.py
index e55272402cc..c9404ca527f 100644
--- a/megatron/training/training.py
+++ b/megatron/training/training.py
@@ -52,7 +52,10 @@ def set_startup_timestamps(program_start=None, main_entry=None):
 import torch.distributed
 
 from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer
-from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer
+from megatron.core.optimizer.layer_wise_optimizer import (
+    LayerWiseDistributedOptimizer,
+    tag_params_for_buffer_routing,
+)
 from megatron.core.optimizer_param_scheduler import get_canonical_lr_for_logging
 from .log_handler import CustomHandler
 
@@ -1375,6 +1378,11 @@ def wrap_model_chunks_with_ddp(
         if use_layer_wise_distributed_optimizer and use_layer_wise_param_layout:
             ddp_config.use_distributed_optimizer = True
             compute_layout = LayerWiseDistributedOptimizer.compute_full_param_layout
+            # Tag params so DDP buffer grouping routes LayerWise-managed matrices
+            # (Muon's Newton-Schulz domain) to a shard-aligned buffer and routes
+            # everything else (embeddings, biases, layernorm) to a separate
+            # DistOpt-style buffer.
+            tag_params_for_buffer_routing(model_chunks)
         elif not use_layer_wise_distributed_optimizer and ddp_config.use_distributed_optimizer:
             compute_layout = DistributedOptimizer.compute_full_param_layout
         else:
@@ -1601,7 +1609,9 @@ def build_model():
                 use_layer_wise_distributed_optimizer=getattr(
                     args, 'use_layer_wise_distributed_optimizer', False
                 ),
-                use_layer_wise_param_layout=False,
+                use_layer_wise_param_layout=getattr(
+                    args, 'use_layer_wise_param_layout', True
+                ),
                 DP=DP,
                 pg_collection=pg_collection if args.use_megatron_fsdp else None,
                 bucket_sizes=per_chunk_bucket_sizes,
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml
index 028074bd34f..fffad86a016 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml
@@ -66,4 +66,5 @@ MODEL_ARGS:
   --async-save: true
   --use-persistent-ckpt-worker: true
   --use-distributed-optimizer: true
+  --no-use-layer-wise-param-layout: true
 TEST_TYPE: ckpt-resume
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/model_config.yaml
index 40ac94eed9b..874e7dccf8d 100644
--- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/model_config.yaml
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/model_config.yaml
@@ -66,4 +66,5 @@ MODEL_ARGS:
   --async-save: true
   --use-persistent-ckpt-worker: true
   --use-distributed-optimizer: true
+  --no-use-layer-wise-param-layout: true
 TEST_TYPE: ckpt-resume
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/golden_values_dev_dgx_h100.json
new file mode 100644
index 00000000000..bfb742c4407
--- /dev/null
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/golden_values_dev_dgx_h100.json
@@ -0,0 +1,537 @@
+{
+    "lm loss": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 10.90111,
+            "2": 10.8957,
+            "3": 10.90636,
+            "4": 10.90598,
+            "5": 10.91363,
+            "6": 10.89758,
+            "7": 10.89741,
+            "8": 10.91392,
+            "9": 10.89606,
+            "10": 10.90437,
+            "11": 10.9006,
+            "12": 10.91884,
+            "13": 10.88464,
+            "14": 10.87537,
+            "15": 10.90982,
+            "16": 10.90335,
+            "17": 10.88689,
+            "18": 10.89443,
+            "19": 10.89475,
+            "20": 10.89108,
+            "21": 10.88678,
+            "22": 10.90366,
+            "23": 10.91203,
+            "24": 10.88567,
+            "25": 10.90092,
+            "26": 10.88848,
+            "27": 10.89687,
+            "28": 10.88207,
+            "29": 10.88987,
+            "30": 10.91122,
+            "31": 10.89171,
+            "32": 10.889,
+            "33": 10.89725,
+            "34": 10.87363,
+            "35": 10.89659,
+            "36": 10.90839,
+            "37": 10.8708,
+            "38": 10.87871,
+            "39": 10.88589,
+            "40": 10.8922,
+            "41": 10.88393,
+            "42": 10.89545,
+            "43": 10.88027,
+            "44": 10.88452,
+            "45": 10.88357,
+            "46": 10.88635,
+            "47": 10.88548,
+            "48": 10.86363,
+            "49": 10.87585,
+            "50": 10.88305,
+            "51": 10.89289,
+            "52": 10.87105,
+            "53": 10.85916,
+            "54": 10.8708,
+            "55": 10.8682,
+            "56": 10.87334,
+            "57": 10.84465,
+            "58": 10.85629,
+            "59": 10.84681,
+            "60": 10.84338,
+            "61": 10.86051,
+            "62": 10.85659,
+            "63": 10.86048,
+            "64": 10.83796,
+            "65": 10.82775,
+            "66": 10.8461,
+            "67": 10.82815,
+            "68": 10.82996,
+            "69": 10.81904,
+            "70": 10.82857,
+            "71": 10.82638,
+            "72": 10.80839,
+            "73": 10.8073,
+            "74": 10.8066,
+            "75": 10.81434,
+            "76": 10.81156,
+            "77": 10.80737,
+            "78": 10.79215,
+            "79": 10.80177,
+            "80": 10.78932,
+            "81": 10.79557,
+            "82": 10.79565,
+            "83": 10.78673,
+            "84": 10.75649,
+            "85": 10.76263,
+            "86": 10.77681,
+            "87": 10.79626,
+            "88": 10.7749,
+            "89": 10.77573,
+            "90": 10.76463,
+            "91": 10.74072,
+            "92": 10.76008,
+            "93": 10.74645,
+            "94": 10.73453,
+            "95": 10.75255,
+            "96": 10.72309,
+            "97": 10.7149,
+            "98": 10.72599,
+            "99": 10.74586,
+            "100": 10.69546
+        }
+    },
+    "num-zeros": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 1133.0,
+            "2": 1120.0,
+            "3": 1414.0,
+            "4": 1257.0,
+            "5": 1241.0,
+            "6": 1300.0,
+            "7": 1244.0,
+            "8": 1140.0,
+            "9": 1222.0,
+            "10": 1194.0,
+            "11": 1259.0,
+            "12": 1267.0,
+            "13": 1267.0,
+            "14": 1269.0,
+            "15": 1222.0,
+            "16": 1210.0,
+            "17": 1184.0,
+            "18": 1083.0,
+            "19": 1227.0,
+            "20": 1146.0,
+            "21": 1219.0,
+            "22": 1242.0,
+            "23": 1261.0,
+            "24": 1183.0,
+            "25": 1228.0,
+            "26": 1232.0,
+            "27": 1315.0,
+            "28": 1194.0,
+            "29": 1253.0,
+            "30": 1252.0,
+            "31": 1195.0,
+            "32": 1299.0,
+            "33": 1218.0,
+            "34": 1220.0,
+            "35": 1318.0,
+            "36": 1248.0,
+            "37": 1181.0,
+            "38": 1257.0,
+            "39": 1375.0,
+            "40": 1282.0,
+            "41": 1342.0,
+            "42": 1140.0,
+            "43": 1219.0,
+            "44": 1218.0,
+            "45": 1157.0,
+            "46": 1431.0,
+            "47": 1240.0,
+            "48": 1132.0,
+            "49": 1372.0,
+            "50": 1321.0,
+            "51": 1254.0,
+            "52": 1226.0,
+            "53": 1289.0,
+            "54": 1196.0,
+            "55": 1183.0,
+            "56": 1150.0,
+            "57": 1149.0,
+            "58": 1329.0,
+            "59": 1119.0,
+            "60": 1187.0,
+            "61": 1225.0,
+            "62": 1243.0,
+            "63": 1273.0,
+            "64": 1299.0,
+            "65": 1196.0,
+            "66": 1261.0,
+            "67": 1200.0,
+            "68": 1281.0,
+            "69": 1144.0,
+            "70": 1293.0,
+            "71": 1188.0,
+            "72": 1208.0,
+            "73": 1229.0,
+            "74": 1275.0,
+            "75": 1326.0,
+            "76": 1295.0,
+            "77": 1223.0,
+            "78": 1271.0,
+            "79": 1274.0,
+            "80": 1202.0,
+            "81": 1269.0,
+            "82": 1127.0,
+            "83": 1254.0,
+            "84": 1180.0,
+            "85": 1257.0,
+            "86": 1408.0,
+            "87": 1070.0,
+            "88": 1192.0,
+            "89": 1162.0,
+            "90": 1383.0,
+            "91": 1252.0,
+            "92": 1274.0,
+            "93": 1343.0,
+            "94": 1296.0,
+            "95": 1061.0,
+            "96": 1138.0,
+            "97": 1182.0,
+            "98": 1371.0,
+            "99": 1108.0,
+            "100": 1211.0
+        }
+    },
+    "mem-allocated-bytes": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 742158848.0,
+            "2": 742149632.0,
+            "3": 742124544.0,
+            "4": 742104576.0,
+            "5": 742155776.0,
+            "6": 742138368.0,
+            "7": 742066176.0,
+            "8": 742102528.0,
+            "9": 742105088.0,
+            "10": 742136320.0,
+            "11": 742110208.0,
+            "12": 742118912.0,
+            "13": 742122496.0,
+            "14": 742126592.0,
+            "15": 742103040.0,
+            "16": 742086656.0,
+            "17": 742095872.0,
+            "18": 742097920.0,
+            "19": 742127616.0,
+            "20": 742092800.0,
+            "21": 742147584.0,
+            "22": 742172160.0,
+            "23": 742144000.0,
+            "24": 742061568.0,
+            "25": 742142464.0,
+            "26": 742171648.0,
+            "27": 742126080.0,
+            "28": 742078976.0,
+            "29": 742151168.0,
+            "30": 742084096.0,
+            "31": 742091264.0,
+            "32": 742098432.0,
+            "33": 742106624.0,
+            "34": 742085632.0,
+            "35": 742069248.0,
+            "36": 742135296.0,
+            "37": 742082560.0,
+            "38": 742140416.0,
+            "39": 742149120.0,
+            "40": 742096896.0,
+            "41": 742116352.0,
+            "42": 742060544.0,
+            "43": 742126592.0,
+            "44": 742111744.0,
+            "45": 742084096.0,
+            "46": 742150144.0,
+            "47": 742106624.0,
+            "48": 742136832.0,
+            "49": 742095872.0,
+            "50": 742143488.0,
+            "51": 742125568.0,
+            "52": 742096896.0,
+            "53": 742048256.0,
+            "54": 742146048.0,
+            "55": 742135296.0,
+            "56": 742134272.0,
+            "57": 742192128.0,
+            "58": 742129664.0,
+            "59": 742134272.0,
+            "60": 742113280.0,
+            "61": 742084096.0,
+            "62": 742074368.0,
+            "63": 742084096.0,
+            "64": 742110720.0,
+            "65": 742118400.0,
+            "66": 742132736.0,
+            "67": 742115328.0,
+            "68": 742116864.0,
+            "69": 742080000.0,
+            "70": 742163968.0,
+            "71": 742141952.0,
+            "72": 742139392.0,
+            "73": 742070272.0,
+            "74": 742113792.0,
+            "75": 742156288.0,
+            "76": 742112256.0,
+            "77": 742155264.0,
+            "78": 742101504.0,
+            "79": 742167552.0,
+            "80": 742088192.0,
+            "81": 742142976.0,
+            "82": 742090752.0,
+            "83": 742152192.0,
+            "84": 742119936.0,
+            "85": 742129664.0,
+            "86": 742112256.0,
+            "87": 742126080.0,
+            "88": 742107648.0,
+            "89": 742098432.0,
+            "90": 742119936.0,
+            "91": 742105600.0,
+            "92": 742105088.0,
+            "93": 742134784.0,
+            "94": 742109696.0,
+            "95": 742133760.0,
+            "96": 742124032.0,
+            "97": 742093824.0,
+            "98": 742130176.0,
+            "99": 742147584.0,
+            "100": 742176768.0
+        }
+    },
+    "mem-max-allocated-bytes": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 3131866112.0,
+            "2": 3239834624.0,
+            "3": 3239834624.0,
+            "4": 3239834624.0,
+            "5": 3243102208.0,
+            "6": 3243102208.0,
+            "7": 3243102208.0,
+            "8": 3243102208.0,
+            "9": 3243102208.0,
+            "10": 3243102208.0,
+            "11": 3243102208.0,
+            "12": 3243102208.0,
+            "13": 3243102208.0,
+            "14": 3243102208.0,
+            "15": 3243102208.0,
+            "16": 3243102208.0,
+            "17": 3243102208.0,
+            "18": 3243102208.0,
+            "19": 3243102208.0,
+            "20": 3243102208.0,
+            "21": 3243102208.0,
+            "22": 3253252096.0,
+            "23": 3253252096.0,
+            "24": 3253252096.0,
+            "25": 3253252096.0,
+            "26": 3253252096.0,
+            "27": 3253252096.0,
+            "28": 3253252096.0,
+            "29": 3253252096.0,
+            "30": 3253252096.0,
+            "31": 3253252096.0,
+            "32": 3253252096.0,
+            "33": 3253252096.0,
+            "34": 3253252096.0,
+            "35": 3253252096.0,
+            "36": 3253252096.0,
+            "37": 3253252096.0,
+            "38": 3253252096.0,
+            "39": 3253252096.0,
+            "40": 3253252096.0,
+            "41": 3253252096.0,
+            "42": 3253252096.0,
+            "43": 3253252096.0,
+            "44": 3253252096.0,
+            "45": 3253252096.0,
+            "46": 3253252096.0,
+            "47": 3253252096.0,
+            "48": 3253252096.0,
+            "49": 3253252096.0,
+            "50": 3253252096.0,
+            "51": 3253252096.0,
+            "52": 3253252096.0,
+            "53": 3253252096.0,
+            "54": 3253252096.0,
+            "55": 3253252096.0,
+            "56": 3253252096.0,
+            "57": 3278233088.0,
+            "58": 3278233088.0,
+            "59": 3278233088.0,
+            "60": 3278233088.0,
+            "61": 3278233088.0,
+            "62": 3278233088.0,
+            "63": 3278233088.0,
+            "64": 3278233088.0,
+            "65": 3278233088.0,
+            "66": 3278233088.0,
+            "67": 3278233088.0,
+            "68": 3278233088.0,
+            "69": 3278233088.0,
+            "70": 3278233088.0,
+            "71": 3278233088.0,
+            "72": 3278233088.0,
+            "73": 3278233088.0,
+            "74": 3278233088.0,
+            "75": 3278233088.0,
+            "76": 3278233088.0,
+            "77": 3278233088.0,
+            "78": 3278233088.0,
+            "79": 3278233088.0,
+            "80": 3278233088.0,
+            "81": 3278233088.0,
+            "82": 3278233088.0,
+            "83": 3278233088.0,
+            "84": 3278233088.0,
+            "85": 3278233088.0,
+            "86": 3278233088.0,
+            "87": 3278233088.0,
+            "88": 3278233088.0,
+            "89": 3278233088.0,
+            "90": 3278233088.0,
+            "91": 3278233088.0,
+            "92": 3278233088.0,
+            "93": 3278233088.0,
+            "94": 3278233088.0,
+            "95": 3278233088.0,
+            "96": 3278233088.0,
+            "97": 3278233088.0,
+            "98": 3278233088.0,
+            "99": 3278233088.0,
+            "100": 3278233088.0
+        }
+    },
+    "iteration-time": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": "nan",
+            "2": 4.59241,
+            "3": 0.25053,
+            "4": 0.22419,
+            "5": 0.23038,
+            "6": 0.22482,
+            "7": 0.22216,
+            "8": 0.23153,
+            "9": 0.21307,
+            "10": 0.21854,
+            "11": 0.22355,
+            "12": 0.21397,
+            "13": 0.20266,
+            "14": 0.21483,
+            "15": 0.22306,
+            "16": 0.20837,
+            "17": 0.22541,
+            "18": 0.20493,
+            "19": 0.21347,
+            "20": 0.21206,
+            "21": 0.20129,
+            "22": 0.20518,
+            "23": 0.21795,
+            "24": 0.20156,
+            "25": 0.21433,
+            "26": 0.21328,
+            "27": 0.20173,
+            "28": 0.19444,
+            "29": 0.20863,
+            "30": 0.20281,
+            "31": 0.19444,
+            "32": 0.19557,
+            "33": 0.19712,
+            "34": 0.20688,
+            "35": 0.20375,
+            "36": 0.19114,
+            "37": 0.19802,
+            "38": 0.20284,
+            "39": 0.19627,
+            "40": 0.2057,
+            "41": 0.19683,
+            "42": 0.19719,
+            "43": 0.19229,
+            "44": 0.19583,
+            "45": 0.20139,
+            "46": 0.20022,
+            "47": 0.20197,
+            "48": 0.19785,
+            "49": 0.19701,
+            "50": 0.19251,
+            "51": 0.26454,
+            "52": 0.23231,
+            "53": 0.20298,
+            "54": 0.19903,
+            "55": 0.18945,
+            "56": 0.18994,
+            "57": 0.20554,
+            "58": 0.1945,
+            "59": 0.18888,
+            "60": 0.19456,
+            "61": 0.20203,
+            "62": 0.19831,
+            "63": 0.18835,
+            "64": 0.1974,
+            "65": 0.18951,
+            "66": 0.19088,
+            "67": 0.20067,
+            "68": 0.19509,
+            "69": 0.19122,
+            "70": 0.19079,
+            "71": 0.18734,
+            "72": 0.19512,
+            "73": 0.19257,
+            "74": 0.18802,
+            "75": 0.18904,
+            "76": 0.19175,
+            "77": 0.19111,
+            "78": 0.19207,
+            "79": 0.19136,
+            "80": 0.19355,
+            "81": 0.19454,
+            "82": 0.19598,
+            "83": 0.18855,
+            "84": 0.19641,
+            "85": 0.19109,
+            "86": 0.20118,
+            "87": 0.19366,
+            "88": 0.19951,
+            "89": 0.19552,
+            "90": 0.18723,
+            "91": 0.19973,
+            "92": 0.1958,
+            "93": 0.1899,
+            "94": 0.19261,
+            "95": 0.1964,
+            "96": 0.19581,
+            "97": 0.20057,
+            "98": 0.19693,
+            "99": 0.19173,
+            "100": 0.19609
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/model_config.yaml
new file mode 100644
index 00000000000..028074bd34f
--- /dev/null
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/model_config.yaml
@@ -0,0 +1,69 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: :4096:8
+MODEL_ARGS:
+  --num-layers: 12
+  --hidden-size: 512
+  --num-attention-heads: 8
+  --log-params-norm: true
+  --log-num-zeros-in-grad: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-timers-to-tensorboard: true
+  --tensorboard-dir: ${TENSORBOARD_PATH}
+  --micro-batch-size: 4
+  --global-batch-size: 32
+  --seq-length: 1024
+  --max-position-embeddings: 1024
+  --disable-bias-linear: true
+  --train-iters: 100
+  --timing-log-level: 0
+  --lr-decay-iters: 320000
+  --save: ${CHECKPOINT_SAVE_PATH}
+  --load: ${CHECKPOINT_LOAD_PATH}
+  --data-path: ${DATA_PATH}/text/common_pile/v01_filtered_data/my-gpt3_00_text_document
+  --vocab-file: ${DATA_PATH}/text/common_pile/v01_filtered_data/bpe/vocab.json
+  --merge-file: ${DATA_PATH}/text/common_pile/v01_filtered_data/bpe/merges.txt
+  --split: 949,50,1
+  --distributed-backend: nccl
+  --lr: 0.00015
+  --lr-decay-style: cosine
+  --min-lr: 1.0e-5
+  --weight-decay: 1e-2
+  --clip-grad: 1.0
+  --lr-warmup-fraction: .01
+  --log-interval: 1
+  --save-interval: 50
+  --eval-interval: 1000
+  --eval-iters: 10
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 1
+  --pipeline-model-parallel-size: 1
+  --expert-model-parallel-size: 8
+  --num-experts: 8
+  --moe-token-dispatcher-type: allgather
+  --moe-router-load-balancing-type: aux_loss
+  --moe-router-topk: 2
+  --moe-router-dtype: fp32
+  --moe-ffn-hidden-size: 1024
+  --moe-grouped-gemm: true
+  --ckpt-fully-parallel-load: true
+  --deterministic-mode: true
+  --no-gradient-accumulation-fusion: true
+  --attention-softmax-in-fp32: true
+  --use-checkpoint-opt_param-scheduler: true
+  --use-mcore-models: true
+  --ckpt-format: torch_dist
+  --data-cache-path: ${DATA_CACHE_PATH}
+  --bf16: true
+  --no-bias-gelu-fusion: true
+  --log-memory-to-tensorboard: true
+  --optimizer: muon
+  --muon-momentum: 0.9
+  --muon-extra-scale-factor: 0.2
+  --muon-scale-mode: spectral
+  --async-save: true
+  --use-persistent-ckpt-worker: true
+  --use-distributed-optimizer: true
+TEST_TYPE: ckpt-resume
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/golden_values_dev_dgx_gb200.json
new file mode 100644
index 00000000000..569607281bf
--- /dev/null
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/golden_values_dev_dgx_gb200.json
@@ -0,0 +1,537 @@
+{
+    "lm loss": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 10.93064,
+            "2": 10.92264,
+            "3": 10.92452,
+            "4": 10.91698,
+            "5": 10.92714,
+            "6": 10.92087,
+            "7": 10.91619,
+            "8": 10.91819,
+            "9": 10.9381,
+            "10": 10.92374,
+            "11": 10.91546,
+            "12": 10.92551,
+            "13": 10.9344,
+            "14": 10.91609,
+            "15": 10.92842,
+            "16": 10.91843,
+            "17": 10.92694,
+            "18": 10.91283,
+            "19": 10.93032,
+            "20": 10.90036,
+            "21": 10.94024,
+            "22": 10.92585,
+            "23": 10.92991,
+            "24": 10.90536,
+            "25": 10.91474,
+            "26": 10.92588,
+            "27": 10.92819,
+            "28": 10.92039,
+            "29": 10.92618,
+            "30": 10.91402,
+            "31": 10.91028,
+            "32": 10.91864,
+            "33": 10.91978,
+            "34": 10.90072,
+            "35": 10.91939,
+            "36": 10.90729,
+            "37": 10.91269,
+            "38": 10.92345,
+            "39": 10.91192,
+            "40": 10.91091,
+            "41": 10.91315,
+            "42": 10.90865,
+            "43": 10.90552,
+            "44": 10.90735,
+            "45": 10.89775,
+            "46": 10.90666,
+            "47": 10.90179,
+            "48": 10.88408,
+            "49": 10.90017,
+            "50": 10.89886,
+            "51": 10.89801,
+            "52": 10.90418,
+            "53": 10.90136,
+            "54": 10.88777,
+            "55": 10.88666,
+            "56": 10.88882,
+            "57": 10.88709,
+            "58": 10.88849,
+            "59": 10.88262,
+            "60": 10.87372,
+            "61": 10.87294,
+            "62": 10.86897,
+            "63": 10.87602,
+            "64": 10.85859,
+            "65": 10.86061,
+            "66": 10.85694,
+            "67": 10.86217,
+            "68": 10.86035,
+            "69": 10.84972,
+            "70": 10.86675,
+            "71": 10.84874,
+            "72": 10.84042,
+            "73": 10.84832,
+            "74": 10.83856,
+            "75": 10.83646,
+            "76": 10.83499,
+            "77": 10.82935,
+            "78": 10.82313,
+            "79": 10.82891,
+            "80": 10.82946,
+            "81": 10.81756,
+            "82": 10.81674,
+            "83": 10.80491,
+            "84": 10.78233,
+            "85": 10.78341,
+            "86": 10.8006,
+            "87": 10.79552,
+            "88": 10.79515,
+            "89": 10.77697,
+            "90": 10.78319,
+            "91": 10.78928,
+            "92": 10.78346,
+            "93": 10.75504,
+            "94": 10.76051,
+            "95": 10.77292,
+            "96": 10.7352,
+            "97": 10.74134,
+            "98": 10.74715,
+            "99": 10.7623,
+            "100": 10.74254
+        }
+    },
+    "num-zeros": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 49180.0,
+            "2": 50476.0,
+            "3": 49424.0,
+            "4": 49432.0,
+            "5": 50158.0,
+            "6": 49809.0,
+            "7": 51674.0,
+            "8": 49844.0,
+            "9": 50774.0,
+            "10": 51711.0,
+            "11": 50055.0,
+            "12": 48173.0,
+            "13": 51987.0,
+            "14": 50012.0,
+            "15": 47793.0,
+            "16": 49432.0,
+            "17": 49553.0,
+            "18": 52046.0,
+            "19": 51104.0,
+            "20": 52161.0,
+            "21": 49606.0,
+            "22": 49387.0,
+            "23": 53274.0,
+            "24": 51418.0,
+            "25": 49019.0,
+            "26": 52564.0,
+            "27": 51550.0,
+            "28": 48436.0,
+            "29": 50200.0,
+            "30": 51269.0,
+            "31": 52612.0,
+            "32": 49440.0,
+            "33": 52159.0,
+            "34": 50829.0,
+            "35": 51311.0,
+            "36": 51173.0,
+            "37": 49650.0,
+            "38": 51020.0,
+            "39": 50641.0,
+            "40": 51446.0,
+            "41": 47582.0,
+            "42": 50391.0,
+            "43": 50658.0,
+            "44": 47747.0,
+            "45": 54687.0,
+            "46": 47497.0,
+            "47": 50374.0,
+            "48": 50491.0,
+            "49": 54486.0,
+            "50": 50450.0,
+            "51": 48947.0,
+            "52": 50771.0,
+            "53": 49531.0,
+            "54": 49393.0,
+            "55": 48915.0,
+            "56": 48727.0,
+            "57": 48700.0,
+            "58": 52622.0,
+            "59": 50435.0,
+            "60": 49573.0,
+            "61": 47247.0,
+            "62": 50690.0,
+            "63": 48469.0,
+            "64": 56352.0,
+            "65": 50850.0,
+            "66": 49341.0,
+            "67": 52467.0,
+            "68": 49616.0,
+            "69": 55080.0,
+            "70": 49654.0,
+            "71": 49521.0,
+            "72": 52208.0,
+            "73": 52701.0,
+            "74": 49718.0,
+            "75": 50138.0,
+            "76": 51332.0,
+            "77": 49802.0,
+            "78": 49721.0,
+            "79": 48872.0,
+            "80": 52104.0,
+            "81": 49959.0,
+            "82": 47202.0,
+            "83": 52816.0,
+            "84": 48777.0,
+            "85": 50514.0,
+            "86": 49122.0,
+            "87": 46396.0,
+            "88": 48097.0,
+            "89": 49004.0,
+            "90": 52069.0,
+            "91": 50371.0,
+            "92": 50285.0,
+            "93": 50754.0,
+            "94": 51462.0,
+            "95": 48801.0,
+            "96": 49986.0,
+            "97": 49690.0,
+            "98": 48768.0,
+            "99": 49360.0,
+            "100": 47284.0
+        }
+    },
+    "mem-allocated-bytes": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 946676736.0,
+            "2": 946695680.0,
+            "3": 946680320.0,
+            "4": 946677760.0,
+            "5": 946689024.0,
+            "6": 946675712.0,
+            "7": 946671104.0,
+            "8": 946700288.0,
+            "9": 946647552.0,
+            "10": 946648064.0,
+            "11": 946632192.0,
+            "12": 946664960.0,
+            "13": 946663936.0,
+            "14": 946653696.0,
+            "15": 946683904.0,
+            "16": 946672128.0,
+            "17": 946682880.0,
+            "18": 946660864.0,
+            "19": 946656256.0,
+            "20": 946669056.0,
+            "21": 946690048.0,
+            "22": 946625536.0,
+            "23": 946667520.0,
+            "24": 946681856.0,
+            "25": 946646016.0,
+            "26": 946681856.0,
+            "27": 946686464.0,
+            "28": 946650624.0,
+            "29": 946669056.0,
+            "30": 946678784.0,
+            "31": 946642944.0,
+            "32": 946661376.0,
+            "33": 946648064.0,
+            "34": 946691072.0,
+            "35": 946619904.0,
+            "36": 946692608.0,
+            "37": 946660352.0,
+            "38": 946685440.0,
+            "39": 946648576.0,
+            "40": 946700800.0,
+            "41": 946663424.0,
+            "42": 946643968.0,
+            "43": 946649088.0,
+            "44": 946657280.0,
+            "45": 946673664.0,
+            "46": 946622976.0,
+            "47": 946645504.0,
+            "48": 946635264.0,
+            "49": 946665984.0,
+            "50": 946649600.0,
+            "51": 946617344.0,
+            "52": 946663936.0,
+            "53": 946621952.0,
+            "54": 946682880.0,
+            "55": 946628096.0,
+            "56": 946677248.0,
+            "57": 946650112.0,
+            "58": 946654208.0,
+            "59": 946659328.0,
+            "60": 946620928.0,
+            "61": 946602496.0,
+            "62": 946637312.0,
+            "63": 946641408.0,
+            "64": 946655232.0,
+            "65": 946661376.0,
+            "66": 946580480.0,
+            "67": 946668032.0,
+            "68": 946612736.0,
+            "69": 946643456.0,
+            "70": 946609664.0,
+            "71": 946628608.0,
+            "72": 946625536.0,
+            "73": 946625536.0,
+            "74": 946591232.0,
+            "75": 946616320.0,
+            "76": 946619392.0,
+            "77": 946655232.0,
+            "78": 946652160.0,
+            "79": 946620928.0,
+            "80": 946565120.0,
+            "81": 946611200.0,
+            "82": 946585600.0,
+            "83": 946577920.0,
+            "84": 946618368.0,
+            "85": 946592768.0,
+            "86": 946585088.0,
+            "87": 946575872.0,
+            "88": 946589696.0,
+            "89": 946577408.0,
+            "90": 946614272.0,
+            "91": 946577920.0,
+            "92": 946578432.0,
+            "93": 946593280.0,
+            "94": 946576896.0,
+            "95": 946573824.0,
+            "96": 946578432.0,
+            "97": 946571776.0,
+            "98": 946619904.0,
+            "99": 946559488.0,
+            "100": 946562560.0
+        }
+    },
+    "mem-max-allocated-bytes": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": 3364709376.0,
+            "2": 3562544128.0,
+            "3": 3562544128.0,
+            "4": 3562544128.0,
+            "5": 3562544128.0,
+            "6": 3562544128.0,
+            "7": 3562544128.0,
+            "8": 3562544128.0,
+            "9": 3562544128.0,
+            "10": 3562544128.0,
+            "11": 3562544128.0,
+            "12": 3562544128.0,
+            "13": 3562544128.0,
+            "14": 3567309824.0,
+            "15": 3567309824.0,
+            "16": 3567309824.0,
+            "17": 3567309824.0,
+            "18": 3567309824.0,
+            "19": 3567309824.0,
+            "20": 3567309824.0,
+            "21": 3567309824.0,
+            "22": 3567309824.0,
+            "23": 3567309824.0,
+            "24": 3567309824.0,
+            "25": 3567309824.0,
+            "26": 3567309824.0,
+            "27": 3567309824.0,
+            "28": 3567309824.0,
+            "29": 3567309824.0,
+            "30": 3567309824.0,
+            "31": 3567309824.0,
+            "32": 3567309824.0,
+            "33": 3567309824.0,
+            "34": 3567309824.0,
+            "35": 3567309824.0,
+            "36": 3567309824.0,
+            "37": 3567309824.0,
+            "38": 3567309824.0,
+            "39": 3567309824.0,
+            "40": 3567309824.0,
+            "41": 3567309824.0,
+            "42": 3567309824.0,
+            "43": 3567309824.0,
+            "44": 3567309824.0,
+            "45": 3567309824.0,
+            "46": 3567309824.0,
+            "47": 3567309824.0,
+            "48": 3567309824.0,
+            "49": 3567309824.0,
+            "50": 3567309824.0,
+            "51": 3567309824.0,
+            "52": 3567309824.0,
+            "53": 3567309824.0,
+            "54": 3567309824.0,
+            "55": 3567309824.0,
+            "56": 3567309824.0,
+            "57": 3567309824.0,
+            "58": 3567309824.0,
+            "59": 3567309824.0,
+            "60": 3567309824.0,
+            "61": 3567309824.0,
+            "62": 3567309824.0,
+            "63": 3567309824.0,
+            "64": 3567309824.0,
+            "65": 3567309824.0,
+            "66": 3567309824.0,
+            "67": 3567309824.0,
+            "68": 3567309824.0,
+            "69": 3567309824.0,
+            "70": 3567309824.0,
+            "71": 3567309824.0,
+            "72": 3567309824.0,
+            "73": 3567309824.0,
+            "74": 3567309824.0,
+            "75": 3567309824.0,
+            "76": 3567309824.0,
+            "77": 3567309824.0,
+            "78": 3567309824.0,
+            "79": 3567309824.0,
+            "80": 3567309824.0,
+            "81": 3567309824.0,
+            "82": 3567309824.0,
+            "83": 3567309824.0,
+            "84": 3567309824.0,
+            "85": 3567309824.0,
+            "86": 3567309824.0,
+            "87": 3567309824.0,
+            "88": 3567309824.0,
+            "89": 3567309824.0,
+            "90": 3567309824.0,
+            "91": 3567309824.0,
+            "92": 3567309824.0,
+            "93": 3567309824.0,
+            "94": 3567309824.0,
+            "95": 3567309824.0,
+            "96": 3567309824.0,
+            "97": 3567309824.0,
+            "98": 3567309824.0,
+            "99": 3567309824.0,
+            "100": 3567309824.0
+        }
+    },
+    "iteration-time": {
+        "start_step": 1,
+        "end_step": 100,
+        "step_interval": 1,
+        "values": {
+            "1": "nan",
+            "2": 6.05801,
+            "3": 0.39062,
+            "4": 0.87277,
+            "5": 0.51828,
+            "6": 0.69473,
+            "7": 0.71957,
+            "8": 1.27965,
+            "9": 0.49964,
+            "10": 0.38429,
+            "11": 1.30501,
+            "12": 0.35422,
+            "13": 0.54857,
+            "14": 1.02604,
+            "15": 0.37296,
+            "16": 0.95544,
+            "17": 0.61093,
+            "18": 1.05287,
+            "19": 0.34997,
+            "20": 0.53359,
+            "21": 0.91265,
+            "22": 0.72092,
+            "23": 0.75137,
+            "24": 0.35209,
+            "25": 0.45237,
+            "26": 1.03503,
+            "27": 0.34594,
+            "28": 0.49593,
+            "29": 0.72384,
+            "30": 0.34985,
+            "31": 0.34442,
+            "32": 1.1492,
+            "33": 0.66678,
+            "34": 0.75815,
+            "35": 0.3493,
+            "36": 0.67982,
+            "37": 0.67449,
+            "38": 0.4266,
+            "39": 0.5845,
+            "40": 0.60572,
+            "41": 0.61727,
+            "42": 0.80642,
+            "43": 0.72708,
+            "44": 0.98036,
+            "45": 0.34515,
+            "46": 0.44476,
+            "47": 0.64953,
+            "48": 0.8351,
+            "49": 0.75792,
+            "50": 0.67539,
+            "51": 0.42011,
+            "52": 0.39133,
+            "53": 0.40794,
+            "54": 1.09969,
+            "55": 0.78503,
+            "56": 0.90347,
+            "57": 0.45729,
+            "58": 0.39823,
+            "59": 0.83852,
+            "60": 0.75307,
+            "61": 0.59767,
+            "62": 0.66194,
+            "63": 0.58181,
+            "64": 1.01678,
+            "65": 0.54766,
+            "66": 1.2761,
+            "67": 0.82187,
+            "68": 0.34682,
+            "69": 1.60066,
+            "70": 1.62769,
+            "71": 0.34481,
+            "72": 0.69576,
+            "73": 0.81927,
+            "74": 0.45901,
+            "75": 1.0719,
+            "76": 0.47771,
+            "77": 0.46301,
+            "78": 1.44906,
+            "79": 0.59493,
+            "80": 0.57479,
+            "81": 1.1542,
+            "82": 0.73537,
+            "83": 0.49291,
+            "84": 0.67403,
+            "85": 0.62403,
+            "86": 0.71952,
+            "87": 0.94935,
+            "88": 0.52739,
+            "89": 0.5562,
+            "90": 0.51513,
+            "91": 0.54768,
+            "92": 1.18744,
+            "93": 0.48268,
+            "94": 0.62455,
+            "95": 0.69947,
+            "96": 0.43952,
+            "97": 0.50572,
+            "98": 0.76873,
+            "99": 1.25479,
+            "100": 0.51265
+        }
+    }
+}
\ No newline at end of file
diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/model_config.yaml
new file mode 100644
index 00000000000..40ac94eed9b
--- /dev/null
+++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/model_config.yaml
@@ -0,0 +1,69 @@
+ENV_VARS:
+  CUDA_DEVICE_MAX_CONNECTIONS: 1
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0
+  NCCL_ALGO: Ring
+  CUBLAS_WORKSPACE_CONFIG: :4096:8
+MODEL_ARGS:
+  --num-layers: 12
+  --hidden-size: 512
+  --num-attention-heads: 8
+  --log-params-norm: true
+  --log-num-zeros-in-grad: true
+  --log-validation-ppl-to-tensorboard: true
+  --log-timers-to-tensorboard: true
+  --tensorboard-dir: ${TENSORBOARD_PATH}
+  --micro-batch-size: 4
+  --global-batch-size: 32
+  --seq-length: 1024
+  --max-position-embeddings: 1024
+  --disable-bias-linear: true
+  --train-iters: 100
+  --timing-log-level: 0
+  --lr-decay-iters: 320000
+  --save: ${CHECKPOINT_SAVE_PATH}
+  --load: ${CHECKPOINT_LOAD_PATH}
+  --data-path: ${DATA_PATH}/text/common_pile/v01_filtered_data/my-gpt3_00_text_document
+  --vocab-file: ${DATA_PATH}/text/common_pile/v01_filtered_data/bpe/vocab.json
+  --merge-file: ${DATA_PATH}/text/common_pile/v01_filtered_data/bpe/merges.txt
+  --split: 949,50,1
+  --distributed-backend: nccl
+  --lr: 0.00015
+  --lr-decay-style: cosine
+  --min-lr: 1.0e-5
+  --weight-decay: 1e-2
+  --clip-grad: 1.0
+  --lr-warmup-fraction: .01
+  --log-interval: 1
+  --save-interval: 50
+  --eval-interval: 1000
+  --eval-iters: 10
+  --transformer-impl: transformer_engine
+  --tensor-model-parallel-size: 1
+  --pipeline-model-parallel-size: 1
+  --expert-model-parallel-size: 4
+  --num-experts: 8
+  --moe-token-dispatcher-type: allgather
+  --moe-router-load-balancing-type: aux_loss
+  --moe-router-topk: 2
+  --moe-router-dtype: fp32
+  --moe-ffn-hidden-size: 1024
+  --moe-grouped-gemm: true
+  --ckpt-fully-parallel-load: true
+  --deterministic-mode: true
+  --no-gradient-accumulation-fusion: true
+  --attention-softmax-in-fp32: true
+  --use-checkpoint-opt_param-scheduler: true
+  --use-mcore-models: true
+  --ckpt-format: torch_dist
+  --data-cache-path: ${DATA_CACHE_PATH}
+  --bf16: true
+  --no-bias-gelu-fusion: true
+  --log-memory-to-tensorboard: true
+  --optimizer: muon
+  --muon-momentum: 0.9
+  --muon-extra-scale-factor: 0.2
+  --muon-scale-mode: spectral
+  --async-save: true
+  --use-persistent-ckpt-worker: true
+  --use-distributed-optimizer: true
+TEST_TYPE: ckpt-resume
diff --git a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py
index 3f60658a005..42ef0a401ee 100644
--- a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py
+++ b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py
@@ -603,4 +603,75 @@ def test_optimizer_common_state_dict(
                 # Test both param state dicts are equal
                 check_equal(optim_param_state_A, optim_param_state_B)
 
+    @pytest.mark.parametrize('tp', [1, 2])
+    @pytest.mark.parametrize('pp', [1, 2])
+    def test_optimizer_common_state_dict_hybrid(self, tmp_path_dist_ckpt, tp, pp):
+        """End-to-end ``save_checkpoint``/``load_checkpoint`` roundtrip on the
+        hybrid LayerWise + DistributedOptimizer path.
+
+        Muon matrix params live in :class:`LayerWiseDistributedOptimizer`
+        while non-Muon params (embeddings, biases, layernorm) go through a
+        real :class:`DistributedOptimizer` sub-optimizer. Catches the case
+        where ``_build_sharded_state_dict_metadata`` skipped populating
+        ``distrib_optim_sharding_type`` because the arg parser flips
+        ``use_distributed_optimizer`` off in Muon mode -- the DistOpt
+        sub-optimizer then defaulted to the deprecated
+        ``fully_sharded_model_space`` save path which is incompatible with
+        the post-5ab481cb45 ShardedTensor validation.
+        """
+        if tp * pp > 8:
+            pytest.skip("TP*PP > 8 is larger than world size")
+
+        Utils.initialize_model_parallel(tp, pp)
+
+        with TempNamedDir(
+            tmp_path_dist_ckpt / 'test_optimizer_common_state_dict_hybrid', sync=True
+        ) as ckpt_dir:
+            mock_args = parse_args(ignore_unknown_args=True)
+            # Mirror the arg-parser's Muon path: ``use_distributed_optimizer``
+            # is flipped off and ``use_layer_wise_distributed_optimizer`` is
+            # the surviving flag.
+            mock_args.use_distributed_optimizer = False
+            mock_args.use_layer_wise_distributed_optimizer = True
+            with mock.patch('megatron.training.checkpointing.get_args', new=lambda: mock_args):
+                model, optimizer_A = setup_model_and_optimizer(
+                    seed=2,
+                    tp=tp,
+                    pp=pp,
+                    initialize_fn=initialize_gpt_model,
+                    dist_opt=True,
+                    optimizer='dist_muon',
+                    use_param_layout=True,
+                )
+
+                init_checkpointing_mock_args(mock_args, ckpt_dir, fully_parallel=True)
+                from megatron.training.training import preprocess_common_state_dict
+
+                save_checkpoint(
+                    10,
+                    model,
+                    optimizer_A,
+                    None,
+                    0,
+                    preprocess_common_state_dict_fn=preprocess_common_state_dict,
+                )
+
+                optim_param_state_A = optimizer_A.state_dict()
+
+                model, optimizer_B = setup_model_and_optimizer(
+                    seed=3,
+                    tp=tp,
+                    pp=pp,
+                    initialize_fn=initialize_gpt_model,
+                    dist_opt=True,
+                    optimizer='dist_muon',
+                    use_param_layout=True,
+                )
+
+                load_checkpoint_no_arg_checks(model, optimizer_B, None)
+
+                optim_param_state_B = optimizer_B.state_dict()
+
+                check_equal(optim_param_state_A, optim_param_state_B)
+
         Utils.destroy_model_parallel()
diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py
index f5acc373a2f..1a7b1553449 100644
--- a/tests/unit_tests/dist_checkpointing/utils.py
+++ b/tests/unit_tests/dist_checkpointing/utils.py
@@ -235,12 +235,25 @@ def setup_model_and_optimizer(
     torch.manual_seed(seed + 1)
     model_parallel_cuda_manual_seed(seed + 1)
 
+    def _init_states(optimizer):
+        # In hybrid LayerWise + DistOpt mode the top-level ChainedOptimizer
+        # wraps another ChainedOptimizer (LayerWise) alongside DistOpt; recurse
+        # so the Muon Float16 sub-optimizers inside LayerWise still get their
+        # state seeded. Optimizers without ``init_state_fn`` (DistOpt) seed
+        # their state elsewhere and are skipped here.
+        if isinstance(optimizer, ChainedOptimizer):
+            for child_optimizer in optimizer.chained_optimizers:
+                _init_states(child_optimizer)
+            return
+        if not hasattr(optimizer, 'init_state_fn'):
+            return
+        if not hasattr(optimizer, 'optimizer'):
+            optimizer.init_state_fn(optimizer)
+        else:
+            optimizer.init_state_fn(optimizer.optimizer)
+
     if isinstance(optimizer, ChainedOptimizer):
-        for opt in optimizer.chained_optimizers:
-            if not hasattr(opt, 'optimizer'):
-                opt.init_state_fn(opt)
-            else:
-                opt.init_state_fn(opt.optimizer)
+        _init_states(optimizer)
     else:
         for group in optimizer.optimizer.param_groups:
             for p in group['params']:
diff --git a/tests/unit_tests/test_checkpointing.py b/tests/unit_tests/test_checkpointing.py
index c778d58fb7f..9b62c26d674 100644
--- a/tests/unit_tests/test_checkpointing.py
+++ b/tests/unit_tests/test_checkpointing.py
@@ -382,3 +382,71 @@ def test_read_metadata_non_distributed(tmp_path, metadata_content, expected_iter
 
     assert max_iter == expected_iter, f"Expected iteration {expected_iter}, got {max_iter}"
     assert release == expected_release, f"Expected release={expected_release}, got {release}"
+
+
+def _make_metadata_args(
+    use_distributed_optimizer=False,
+    use_layer_wise_distributed_optimizer=False,
+    ckpt_format='torch_dist',
+    dist_ckpt_optim_fully_reshardable=False,
+    distrib_optim_fully_reshardable_mem_efficient=False,
+):
+    args = SimpleNamespace()
+    args.use_distributed_optimizer = use_distributed_optimizer
+    args.use_layer_wise_distributed_optimizer = use_layer_wise_distributed_optimizer
+    args.ckpt_format = ckpt_format
+    args.dist_ckpt_optim_fully_reshardable = dist_ckpt_optim_fully_reshardable
+    args.distrib_optim_fully_reshardable_mem_efficient = (
+        distrib_optim_fully_reshardable_mem_efficient
+    )
+    return args
+
+
+class TestBuildShardedStateDictMetadata:
+    """``_build_sharded_state_dict_metadata`` must set ``distrib_optim_sharding_type``
+    whenever a real :class:`DistributedOptimizer` instance will be used at save
+    time -- otherwise the DistOpt path falls through to the deprecated
+    ``fully_sharded_model_space`` default whose ``flattened_range`` usage is
+    rejected by ``ShardedTensor.validate_metadata_integrity`` post commit
+    5ab481cb45.
+    """
+
+    DUMMY_GROUP = object()
+
+    def test_distributed_optimizer_sets_dp_reshardable_default(self):
+        args = _make_metadata_args(use_distributed_optimizer=True)
+        metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=self.DUMMY_GROUP)
+        assert metadata['distrib_optim_sharding_type'] == 'dp_reshardable'
+
+    def test_distributed_optimizer_fully_reshardable_flag(self):
+        args = _make_metadata_args(
+            use_distributed_optimizer=True, dist_ckpt_optim_fully_reshardable=True
+        )
+        metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=self.DUMMY_GROUP)
+        assert metadata['distrib_optim_sharding_type'] == 'fully_reshardable'
+        assert metadata['distrib_optim_fully_reshardable_mem_efficient'] is False
+
+    def test_distributed_optimizer_fsdp_dtensor(self):
+        args = _make_metadata_args(use_distributed_optimizer=True, ckpt_format='fsdp_dtensor')
+        metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=self.DUMMY_GROUP)
+        assert metadata['distrib_optim_sharding_type'] == 'fsdp_dtensor'
+
+    def test_layer_wise_only_still_sets_sharding_type(self):
+        # Arg parser flips ``use_distributed_optimizer`` off when Muon is in
+        # use, but the LayerWise + DistOpt split path still has a DistOpt
+        # sub-optimizer for non-Muon params, so the metadata is required.
+        args = _make_metadata_args(use_layer_wise_distributed_optimizer=True)
+        metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=self.DUMMY_GROUP)
+        assert metadata['distrib_optim_sharding_type'] == 'dp_reshardable'
+
+    def test_layer_wise_with_fully_reshardable(self):
+        args = _make_metadata_args(
+            use_layer_wise_distributed_optimizer=True, dist_ckpt_optim_fully_reshardable=True
+        )
+        metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=self.DUMMY_GROUP)
+        assert metadata['distrib_optim_sharding_type'] == 'fully_reshardable'
+
+    def test_no_distributed_optimizer_no_sharding_type(self):
+        args = _make_metadata_args()
+        metadata = _build_sharded_state_dict_metadata(args, dp_cp_group=self.DUMMY_GROUP)
+        assert 'distrib_optim_sharding_type' not in metadata
diff --git a/tests/unit_tests/test_layer_wise_optimizer.py b/tests/unit_tests/test_layer_wise_optimizer.py
index 572bdb83c6f..aaa06b2ddbe 100644
--- a/tests/unit_tests/test_layer_wise_optimizer.py
+++ b/tests/unit_tests/test_layer_wise_optimizer.py
@@ -141,7 +141,9 @@ def create_model_and_optimizer(
         pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True)
         pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group()
 
-        optimizer = get_megatron_optimizer(optimizer_config, [model], pg_collection=pg_collection)
+        optimizer = get_megatron_optimizer(
+            optimizer_config, [model], pg_collection=pg_collection, use_gloo_process_groups=False
+        )
         return model, optimizer, pg_collection
 
     def create_model_and_optimizer_with_overlap_param_gather(
@@ -237,7 +239,7 @@ def create_model_and_optimizer_with_overlap_param_gather(
         optimizer = get_megatron_optimizer(
             config=optimizer_config,
             model_chunks=[model],
-            use_gloo_process_groups=True,
+            use_gloo_process_groups=False,
             pg_collection=pg_collection,
         )
         return model, optimizer, pg_collection
@@ -588,7 +590,18 @@ def test_overlap_param_gather_basic(self, use_param_layout):
         )
 
         assert optimizer is not None, "Optimizer should not be None"
-        assert optimizer.overlap_param_gather, "overlap_param_gather should be True"
+        # ``optimizer`` may be a ChainedOptimizer wrapping the LayerWise +
+        # DistOpt pair when non-Muon params are routed to DistOpt. Find the
+        # LayerWise instance to check its ``overlap_param_gather`` flag.
+        layer_wise_optimizer = next(
+            (
+                sub_optimizer
+                for sub_optimizer in getattr(optimizer, 'chained_optimizers', [optimizer])
+                if isinstance(sub_optimizer, LayerWiseDistributedOptimizer)
+            ),
+            optimizer,
+        )
+        assert layer_wise_optimizer.overlap_param_gather, "overlap_param_gather should be True"
 
         reference_model = self.create_reference_model(model)