[dev] Add experimental decoupled compact LayerWise DDP layout for Muon

megatron/core/distributed/distributed_data_parallel_config.py

@@ -36,6 +36,12 @@ class DistributedDataParallelConfig:
        enabled. Defaults to 1, which means DistOpt is across entire DP domain.
     """
 
+    use_layer_wise_param_layout: bool = False
+    """Layer-wise (Muon) optimizer only. When True, LayerWise-managed buffers use
+       the shard-aligned padded LayerWise param layout. When False (default), the compact
+       decoupled layout is selected instead.
+    """
+
     check_for_nan_in_grad: bool = False
     """
     If true, check for NaNs and Infs in gradients _before_ communication collective.

megatron/core/distributed/param_and_grad_buffer.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
 
+import dataclasses
 import fnmatch
 import functools
 import logging
@@ -1006,6 +1007,17 @@ def __init__(
         self.gradient_scaling_factor = gradient_scaling_factor
         self.nccl_ub = nccl_ub
 
+        self._is_layer_wise_buffer = bool(
+            self.params and getattr(self.params[0], "is_managed_by_layer_wise_optimizer", False)
+        )
+        # Bake the per-buffer DistOpt decision into this buffer's ddp_config (single source of
+        # truth; bucket groups inherit it): a LayerWise (Muon) buffer on the compact decoupled
+        # layout disables DistributedOptimizer, while sibling buffers keep the model-level setting.
+        if self._is_layer_wise_buffer and not getattr(
+            self.ddp_config, "use_layer_wise_param_layout", True
+        ):
+            self.ddp_config = dataclasses.replace(self.ddp_config, use_distributed_optimizer=False)
+
         # Data structures to store underlying buckets and relevant indexing data.
         self.buckets = []
         self.param_to_bucket = {}  # Param -> bucket mapping.
@@ -1053,6 +1065,27 @@ def __init__(
             # nvfp4_packed_numel_unpadded is already set by _compute_nvfp4_packed_layout.
 
         assert self.numel_unpadded <= self.numel
+
+        # Diagnostic: log persistent buffer size vs. unpadded payload so the cost of any
+        # optimizer-driven padding (e.g. the LayerWise shard-aligned ``dp_size * max(shard_load)``
+        # layout) is visible per buffer. Emit at INFO only when it is interesting — a
+        # LayerWise-managed buffer or one that actually carries padding — and DEBUG otherwise, so
+        # ordinary (zero-padding) buffers do not spam non-experimental runs.
+        _padding = self.numel - self.numel_unpadded
+        _pad_frac = _padding / max(self.numel_unpadded, 1)
+        log_on_each_pipeline_stage(
+            logger,
+            logging.INFO if (self._is_layer_wise_buffer or _padding > 0) else logging.DEBUG,
+            f"ParamAndGradBuffer layout: param_dtype={self.param_dtype} "
+            f"grad_dtype={self.grad_dtype} dp_world_size={self.data_parallel_world_size} "
+            f"layerwise={self._is_layer_wise_buffer} "
+            f"distopt={self.ddp_config.use_distributed_optimizer} "
+            f"numel={self.numel} numel_unpadded={self.numel_unpadded} "
+            f"padding={_padding} ({_pad_frac:.1%})",
+            tp_group=self.tp_group,
+            dp_cp_group=self.dp_cp_group,
+        )
+
         if self.has_nvfp4_params:
             assert self.nvfp4_packed_numel_unpadded <= self.nvfp4_packed_numel
         if self.ddp_config.use_distributed_optimizer:
@@ -1565,22 +1598,72 @@ def partition_buckets(
             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.
+    # A bucket group performs a single collective type (reduce-scatter for DistOpt buffers,
+    # all-reduce otherwise), so buckets merged into one group must agree on the effective
+    # per-buffer ``use_distributed_optimizer``. The decoupled LayerWise layout
+    # (``use_layer_wise_param_layout=False``) gives LayerWise (Muon) buffers
+    # ``use_distributed_optimizer=False`` while sibling buffers keep True; the no-fp8 Case 2 below
+    # keeps every bucket in its own group so they never mix, but the merging Cases 1/3 must assert
+    # consistency.
+    _ddp_config = buffers[0].ddp_config
+    _decouple = not getattr(_ddp_config, "use_layer_wise_param_layout", True)
+
+    def _bucket_distopt(bucket):
+        """This bucket's effective ``use_distributed_optimizer``."""
+        is_lw = bool(
+            bucket.params_list
+            and getattr(bucket.params_list[0], "is_managed_by_layer_wise_optimizer", False)
+        )
+        if _decouple and is_lw:
+            return False
+        return _ddp_config.use_distributed_optimizer
+
+    def _merged_use_distributed_optimizer(merge_buckets):
+        values = {_bucket_distopt(bucket) for bucket in merge_buckets}
+        assert len(values) == 1, (
+            "Cannot merge buckets with differing effective use_distributed_optimizer into one "
+            "bucket group. This happens when the decoupled LayerWise layout "
+            "(use_layer_wise_param_layout=False) mixes LayerWise (all-reduce) and non-LayerWise "
+            "(reduce-scatter) buffers under a merging bucketing strategy (e.g. the fp8 merge "
+            "path). Disable bucket merging for the decoupled LayerWise path."
+        )
+        return values.pop()
+
+    # Case 1: Put all buckets into a single bucket group if force_single_bucket_group is True
+    # (e.g. disable_bucketing / non-first VPP chunks). A bucket group performs a single
+    # collective type, so when the decoupled LayerWise layout (use_layer_wise_param_layout=False)
+    # mixes LayerWise (all-reduce, non-DistOpt) and non-LayerWise (reduce-scatter, DistOpt)
+    # buffers in one chunk, we cannot
+    # merge them into a single group. Split by the effective per-bucket use_distributed_optimizer
+    # instead, preserving order. When all buckets agree (the non-decoupled case) this collapses
+    # to exactly one group, identical to the previous behavior.
     if force_single_bucket_group:
-        buckets = []
-        ddp_config = buffers[0].ddp_config
         data_parallel_group = buffers[0].data_parallel_group
         data_parallel_world_size = buffers[0].data_parallel_world_size
+        ordered_distopt_values = []
+        buckets_by_distopt = {}
+        # buffer.ddp_config already carries the per-buffer use_distributed_optimizer.
+        ddp_config_by_distopt = {}
         for buffer in buffers:
-            assert ddp_config == buffer.ddp_config
             assert data_parallel_group == buffer.data_parallel_group
             assert data_parallel_world_size == buffer.data_parallel_world_size
-            buckets.extend(buffer.buckets)
-
-        bucket_group = _ParamAndGradBucketGroup(
-            buckets, ddp_config, data_parallel_group, data_parallel_world_size
-        )
-        return [bucket_group]
+            distopt = buffer.ddp_config.use_distributed_optimizer
+            ddp_config_by_distopt.setdefault(distopt, buffer.ddp_config)
+            for bucket in buffer.buckets:
+                if distopt not in buckets_by_distopt:
+                    buckets_by_distopt[distopt] = []
+                    ordered_distopt_values.append(distopt)
+                buckets_by_distopt[distopt].append(bucket)
+
+        return [
+            _ParamAndGradBucketGroup(
+                buckets_by_distopt[distopt],
+                ddp_config_by_distopt[distopt],
+                data_parallel_group,
+                data_parallel_world_size,
+            )
+            for distopt in ordered_distopt_values
+        ]
 
     if fp8_buffer is None:
         # Case 2: When there is no fp8 buffer in the input buffers, let each bucket group have
@@ -1599,11 +1682,12 @@ def partition_buckets(
         return bucket_groups
     else:
         # Case 3: When using fp8 params, merge all non-fp8 buckets into the last fp8 bucket group.
-        non_fp8_buckets = []
+        # Track each non-fp8 bucket with its buffer's (authoritative) ddp_config.
+        non_fp8_buckets = []  # list of (bucket, ddp_config)
         for buffer in buffers:
             if buffer.param_dtype != torch.uint8:
                 for bucket in buffer.buckets:
-                    non_fp8_buckets.append(bucket)
+                    non_fp8_buckets.append((bucket, buffer.ddp_config))
 
         bucket_groups = []
         for bucket in fp8_buffer.buckets:
@@ -1617,32 +1701,37 @@ def partition_buckets(
                     bucket_groups.append(
                         _ParamAndGradBucketGroup(
                             [bucket],
-                            buffer.ddp_config,
+                            fp8_buffer.ddp_config,
                             buffer.data_parallel_group,
                             buffer.data_parallel_world_size,
                         )
                     )
                     if non_fp8_buckets:
-                        for non_fp8_bucket in non_fp8_buckets:
+                        for non_fp8_bucket, non_fp8_ddp_config in non_fp8_buckets:
                             bucket_groups.append(
                                 _ParamAndGradBucketGroup(
                                     [non_fp8_bucket],
-                                    buffer.ddp_config,
+                                    non_fp8_ddp_config,
                                     buffer.data_parallel_group,
                                     buffer.data_parallel_world_size,
                                 )
                             )
 
                     continue  # Skip the default bucket group creation below
                 else:
-                    group_buckets = [bucket] + non_fp8_buckets
+                    group_buckets = [bucket] + [b for b, _ in non_fp8_buckets]
             else:
                 # The first N-1 bucket groups.
                 group_buckets = [bucket]
+            # Merged buckets must share the fp8 group's effective use_distributed_optimizer.
+            assert (
+                _merged_use_distributed_optimizer(group_buckets)
+                == fp8_buffer.ddp_config.use_distributed_optimizer
+            )
             bucket_groups.append(
                 _ParamAndGradBucketGroup(
                     group_buckets,
-                    buffer.ddp_config,
+                    fp8_buffer.ddp_config,
                     buffer.data_parallel_group,
                     buffer.data_parallel_world_size,
                 )

megatron/core/optimizer/layer_wise_optimizer.py

@@ -317,8 +317,15 @@ def compute_full_param_layout(
             :class:`FullParamLayout` with a :class:`PerBufferParamLayout` per buffer group.
         """
         # Avoid a circular import: DistributedOptimizer imports LayerWise indirectly.
+        from ..distributed.param_and_grad_buffer import _compute_default_per_buffer_param_layout
         from .distrib_optimizer import DistributedOptimizer
 
+        # Decoupled layout (use_layer_wise_param_layout=False): LayerWise (Muon) buffers use a
+        # compact no-padding DDP layout (and locally disable DistributedOptimizer semantics in
+        # DDP), so they must NOT receive the shard-aligned ``dp_size * max(shard_load)`` padded
+        # layout here. Non-LayerWise buffers keep DistOpt's byte-level layout regardless.
+        decouple_ddp_layout = not ddp_config.use_layer_wise_param_layout
+
         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():
@@ -334,6 +341,15 @@ def compute_full_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 and decouple_ddp_layout:
+                # Compact no-padding layout (DDP treats this buffer as non-DistOpt). Attach
+                # param_indices so DDP's layout/grouping consistency check passes.
+                per_buffer_layout = _compute_default_per_buffer_param_layout(
+                    group_params, bucket_size
+                )
+                per_buffer_layout.param_indices = param_indices
+                layouts[buffer_key] = per_buffer_layout
+                continue
             if buffer_key.is_managed_by_layer_wise_optimizer:
                 compute_per_buffer_layout = (
                     LayerWiseDistributedOptimizer._compute_per_buffer_param_layout
@@ -365,22 +381,24 @@ def __init__(
         """
 
         self.pg_collection = pg_collection
+        self.decouple_ddp_layout = not config.use_layer_wise_param_layout
 
         full_param_layouts = None
-        if model_chunks is not None:
+        if model_chunks is not None and not self.decouple_ddp_layout:
             full_param_layouts = [
                 chunk.full_param_layout
                 for chunk in model_chunks
                 if hasattr(chunk, 'full_param_layout') and chunk.full_param_layout is not None
             ] or None
         self.shard_params(optimizers, full_param_layouts)
 
-        # When a full_param_layout is available, ddp_config.use_distributed_optimizer
+        # When a full_param_layout is available (no decoupling), use_distributed_optimizer
         # is True and model params are views into the DDP param buffer.  After the
         # optimizer step copies updated fp32 main params → bf16 model params, the
         # buffer is already up-to-date in-place.  We can use DDP's buffer-based
         # all-gather (start_param_sync) instead of the flatten/unflatten allgather_params
-        # path.
+        # In the decouple path, Muon buffers are non-DistOpt and own
+        # whole params via ping-pong, so we use the legacy allgather_params path instead.
         self.use_buffer_param_sync = full_param_layouts is not None
 
         # Set up overlap param gather using DDP bucket infrastructure.

megatron/core/optimizer/optimizer_config.py

@@ -323,6 +323,13 @@ class OptimizerConfig:
     arguments layer sets this flag and resets ``use_distributed_optimizer`` to False so
     that the standard distributed-optimizer path is not triggered."""
 
+    use_layer_wise_param_layout: bool = False
+    """Layer-wise (Muon) optimizer only. When True, LayerWise-managed buffers use
+    the shard-aligned padded LayerWise param layout. When False (default), the compact decoupled
+    layout is selected: LayerWise-managed (Muon) buffers use a compact no-padding DDP layout with
+    all-reduce gradients and legacy whole-param ping-pong ownership + ``allgather_params``
+    param sync."""
+
     overlap_param_gather: bool = False
     """If true, overlap param all-gather with forward compute. 
         This argument is intended to have the same value as the "overlap_param_gather" argument 

megatron/training/arguments.py

@@ -1893,6 +1893,21 @@ def validate_args(args, defaults={}):
             "torch_dist",
         ], "Emerging optimizer supports torch and torch_dist checkpoint format."
 
+    if args.use_layer_wise_distributed_optimizer:
+        assert not args.fp8_param_gather and not getattr(args, 'fp4_param_gather', False), (
+            "Layer-wise (Muon) distributed optimizer does not support FP8/FP4 parameter gather "
+            "(fp8_param_gather / fp4_param_gather). Use fp8_param_gather=False (e.g. blockwise/"
+            "MXFP8 compute with parameters persisted in bf16)."
+        )
+        if not args.use_layer_wise_param_layout:
+            assert args.num_distributed_optimizer_instances == 1, (
+                "the decoupled compact LayerWise DDP layout (the default; pass "
+                "--use-layer-wise-param-layout for the padded layout) requires "
+                "num_distributed_optimizer_instances == 1: the non-DistOpt LayerWise (Muon) buffers "
+                "only all-reduce within a single optimizer instance, so partial DistOpt (>1 "
+                "instance) would under-reduce Muon gradients across the full data-parallel domain."
+            )
+
     # Make sure all functionality that requires Gloo process groups is disabled.
     if not args.use_gloo_process_groups:
         if args.use_distributed_optimizer:
@@ -4204,17 +4219,16 @@ def _add_distributed_args(parser):
                        Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.',
     )
     group.add_argument(
-        '--no-use-layer-wise-param-layout',
-        action='store_false',
+        '--use-layer-wise-param-layout',
+        action='store_true',
         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.',
+        default=False,
+        help='Opt INTO the padded shard-aligned LayerWise param layout. The default is the compact '
+        'decoupled layout, where LayerWise (Muon 2D) buffers use a no-padding DDP layout and locally '
+        'disable DistributedOptimizer (all-reduce grads + whole-param ping-pong + allgather_params), '
+        'while sibling buffers keep the byte-level DistributedOptimizer; this avoids the persistent '
+        'dp_size * max(shard_load) padding. Pass this flag to restore the padded layout (e.g. for '
+        'bit-for-bit comparison; it uses a different bf16 reduction ordering).',
     )
     return parser