diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py index cfd91031b44..f743b044064 100644 --- a/megatron/core/distributed/distributed_data_parallel.py +++ b/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) diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 309ac7ef6d5..6c61f29e193 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/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 diff --git a/megatron/training/training.py b/megatron/training/training.py index f69e4f30f6a..77a2ddfdf52 100644 --- a/megatron/training/training.py +++ b/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,