Fix layer-wise distributed optimizer MXFP8 wgrad failure by padding per-param start

megatron/core/distributed/distributed_data_parallel.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import logging
 from contextlib import contextmanager
@@ -38,6 +38,11 @@ class DistributedDataParallel(_BaseDataParallel):
         full_param_layout: Optional FullParamLayout providing pre-computed layouts for all
             dtype groups. When provided, each buffer uses the corresponding PerBufferParamLayout
             instead of computing a default one.
+        pad_param_starts: If True (and ``full_param_layout`` is not provided), each buffer's
+            default layout rounds per-param start indices up to a 64-element boundary. The
+            layer-wise distributed optimizer relies on this for cuBLAS MXFP8 wgrad D-pointer
+            alignment; standard distributed optimizer doesn't need it here because it supplies
+            a pre-padded ``full_param_layout`` instead.
 
     """
 
@@ -49,6 +54,7 @@ def __init__(
         disable_bucketing: bool = False,
         pg_collection: Optional[ProcessGroupCollection] = None,
         full_param_layout: Optional[FullParamLayout] = None,
+        pad_param_starts: bool = False,
     ):
         super().__init__(config=config, module=module)
         if has_config_logger_enabled(config):
@@ -250,6 +256,7 @@ def __init__(
                 self.ddp_config.nccl_ub,
                 pg_collection,
                 param_layout=param_layout,
+                pad_param_starts=pad_param_starts,
             )
             if buffer_key.is_expert_parallel:
                 self.expert_parallel_buffers.append(buffer)

megatron/core/distributed/param_and_grad_buffer.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import fnmatch
 import functools
@@ -847,21 +847,33 @@ def group_params_for_buffers(
 
 
 def _compute_default_per_buffer_param_layout(
-    params: List[torch.nn.Parameter], bucket_size: Optional[int]
+    params: List[torch.nn.Parameter], bucket_size: Optional[int], pad_param_starts: bool = False
 ) -> 'PerBufferParamLayout':
     """Compute parameter layout for the non-distributed-optimizer case.
 
-    No padding is applied. Parameters are iterated in reverse order (backprop order)
-    and grouped into buckets of approximately `bucket_size` elements.
+    By default no padding is applied. When ``pad_param_starts`` is True, each parameter's
+    start index is rounded up to a 64-element boundary, mirroring what
+    ``DistributedOptimizer._compute_per_buffer_param_layout`` does. This is needed by the
+    layer-wise distributed optimizer (Muon and other non-Adam/SGD optimizers): without
+    per-param alignment, mid-bucket params' ``main_grad`` slices land at low (16-byte)
+    D-pointer alignment in the MXFP8 wgrad path on cuBLASLt 12.8.x, and
+    ``cublasLtMatmulAlgoGetHeuristic`` returns ``CUBLAS_STATUS_NOT_SUPPORTED``.
+    Memory cost: <= 63 * 2 bytes = 126 bytes per param (trivial).
+
+    Parameters are iterated in reverse order (backprop order) and grouped into buckets of
+    approximately ``bucket_size`` elements.
 
     Args:
         params: List of parameters to lay out.
         bucket_size: Approximate number of elements per bucket, or None for a single bucket.
+        pad_param_starts: If True, round each param's start index up to a 64-element
+            boundary. Caller is responsible for deciding when this is needed (e.g.,
+            layer-wise distributed optimizer with MXFP8 wgrad).
 
     Returns:
         PerBufferParamLayout with the computed mapping.
     """
-    from ..optimizer.param_layout import PerBufferParamLayout
+    from ..optimizer.param_layout import PerBufferParamLayout, pad_param_start
 
     param_index_map = {}
     bucket_indices = []
@@ -873,6 +885,8 @@ def _compute_default_per_buffer_param_layout(
     bucket_id = 0
 
     for param in params[::-1]:
+        if pad_param_starts:
+            param_start_index = pad_param_start(param_start_index)
         this_numel = param.data.nelement()
         param_end_index = param_start_index + this_numel
         param_index_map[param] = (param_start_index, param_end_index, bucket_id)
@@ -917,6 +931,10 @@ class _ParamAndGradBuffer:
         param_indices: The index of each param among the params with same dtype, if a param is fp8,
             use its "fake" high precision dtype to determine which params have same dtype with it.
             These indices are needed when loading a non-native-fp8 checkpoint in native-fp8 mode.
+        pad_param_starts: If True and ``param_layout`` is not supplied, the default layout
+            rounds each param's start index up to a 64-element boundary (needed for cuBLAS
+            MXFP8 wgrad D-pointer alignment under the layer-wise distributed optimizer).
+            Ignored when ``param_layout`` is provided (the supplied layout decides padding).
     """
 
     def __init__(
@@ -933,6 +951,7 @@ def __init__(
         nccl_ub: bool,
         pg_collection: Optional[ProcessGroupCollection] = None,
         param_layout: Optional['PerBufferParamLayout'] = None,
+        pad_param_starts: bool = False,
     ):
 
         if pg_collection is None:
@@ -968,9 +987,14 @@ def __init__(
         self.buckets = []
         self.param_to_bucket = {}  # Param -> bucket mapping.
 
-        # Use the provided layout if given, otherwise compute the default (no-padding) layout.
+        # Use the provided layout if given, otherwise compute the default layout. The default
+        # layout applies per-param 64-element start padding only when pad_param_starts is set
+        # (needed for cuBLAS MXFP8 wgrad D-pointer alignment under the layer-wise distributed
+        # optimizer); otherwise no padding.
         if param_layout is None:
-            param_layout = _compute_default_per_buffer_param_layout(self.params, bucket_size)
+            param_layout = _compute_default_per_buffer_param_layout(
+                self.params, bucket_size, pad_param_starts=pad_param_starts
+            )
         self.param_index_map = param_layout.param_index_map
         self.bucket_indices = param_layout.bucket_indices
         per_bucket_numel_unpadded = param_layout.per_bucket_numel_unpadded

megatron/training/training.py

@@ -1597,6 +1597,18 @@ def wrap_model_chunks_with_ddp(
                     expert_data_parallel_world_size=expert_data_parallel_world_size,
                 )
 
+    # Under the layer-wise distributed optimizer with no pre-computed layout (Muon and other
+    # non-Adam/SGD optimizers using the legacy ``allgather_params`` sync path), DDP falls back
+    # to the default layout. Without per-param start padding, cuBLAS MXFP8 wgrad on
+    # cuBLASLt 12.8.x rejects the heuristic for mid-bucket params whose ``main_grad`` slice
+    # lands at low (16-byte) D-pointer alignment. Request 64-element start padding for that
+    # path; ignored when ``full_param_layout`` is supplied.
+    pad_param_starts = (
+        DP is DDP
+        and use_layer_wise_distributed_optimizer
+        and not use_layer_wise_param_layout
+    )
+
     # Wrap each chunk.
     wrapped = []
     for chunk, layout, disable_bucketing in zip(
@@ -1607,6 +1619,8 @@ def wrap_model_chunks_with_ddp(
             chunk_kwargs["pg_collection"] = pg_collection
         if layout is not None:
             chunk_kwargs["full_param_layout"] = layout
+        if pad_param_starts:
+            chunk_kwargs["pad_param_starts"] = True
         wrapped.append(
             DP(
                 config=config,