DDP refactoring: Extract parameter layout computation into optimizer classmethod

megatron/core/optimizer/distrib_optimizer.py

@@ -47,14 +47,19 @@
     ShardedTensorFactory,
 )
 from ..dist_checkpointing.utils import extract_sharded_tensors_and_factories
-from ..distributed.param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets
+from ..distributed.param_and_grad_buffer import (
+    _ParamAndGradBuffer,
+    group_params_for_buffers,
+    partition_buckets,
+)
 from ..fp4_utils import is_nvfp4tensor, quantize_nvfp4_param_shard
 from ..fp8_utils import dequantize_fp8_tensor, is_float8tensor, quantize_param_shard
 from ..transformer.fsdp_dtensor_checkpoint import handle_experts_in_state_dict
 from ..transformer.module import MegatronModule
 from .grad_scaler import MegatronGradScaler
 from .optimizer import MixedPrecisionOptimizer, _zero_grad_group_helper, param_group_identifier_keys
 from .optimizer_config import OptimizerConfig
+from .param_layout import FullParamLayout, PerBufferParamLayout, pad_bucket_end, pad_param_start
 
 logger = getLogger(__name__)
 
@@ -470,6 +475,133 @@ def _build_model_and_main_param_groups(
             shard_fp32_from_float16_groups,
         )
 
+    @staticmethod
+    def _compute_per_buffer_param_layout(
+        params: List[torch.nn.Parameter],
+        bucket_size: Optional[int],
+        data_parallel_world_size: int,
+        ddp_config,
+        param_indices: Optional[List[int]] = None,
+    ) -> 'PerBufferParamLayout':
+        """Compute how parameters should be laid out in the contiguous buffer.
+
+        Iterates params in reverse order (backprop order), applies 64-byte param
+        alignment, bucket-end padding for DP divisibility, and shared-embedding
+        bucket splitting.
+
+        Args:
+            params: List of parameters to lay out.
+            bucket_size: Approximate number of elements per bucket, or None for single bucket.
+            data_parallel_world_size: Size of the data-parallel group.
+            ddp_config: DistributedDataParallel config object.
+            param_indices: Optional indices for each param among same-dtype params.
+
+        Returns:
+            PerBufferParamLayout with the computed mapping.
+        """
+
+        def _does_param_require_new_bucket(param):
+            return getattr(param, "shared_embedding", False)
+
+        param_index_map = {}
+        bucket_indices = []
+        per_bucket_numel_unpadded = []
+
+        param_start_index = 0
+        bucket_start_index = 0
+        bucket_params = set()
+        bucket_id = 0
+
+        def _finalize_bucket(param_end_index, bucket_start_index, bucket_id):
+            per_bucket_numel_unpadded.append(param_end_index - bucket_start_index)
+            bucket_end_index = pad_bucket_end(
+                param_end_index,
+                data_parallel_world_size,
+                ddp_config.pad_buckets_for_high_nccl_busbw,
+            )
+            bucket_indices.append((bucket_start_index, bucket_end_index))
+            return bucket_end_index, bucket_id + 1
+
+        for param in params[::-1]:
+            param_start_index = pad_param_start(param_start_index)
+
+            # Split shared embedding params into separate bucket.
+            if _does_param_require_new_bucket(param) and len(bucket_params) > 0:
+                bucket_start_index, bucket_id = _finalize_bucket(
+                    param_start_index, bucket_start_index, bucket_id
+                )
+                bucket_params = set()
+                param_start_index = bucket_start_index
+
+            param_numel = param.data.nelement()
+            param_end_index = param_start_index + param_numel
+            param_index_map[param] = (param_start_index, param_end_index, bucket_id)
+            bucket_params.add(param)
+
+            if (
+                bucket_size is not None and (param_end_index - bucket_start_index) >= bucket_size
+            ) or _does_param_require_new_bucket(param):
+                bucket_start_index, bucket_id = _finalize_bucket(
+                    param_end_index, bucket_start_index, bucket_id
+                )
+                bucket_params = set()
+                param_start_index = bucket_start_index
+            else:
+                param_start_index = param_end_index
+
+        if len(bucket_params) > 0:
+            _finalize_bucket(param_end_index, bucket_start_index, bucket_id)
+
+        return PerBufferParamLayout(
+            param_index_map=param_index_map,
+            bucket_indices=bucket_indices,
+            per_bucket_numel_unpadded=per_bucket_numel_unpadded,
+            param_indices=param_indices if param_indices is not None else [],
+        )
+
+    @staticmethod
+    def compute_full_param_layout(
+        params: List[torch.nn.Parameter],
+        bucket_size: Optional[int],
+        data_parallel_world_size: int,
+        ddp_config,
+        expert_data_parallel_world_size: Optional[int] = None,
+    ) -> 'FullParamLayout':
+        """Compute parameter layouts for all buffer groups.
+
+        Groups parameters by (param_dtype, grad_dtype, is_expert_parallel), then
+        computes a padded PerBufferParamLayout for each group. Expert-parallel groups use
+        expert_data_parallel_world_size for padding alignment.
+
+        Args:
+            params: List of all parameters to lay out.
+            bucket_size: Approximate number of elements per bucket, or None for single bucket.
+            data_parallel_world_size: Size of the data-parallel group for dense params.
+            ddp_config: DistributedDataParallel config object.
+            expert_data_parallel_world_size: Size of the expert data-parallel group.
+                Required if any expert-parallel params are present. Defaults to
+                data_parallel_world_size if not provided.
+
+        Returns:
+            FullParamLayout with a PerBufferParamLayout per buffer group.
+        """
+        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():
+            if buffer_key.is_expert_parallel:
+                dp_world_size = (
+                    expert_data_parallel_world_size
+                    if expert_data_parallel_world_size is not None
+                    else data_parallel_world_size
+                )
+            else:
+                dp_world_size = data_parallel_world_size
+            layout = DistributedOptimizer._compute_per_buffer_param_layout(
+                group_params, bucket_size, dp_world_size, ddp_config, param_indices
+            )
+            layouts[buffer_key] = layout
+        return FullParamLayout(layouts=layouts)
+
     def __init__(
         self,
         optimizer: torch.optim.Optimizer,

megatron/core/optimizer/param_layout.py

@@ -0,0 +1,92 @@
+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+
+"""Parameter layout dataclasses for optimizer-driven buffer layout.
+
+These dataclasses describe how parameters are laid out in contiguous buffers.
+Each distributed optimizer implementation (e.g., DistributedOptimizer) is
+responsible for computing these layouts via a _compute_per_buffer_param_layout method,
+applying its own padding, alignment, and bucket splitting rules. DDP and
+buffers consume the resulting layouts without any optimizer-specific knowledge.
+"""
+
+import math
+from dataclasses import dataclass, field
+from typing import Dict, List, Tuple
+
+import torch
+
+
+def pad_to_divisor(value: int, divisor: int) -> int:
+    """Round up ``value`` to the nearest multiple of ``divisor``."""
+    return int(math.ceil(value / divisor) * divisor)
+
+
+def pad_param_start(param_start_index: int) -> int:
+    """Align parameter start index to a 64-element boundary."""
+    return pad_to_divisor(param_start_index, 64)
+
+
+def pad_bucket_end(
+    bucket_end_index: int, data_parallel_world_size: int, pad_for_high_nccl_busbw: bool
+) -> int:
+    """Pad bucket end for DP-divisibility (and optionally high NCCL bus bandwidth)."""
+    if pad_for_high_nccl_busbw:
+        divisor = math.lcm(data_parallel_world_size, 128, 2**16)
+    else:
+        divisor = math.lcm(data_parallel_world_size, 128)
+    return pad_to_divisor(bucket_end_index, divisor)
+
+
+@dataclass(frozen=True)
+class BufferKey:
+    """Identifies a distinct parameter buffer.
+
+    Each unique combination of these fields corresponds to a separate contiguous
+    buffer in DDP. Parameters are grouped into buffers by these dimensions.
+
+    Attributes:
+        param_dtype: Storage dtype (torch.uint8 for FP8/NVFP4 parameters, else param.dtype).
+        grad_dtype: Gradient reduction dtype.
+        is_expert_parallel: Whether the buffer holds expert-parallel parameters,
+            which use a separate data-parallel group.
+    """
+
+    param_dtype: torch.dtype
+    grad_dtype: torch.dtype
+    is_expert_parallel: bool
+
+
+@dataclass
+class PerBufferParamLayout:
+    """Layout for parameters within a single contiguous buffer.
+
+    Describes how parameters should be laid out in the contiguous buffer.
+
+    Attributes:
+        param_index_map: Mapping from parameter to (start_index, end_index, bucket_id) in buffer.
+        bucket_indices: List of (start_index, end_index) for each bucket.
+        per_bucket_numel_unpadded: Number of unpadded elements per bucket.
+        param_indices: The index of each param among same-dtype params (using the "fake"
+            high-precision dtype for FP8/NVFP4 params). Needed for loading non-native-fp8
+            checkpoints in native-fp8 mode. Order matches param_index_map iteration order.
+    """
+
+    param_index_map: Dict[torch.nn.Parameter, Tuple[int, int, int]] = field(default_factory=dict)
+    bucket_indices: List[Tuple[int, int]] = field(default_factory=list)
+    per_bucket_numel_unpadded: List[int] = field(default_factory=list)
+    param_indices: List[int] = field(default_factory=list)
+
+
+@dataclass
+class FullParamLayout:
+    """Layout for all parameters across all buffer groups in a model chunk.
+
+    Maps BufferKey to per-buffer PerBufferParamLayout objects. Each PerBufferParamLayout has its
+    own independent index space since different buffer groups are physically
+    separate buffers.
+
+    Attributes:
+        layouts: Mapping from BufferKey to PerBufferParamLayout.
+    """
+
+    layouts: Dict[BufferKey, PerBufferParamLayout] = field(default_factory=dict)

megatron/training/training.py

@@ -1506,18 +1506,49 @@ def build_model():
             dp_init_kwargs = {}
             if args.use_megatron_fsdp:
                 dp_init_kwargs["pg_collection"] = pg_collection
-            model = [
-                DP(
-                    config=config,
-                    ddp_config=ddp_config,
-                    module=model_chunk,
-                    # Turn off bucketing for model_chunk 2 onwards, since communication
-                    # for these model chunks is overlapped with compute anyway.
-                    disable_bucketing=(model_chunk_idx > 0) or args.overlap_param_gather_with_optimizer_step,
-                    **dp_init_kwargs,
+
+            wrapped_model = []
+            for model_chunk_idx, model_chunk in enumerate(model):
+                chunk_kwargs = dict(dp_init_kwargs)
+                disable_bucketing = (
+                    (model_chunk_idx > 0)
+                    or args.overlap_param_gather_with_optimizer_step
+                )
+
+                # Pre-compute parameter layouts for the distributed optimizer.
+                # Only pass to DDP; FSDP variants don't accept full_param_layout.
+                if args.use_distributed_optimizer and DP is DDP:
+                    all_params = [
+                        p for p in model_chunk.parameters() if p.requires_grad
+                    ]
+                    pp_rank = mpu.get_pipeline_model_parallel_rank()
+                    effective_bucket_size = (
+                        None
+                        if disable_bucketing or pp_rank > 0
+                        else ddp_config.bucket_size
+                    )
+                    chunk_kwargs["full_param_layout"] = (
+                        DistributedOptimizer.compute_full_param_layout(
+                            all_params,
+                            effective_bucket_size,
+                            mpu.get_data_parallel_world_size(with_context_parallel=True),
+                            ddp_config,
+                            expert_data_parallel_world_size=(
+                                mpu.get_expert_data_parallel_world_size()
+                            ),
+                        )
+                    )
+
+                wrapped_model.append(
+                    DP(
+                        config=config,
+                        ddp_config=ddp_config,
+                        module=model_chunk,
+                        disable_bucketing=disable_bucketing,
+                        **chunk_kwargs,
+                    )
                 )
-                for (model_chunk_idx, model_chunk) in enumerate(model)
-            ]
+            model = wrapped_model
         # End of setup_stream
         # Critical: ensure side-stream work completes before touching params on default stream
         torch.cuda.current_stream().wait_stream(ddp_stream)

tests/unit_tests/dist_checkpointing/test_optimizer.py

@@ -589,7 +589,7 @@ def test_bucket_space_optimizer_save_load(
             ) as ckpt_dir_B,
         ):
             # Init model and optimizer with "src" bucket padding
-            with patch('megatron.core.distributed.param_and_grad_buffer.math.lcm') as lcm_mock:
+            with patch('megatron.core.optimizer.param_layout.math.lcm') as lcm_mock:
                 lcm_mock.return_value = src_bucket_pad_divisor
 
                 model_A, optimizer_A = setup_model_and_optimizer(
@@ -615,7 +615,7 @@ def test_bucket_space_optimizer_save_load(
                 parallel_state.get_model_parallel_group()
             )
             # Init model and optimizer with "dest" bucket padding
-            with patch('megatron.core.distributed.param_and_grad_buffer.math.lcm') as lcm_mock:
+            with patch('megatron.core.optimizer.param_layout.math.lcm') as lcm_mock:
                 lcm_mock.return_value = dest_bucket_pad_divisor
 
                 model_B, optimizer_B = setup_model_and_optimizer(