diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py index 35325d70ce9..06f3fb32238 100644 --- a/megatron/core/distributed/distributed_data_parallel.py +++ b/megatron/core/distributed/distributed_data_parallel.py @@ -7,15 +7,15 @@ import torch from ..config_logger import has_config_logger_enabled, log_config_to_disk -from ..fp4_utils import is_nvfp4tensor from ..fp8_utils import is_float8tensor, post_all_gather_processing +from ..optimizer.param_layout import FullParamLayout from ..process_groups_config import ProcessGroupCollection from ..transformer.cuda_graphs import is_graph_capturing from ..transformer.transformer_config import TransformerConfig from ..utils import log_single_rank from .data_parallel_base import _BaseDataParallel from .distributed_data_parallel_config import DistributedDataParallelConfig -from .param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets +from .param_and_grad_buffer import _ParamAndGradBuffer, group_params_for_buffers, partition_buckets logger = logging.getLogger(__name__) @@ -36,6 +36,9 @@ class DistributedDataParallel(_BaseDataParallel): use standard bucketing policy: assign parameters to smaller buckets and all-reduce per bucket _if_ overlap_grad_reduce is True and pp_rank is 0. pg_collection: Optional unified process group for distributed training. + 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. """ @@ -46,6 +49,7 @@ def __init__( module: torch.nn.Module, disable_bucketing: bool = False, pg_collection: Optional[ProcessGroupCollection] = None, + full_param_layout: Optional[FullParamLayout] = None, ): super().__init__(config=config, module=module) if has_config_logger_enabled(config): @@ -104,11 +108,10 @@ def __init__( self.param_to_bucket_group = {} - # Group parameters by their gradient type. + # Collect all trainable parameters. param_to_name = {} - dense_params = [] - expert_parallel_params = [] self.params_with_grad = [] + all_params = [] for name, param in self.module.named_parameters(): if not param.requires_grad: continue @@ -119,142 +122,50 @@ def __init__( param.grad_added_to_main_grad = False param_to_name[param] = name + all_params.append(param) + + # Group parameters by (param_dtype, grad_dtype, is_expert_parallel). + buffer_groups = group_params_for_buffers(all_params, self.ddp_config.grad_reduce_in_fp32) + + # Auto-compute layouts when using distributed optimizer but no layout was provided. + # This maintains backward compatibility for callers that create DDP directly + # without pre-computing layouts (e.g., tests, external code). + if full_param_layout is None and self.ddp_config.use_distributed_optimizer: + log_single_rank( + logger, + logging.WARNING, + "DistributedDataParallel: full_param_layout not provided with " + "use_distributed_optimizer=True. Auto-computing layout inside DDP. " + "Callers should pre-compute layouts via " + "DistributedOptimizer.compute_full_param_layout() and pass them in.", + ) + from ..optimizer.distrib_optimizer import DistributedOptimizer + + full_param_layout = DistributedOptimizer.compute_full_param_layout( + all_params, + self.bucket_size, + self.intra_dp_cp_group.size(), + self.ddp_config, + expert_data_parallel_world_size=self.intra_expt_dp_group.size(), + ) - if getattr(param, 'allreduce', True): - dense_params.append((param, name)) - else: - expert_parallel_params.append((param, name)) - - def _allocate_buffers_for_parameters( - input_params, data_parallel_group, gradient_scaling_factor - ): - param_and_grad_dtype_to_params = {} - param_and_grad_dtype_to_offsets = {} - param_and_grad_dtype_to_indices = {} - - # Group parameters by their gradient type. - for param, param_name in input_params: - assert param.requires_grad - - param_dtype = param.dtype - if is_float8tensor(param) or is_nvfp4tensor(param): - # Currently TE's Float8Tensor is a wrapper of torch.Tensor. It has a "fake" - # dtype (usually a higher precision dtype such as bfloat16), but its actual - # data is stored in the form of a torch uint8 tensor within the Float8Tensor's - # ".data" attribute. Therefore, when creating the param buffer for fp8/fp4 - # params,it is necessary to use torch.uint8, not the "fake" dtype got from - # "param.dtype". - param_dtype = torch.uint8 - grad_dtype = torch.float if self.ddp_config.grad_reduce_in_fp32 else param.dtype - - params = param_and_grad_dtype_to_params.get((param_dtype, grad_dtype), []) - params.append((param, param_name)) - param_and_grad_dtype_to_params[(param_dtype, grad_dtype)] = params - - # Get the index of each param among the params with same dtype, if a param is fp8, - # use its "fake" high precision dtype to find which params have same dtype with it. - # For example: - # Case 1: - # params = [p1(bf16), p2(bf16), p3(bf16), p4(bf16)] - # param_and_grad_dtype_to_indices = { - # (torch.bfloat16, torch.float32): [0, 1, 2, 3], - # } - # Case 2: - # params = [p1(bf16), p2(fp8), p3(fp8), p4(bf16)] - # param_and_grad_dtype_to_indices = { - # (torch.bfloat16, torch.float32): [0, 3], - # (torch.uint8, torch.float32): [1, 2], - # } - # We need these indices to load a non-native-fp8 checkpoint in native-fp8 mode. - offset = param_and_grad_dtype_to_offsets.get((param.dtype, grad_dtype), 0) - param_and_grad_dtype_to_offsets[(param.dtype, grad_dtype)] = offset + 1 - indices = param_and_grad_dtype_to_indices.get((param_dtype, grad_dtype), []) - indices.append(offset) - param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)] = indices - - if not config.calculate_per_token_loss: - target_gradient_scaling_factor = 1.0 / self.dp_cp_group.size() - if self.ddp_config.average_in_collective: - if self.ddp_config.num_distributed_optimizer_instances == 1: - # Collective is averaging gradients in collective with data_parallel_group. - assert ( - gradient_scaling_factor / data_parallel_group.size() - == target_gradient_scaling_factor - ) - else: - # For non-expert parameters, gradient_scaling_factor is 1. - # For expert parameters, gradient_scaling_factor is edp_size/dp_size. - assert (gradient_scaling_factor == 1) or ( - gradient_scaling_factor - == (self.expt_dp_group.size() / self.dp_cp_group.size()) - ) - else: - assert gradient_scaling_factor == target_gradient_scaling_factor - - # Allocate the grad buffers and map the grads. - buffers = [] - pg_collection = ProcessGroupCollection() - pg_collection.tp = self.tp_group - pg_collection.dp_cp = self.dp_cp_group - for (param_dtype, grad_dtype), params in param_and_grad_dtype_to_params.items(): - buffers.append( - _ParamAndGradBuffer( - self.ddp_config, - param_dtype, - grad_dtype, - params, - data_parallel_group, - self.bucket_size, - param_to_name, - gradient_scaling_factor, - param_and_grad_dtype_to_indices[(param_dtype, grad_dtype)], - self.ddp_config.nccl_ub, - pg_collection, - ) - ) - - # In some scenarios, we want to put buckets from different buffers into a group so that - # their communication can be aggregated. For example, when there are both fp8 buffers - # and bf16 buffers in the model and vpp is enabled, each model chunk will have an fp8 - # bucket and a bf16 bucket, which doubles the number of communication kernels, and - # because of the use of CUDA_DEVICE_MAX_CONNECTIONS=1, having multiple back-to-back - # communications will prevent the overlap of the communication kernels with computation - # kernels. - # If bucketing is explicitly disabled, then put all buckets in a buffer into a single - # bucket group. - bucket_groups = partition_buckets(buffers, force_single_bucket_group=disable_bucketing) - - if self.ddp_config.num_distributed_optimizer_instances > 1: + # When a full_param_layout is provided, verify that the grouping is consistent + # with the layout (same buffer keys, same params per key, same param_indices). + if full_param_layout is not None: + assert set(buffer_groups.keys()) == set(full_param_layout.layouts.keys()), ( + f"Buffer keys from param grouping {set(buffer_groups.keys())} do not match " + f"full_param_layout keys {set(full_param_layout.layouts.keys())}" + ) + for buffer_key, (params, param_indices) in buffer_groups.items(): + layout = full_param_layout.layouts[buffer_key] + assert set(params) == set( + layout.param_index_map.keys() + ), f"Params for {buffer_key} do not match between grouping and layout" assert ( - self.ddp_config.use_distributed_optimizer - ), 'Partial DistOpt cannot be used without DistOpt' - communication_stream = torch.cuda.Stream(device=torch.cuda.current_device()) - for bucket_group in bucket_groups: - bucket_group.inter_distributed_optimizer_instance_group = ( - self.inter_dist_opt_group - ) - bucket_group.communication_stream = communication_stream - - # Set `next_param_gather_bucket_group` for different bucket groups by iterating through - # buckets in reverse order (since all-gathers happen in reverse order of buckets). - # Note: overlap_param_gather covers both the distributed optimizer and the - # layer-wise optimizer cases; the latter sets overlap_param_gather=True - # without use_distributed_optimizer. - if self.ddp_config.overlap_param_gather: - num_bucket_groups = len(bucket_groups) - for i in range(1, num_bucket_groups): - bucket_groups[num_bucket_groups - i].next_param_gather_bucket_group = ( - bucket_groups[num_bucket_groups - i - 1] - ) - - # Create map from param to bucket group, used in pre_hook. - for bucket_group in bucket_groups: - for bucket in bucket_group.buckets: - for param in bucket.params_list: - self.param_to_bucket_group[param] = bucket_group - - return buffers, bucket_groups + param_indices == layout.param_indices + ), f"param_indices for {buffer_key} do not match between grouping and layout" + # Compute gradient scaling factors. if config.calculate_per_token_loss: assert ( not self.ddp_config.average_in_collective @@ -291,20 +202,107 @@ def _allocate_buffers_for_parameters( gradient_scaling_factor = 1.0 / data_parallel_world_size expert_gradient_scaling_factor = 1.0 / data_parallel_world_size - # Allocate the param+grad buffers for dense params' grads. - self.buffers, self.bucket_groups = _allocate_buffers_for_parameters( - dense_params, self.intra_dp_cp_group, gradient_scaling_factor=gradient_scaling_factor - ) + # Allocate buffers for each group. + self.buffers = [] + self.expert_parallel_buffers = [] + pg_collection = ProcessGroupCollection(tp=self.tp_group, dp_cp=self.dp_cp_group) + for buffer_key, (params, param_indices) in buffer_groups.items(): + if buffer_key.is_expert_parallel: + data_parallel_group = self.intra_expt_dp_group + scaling_factor = expert_gradient_scaling_factor + else: + data_parallel_group = self.intra_dp_cp_group + scaling_factor = gradient_scaling_factor - # Allocate separate param+grad buffers for expert parallel params' grads. - self.expert_parallel_buffers, self.expert_parallel_bucket_groups = ( - _allocate_buffers_for_parameters( - expert_parallel_params, - self.intra_expt_dp_group, - gradient_scaling_factor=expert_gradient_scaling_factor, + if not config.calculate_per_token_loss: + target_gradient_scaling_factor = 1.0 / self.dp_cp_group.size() + if self.ddp_config.average_in_collective: + if self.ddp_config.num_distributed_optimizer_instances == 1: + # Collective is averaging gradients in collective with data_parallel_group. + assert ( + scaling_factor / data_parallel_group.size() + == target_gradient_scaling_factor + ) + else: + # For non-expert parameters, gradient_scaling_factor is 1. + # For expert parameters, gradient_scaling_factor is edp_size/dp_size. + assert (scaling_factor == 1) or ( + scaling_factor == (self.expt_dp_group.size() / self.dp_cp_group.size()) + ) + else: + assert scaling_factor == target_gradient_scaling_factor + + param_layout = ( + full_param_layout.layouts.get(buffer_key) if full_param_layout is not None else None ) + params_with_names = [(p, param_to_name[p]) for p in params] + buffer = _ParamAndGradBuffer( + self.ddp_config, + buffer_key.param_dtype, + buffer_key.grad_dtype, + params_with_names, + data_parallel_group, + self.bucket_size, + param_to_name, + scaling_factor, + param_indices, + self.ddp_config.nccl_ub, + pg_collection, + param_layout=param_layout, + ) + if buffer_key.is_expert_parallel: + self.expert_parallel_buffers.append(buffer) + else: + self.buffers.append(buffer) + + # In some scenarios, we want to put buckets from different buffers into a group so that + # their communication can be aggregated. For example, when there are both fp8 buffers + # and bf16 buffers in the model and vpp is enabled, each model chunk will have an fp8 + # bucket and a bf16 bucket, which doubles the number of communication kernels, and + # because of the use of CUDA_DEVICE_MAX_CONNECTIONS=1, having multiple back-to-back + # communications will prevent the overlap of the communication kernels with computation + # kernels. + # If bucketing is explicitly disabled, then put all buckets in a buffer into a single + # bucket group. + self.bucket_groups = partition_buckets( + self.buffers, force_single_bucket_group=disable_bucketing + ) + self.expert_parallel_bucket_groups = partition_buckets( + self.expert_parallel_buffers, force_single_bucket_group=disable_bucketing ) + if self.ddp_config.num_distributed_optimizer_instances > 1: + assert ( + self.ddp_config.use_distributed_optimizer + ), 'Partial DistOpt cannot be used without DistOpt' + for bucket_groups in [self.bucket_groups, self.expert_parallel_bucket_groups]: + communication_stream = torch.cuda.Stream(device=torch.cuda.current_device()) + for bucket_group in bucket_groups: + bucket_group.inter_distributed_optimizer_instance_group = ( + self.inter_dist_opt_group + ) + bucket_group.communication_stream = communication_stream + + # Set `next_param_gather_bucket_group` for different bucket groups by iterating through + # buckets in reverse order (since all-gathers happen in reverse order of buckets). + # Note: overlap_param_gather covers both the distributed optimizer and the + # layer-wise optimizer cases; the latter sets overlap_param_gather=True + # without use_distributed_optimizer. + if self.ddp_config.overlap_param_gather: + for bucket_groups in [self.bucket_groups, self.expert_parallel_bucket_groups]: + num_bucket_groups = len(bucket_groups) + for i in range(1, num_bucket_groups): + bucket_groups[num_bucket_groups - i].next_param_gather_bucket_group = ( + bucket_groups[num_bucket_groups - i - 1] + ) + + # Create map from param to bucket group, used in pre_hook. + for bucket_groups in [self.bucket_groups, self.expert_parallel_bucket_groups]: + for bucket_group in bucket_groups: + for bucket in bucket_group.buckets: + for param in bucket.params_list: + self.param_to_bucket_group[param] = bucket_group + # Delete references to weight_tensor if they exist since we don't want two parameter copies # if we re-mapped parameters (which happens when we use the distributed optimizer). # This is a temporary workaround around a TE bug that is fixed with diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index fe0a51b86b7..dc3014d72d2 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -27,6 +27,7 @@ modify_underlying_storage, post_all_gather_processing, ) +from ..optimizer.param_layout import pad_bucket_end, pad_param_start from ..utils import is_torch_min_version, log_on_each_pipeline_stage from .distributed_data_parallel_config import DistributedDataParallelConfig from .reduce_scatter_with_fp32_accumulation import reduce_scatter_with_fp32_accumulation @@ -753,6 +754,114 @@ def register_grad_ready( self.start_grad_sync(force_all_reduce=force_all_reduce) +def group_params_for_buffers( + params: List[torch.nn.Parameter], grad_reduce_in_fp32: bool +) -> Dict['BufferKey', Tuple[List[torch.nn.Parameter], List[int]]]: + """Group parameters by buffer identity for buffer allocation. + + Each distinct buffer is identified by a BufferKey with three dimensions: + - param_dtype: storage dtype (torch.uint8 for FP8/NVFP4 parameters, else param.dtype). + - grad_dtype: gradient reduction dtype (torch.float if grad_reduce_in_fp32, else param.dtype). + - is_expert_parallel: whether the parameter is expert-parallel (param.allreduce == False), + which requires a separate buffer with a different data-parallel group. + + The param_indices track each parameter's position 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. + + Args: + params: List of parameters to group. + grad_reduce_in_fp32: Whether gradients are reduced in FP32. + + Returns: + Dict mapping BufferKey to (params_list, param_indices). + """ + from ..optimizer.param_layout import BufferKey + + key_to_params = {} + dtype_to_offsets = {} + key_to_indices = {} + + for param in params: + assert param.requires_grad + + param_dtype = param.dtype + if is_float8tensor(param) or is_nvfp4tensor(param): + param_dtype = torch.uint8 + grad_dtype = torch.float if grad_reduce_in_fp32 else param.dtype + is_expert_parallel = not getattr(param, 'allreduce', True) + + key = BufferKey(param_dtype, grad_dtype, is_expert_parallel) + 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 = dtype_to_offsets.get(offset_key, 0) + dtype_to_offsets[offset_key] = offset + 1 + indices = key_to_indices.get(key, []) + indices.append(offset) + key_to_indices[key] = indices + + result = {} + for key, param_list in key_to_params.items(): + result[key] = (param_list, key_to_indices[key]) + return result + + +def _compute_default_per_buffer_param_layout( + params: List[torch.nn.Parameter], bucket_size: Optional[int] +) -> '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. + + Args: + params: List of parameters to lay out. + bucket_size: Approximate number of elements per bucket, or None for a single bucket. + + Returns: + PerBufferParamLayout with the computed mapping. + """ + from ..optimizer.param_layout import PerBufferParamLayout + + param_index_map = {} + bucket_indices = [] + per_bucket_numel_unpadded = [] + + param_start_index = 0 + bucket_start_index = 0 + bucket_params = set() + bucket_id = 0 + + for param in params[::-1]: + 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) + bucket_params.add(param) + + if bucket_size is not None and (param_end_index - bucket_start_index) >= bucket_size: + per_bucket_numel_unpadded.append(param_end_index - bucket_start_index) + bucket_indices.append((bucket_start_index, param_end_index)) + bucket_start_index = param_end_index + bucket_params = set() + bucket_id += 1 + param_start_index = param_end_index + + if len(bucket_params) > 0: + per_bucket_numel_unpadded.append(param_end_index - bucket_start_index) + bucket_indices.append((bucket_start_index, param_end_index)) + + return PerBufferParamLayout( + param_index_map=param_index_map, + bucket_indices=bucket_indices, + per_bucket_numel_unpadded=per_bucket_numel_unpadded, + ) + + class _ParamAndGradBuffer: """ Groups parameters and gradients into a contiguous buffer, and then breaks the buffer into @@ -788,6 +897,7 @@ def __init__( param_indices: List[int], nccl_ub: bool, pg_collection: Optional[ProcessGroupCollection] = None, + param_layout: Optional['PerBufferParamLayout'] = None, ): if pg_collection is None: @@ -822,121 +932,13 @@ def __init__( # Data structures to store underlying buckets and relevant indexing data. self.buckets = [] self.param_to_bucket = {} # Param -> bucket mapping. - self.param_index_map = {} # Param -> location in buffer mapping (used in dist. optimizer). - - def _pad(number_to_be_padded: int, divisor: int) -> int: - return int(math.ceil(number_to_be_padded / divisor) * divisor) - - def _pad_end_of_bucket_if_needed(bucket_end_index: int) -> int: - """ - Pads end index of bucket if using distributed optimizer (to ensure uniform sharding). - """ - if self.ddp_config.use_distributed_optimizer: - # Workaround for TE bug causing cuBLAS to pick an incompatible algorithm. - # This also helps cuBLAS pick more efficient algorithms for GEMMs. - # We now ensure that all buckets start at a memory address that is 256-byte - # aligned (128 values since params and grads use >= 16-bit precision). - if self.ddp_config.pad_buckets_for_high_nccl_busbw: - # Make sure the bucket size is divisible by a large power of 2 (2^16) to - # ensure NCCL collectives have high bus bandwidth at large DP counts, - # since NCCL message size (which for ring algorithms is bucket_size / - # dp_size) apparently needs to be divisible by a power of 2 for high busbw. - bucket_size_divisor = math.lcm(self.data_parallel_world_size, 128, 2**16) - else: - bucket_size_divisor = math.lcm(self.data_parallel_world_size, 128) - return _pad(bucket_end_index, bucket_size_divisor) - return bucket_end_index - - def _pad_start_of_param_if_needed(param_start_index: int) -> int: - """ - Pads start index of param if using distributed optimizer (to ensure "good" alignment). - """ - if self.ddp_config.use_distributed_optimizer: - # Ensure that params start at 128-byte aligned addresses (64 values - # since params are >= 16-bit precision). - return _pad(param_start_index, 64) - return param_start_index - # First, figure out how many elements should be in the underlying buffer storage. - # Note that if we need to split the buffer into smaller buckets, each of these - # might need to be padded as well (if using the distributed optimizer). - param_start_index = 0 - bucket_start_index = param_start_index - bucket_params = set() - self.bucket_indices = [] - per_bucket_numel_unpadded = [] - bucket_id = 0 - - def _update_bucket_metadata( - param_end_index: int, - bucket_start_index: int, - bucket_indices: list, - numel_unpadded_list: list, - ) -> int: - """ - Record metadata for a bucket. Returns the bucket's (padded) end_index. - - Args: - param_end_index: End index of the last param in this bucket (unpadded). - bucket_start_index: Start index of this bucket. - bucket_indices: List to append (start, end) bucket boundaries to. - numel_unpadded_list: List to append unpadded bucket numel to. - - Returns: - The bucket's end index, padded if using distributed optimizer. - """ - numel_unpadded_list.append(param_end_index - bucket_start_index) - bucket_end_index = _pad_end_of_bucket_if_needed(param_end_index) - bucket_indices.append((bucket_start_index, bucket_end_index)) - return bucket_end_index - - def _finalize_bucket_all_index_spaces( - param_end_index: int, - bucket_start_index: int, - nvfp4_packed_param_end_index: int = None, - nvfp4_packed_bucket_start_index: int = None, - ) -> tuple: - """ - Record metadata for the current bucket across both main and (if applicable) - NVFP4 packed index spaces. Also resets bucket_params and increments bucket_id. - - Args: - param_end_index: End index of the last param in the bucket (full numel). - bucket_start_index: Start index of the bucket (full numel). - nvfp4_packed_param_end_index: End index in packed space (NVFP4 only). - nvfp4_packed_bucket_start_index: Bucket start in packed space (NVFP4 only). - - Returns: - Tuple of (bucket_end_index, nvfp4_packed_bucket_end_index). - """ - nonlocal bucket_params, bucket_id - bucket_end_index = _update_bucket_metadata( - param_end_index, bucket_start_index, self.bucket_indices, per_bucket_numel_unpadded - ) - nvfp4_packed_bucket_end_index = None - if self.has_nvfp4_params: - nvfp4_packed_bucket_end_index = _update_bucket_metadata( - nvfp4_packed_param_end_index, - nvfp4_packed_bucket_start_index, - self.nvfp4_packed_bucket_indices, - nvfp4_packed_per_bucket_numel_unpadded, - ) - bucket_params = set() - bucket_id += 1 - return bucket_end_index, nvfp4_packed_bucket_end_index - - def _does_param_require_new_bucket(param): - """ - Split shared embedding parameters into separate bucket if using distributed - optimizer that makes use of reduce-scatters instead of all-reduces. - This ensures that the first and last pipeline stage partition optimizer state - for the shared embedding parameters the same way across DP replicas, allowing - the DP reduce-scatter to be before the embedding all-reduce. - """ - return ( - getattr(param, "shared_embedding", False) - and self.ddp_config.use_distributed_optimizer - ) + # Use the provided layout if given, otherwise compute the default (no-padding) layout. + if param_layout is None: + param_layout = _compute_default_per_buffer_param_layout(self.params, bucket_size) + 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 # Check if this buffer contains NVFP4 params. # @@ -953,95 +955,17 @@ def _does_param_require_new_bucket(param): # Grad buffer: [g0, g1, g2, g3, ...] numel = N # # We therefore maintain two index maps: - # - param_index_map: offsets using full numel. + # - param_index_map: offsets using full numel (from pre-computed layout). # - nvfp4_packed_param_index_map: offsets into the packed param buffer (numel // 2). # + # The packed index map is derived from param_index_map by iterating through + # the already-computed layout and halving numel for NVFP4 tensors. + # self.has_nvfp4_params = any(is_nvfp4tensor(p) for p in self.params) - # Secondary (packed) index map, counters, and bucket tracking for NVFP4. - self.nvfp4_packed_param_index_map = {} if self.has_nvfp4_params else None - nvfp4_packed_param_start_index = 0 if self.has_nvfp4_params else None - nvfp4_packed_param_end_index = None - nvfp4_packed_bucket_start_index = 0 if self.has_nvfp4_params else None - self.nvfp4_packed_bucket_indices = [] if self.has_nvfp4_params else None - nvfp4_packed_per_bucket_numel_unpadded = [] if self.has_nvfp4_params else None - - for param, _ in params_with_names[::-1]: - # Iterate through parameters in reverse order to roughly follow backprop order. - - param_start_index = _pad_start_of_param_if_needed(param_start_index) - if self.has_nvfp4_params: - nvfp4_packed_param_start_index = _pad_start_of_param_if_needed( - nvfp4_packed_param_start_index - ) - - # Create bucket with collected parameters if current param needs its own bucket. - if _does_param_require_new_bucket(param) and len(bucket_params) > 0: - # Finalize the current bucket and update start indices for the next bucket. - bucket_start_index, nvfp4_packed_bucket_start_index = ( - _finalize_bucket_all_index_spaces( - param_start_index, - bucket_start_index, - nvfp4_packed_param_start_index, - nvfp4_packed_bucket_start_index, - ) - ) - param_start_index = bucket_start_index - if self.has_nvfp4_params: - nvfp4_packed_param_start_index = nvfp4_packed_bucket_start_index - - # Primary index computation: always uses full param numel. - param_numel = param.data.nelement() - param_end_index = param_start_index + param_numel - self.param_index_map[param] = (param_start_index, param_end_index, bucket_id) - - # Secondary (packed) index computation for NVFP4. - if self.has_nvfp4_params: - if is_nvfp4tensor(param): - assert ( - param_numel % 2 == 0 - ), f"NVFP4 requires even numel for packing, got {param_numel}" - # NVFP4 packs two FP4 values into one byte, so packed numel is half. - nvfp4_packed_param_end_index = nvfp4_packed_param_start_index + param_numel // 2 - else: - nvfp4_packed_param_end_index = nvfp4_packed_param_start_index + param_numel - self.nvfp4_packed_param_index_map[param] = ( - nvfp4_packed_param_start_index, - nvfp4_packed_param_end_index, - bucket_id, - ) - - bucket_params.add(param) - - # If we have enough elements already or the current param is part of the shared - # embedding layer and needs a separate bucket, form a new bucket. - if ( - bucket_size is not None and (param_end_index - bucket_start_index) >= bucket_size - ) or _does_param_require_new_bucket(param): - # Finalize the current bucket and update start indices for the next bucket. - bucket_start_index, nvfp4_packed_bucket_start_index = ( - _finalize_bucket_all_index_spaces( - param_end_index, - bucket_start_index, - nvfp4_packed_param_end_index, - nvfp4_packed_bucket_start_index, - ) - ) - param_start_index = bucket_start_index - if self.has_nvfp4_params: - nvfp4_packed_param_start_index = nvfp4_packed_bucket_start_index - else: - param_start_index = param_end_index - if self.has_nvfp4_params: - nvfp4_packed_param_start_index = nvfp4_packed_param_end_index - - # Add remaining params to a new bucket. - if len(bucket_params) > 0: - _finalize_bucket_all_index_spaces( - param_end_index, - bucket_start_index, - nvfp4_packed_param_end_index, - nvfp4_packed_bucket_start_index, - ) + self.nvfp4_packed_param_index_map = None + self.nvfp4_packed_bucket_indices = None + if self.has_nvfp4_params: + self._compute_nvfp4_packed_layout(params_with_names) # Next, create underlying storage for buffer (with numel elements that includes # padding as necessary). @@ -1049,13 +973,14 @@ def _does_param_require_new_bucket(param): self.numel_unpadded = sum(per_bucket_numel_unpadded) if self.has_nvfp4_params: self.nvfp4_packed_numel = self.nvfp4_packed_bucket_indices[-1][1] - self.nvfp4_packed_numel_unpadded = sum(nvfp4_packed_per_bucket_numel_unpadded) + # nvfp4_packed_numel_unpadded is already set by _compute_nvfp4_packed_layout. assert self.numel_unpadded <= self.numel + if self.has_nvfp4_params: + assert self.nvfp4_packed_numel_unpadded <= self.nvfp4_packed_numel if self.ddp_config.use_distributed_optimizer: assert self.numel % self.data_parallel_world_size == 0 if self.has_nvfp4_params: - assert self.nvfp4_packed_numel_unpadded <= self.nvfp4_packed_numel assert self.nvfp4_packed_numel % self.data_parallel_world_size == 0 else: assert self.numel == self.numel_unpadded @@ -1286,6 +1211,93 @@ def _create_bucket(bucket_id, bucket_params, bucket_params_with_extra_main_grads dp_cp_group=self.dp_cp_group, ) + def _compute_nvfp4_packed_layout(self, params_with_names): + """Derive packed NVFP4 index map and bucket indices from the primary layout. + + The primary layout (self.param_index_map, self.bucket_indices) uses full numel + for all params. NVFP4 tensors pack two FP4 values into one byte, so the param + buffer needs a separate "packed" index map where NVFP4 params occupy half the + space. Non-NVFP4 params keep their full numel in the packed space. + + The same padding rules used by the primary layout are applied here: + - 64-element alignment at the start of each param. + - Bucket-end padding for DP-divisibility (when using distributed optimizer). + + Sets: + self.nvfp4_packed_param_index_map: param -> (start, end, bucket_id) + self.nvfp4_packed_bucket_indices: list of (start, end) per bucket + self.nvfp4_packed_numel_unpadded: total unpadded elements across all buckets + """ + + def _pad_start_of_param(param_start_index: int) -> int: + if self.ddp_config.use_distributed_optimizer: + return pad_param_start(param_start_index) + return param_start_index + + def _pad_end_of_bucket(bucket_end_index: int) -> int: + if self.ddp_config.use_distributed_optimizer: + return pad_bucket_end( + bucket_end_index, + self.data_parallel_world_size, + self.ddp_config.pad_buckets_for_high_nccl_busbw, + ) + return bucket_end_index + + self.nvfp4_packed_param_index_map = {} + self.nvfp4_packed_bucket_indices = [] + nvfp4_packed_per_bucket_numel_unpadded = [] + + packed_param_start = 0 + packed_bucket_start = 0 + cur_bucket_id = 0 + + for param, _ in params_with_names[::-1]: + _, _, bucket_id = self.param_index_map[param] + param_numel = param.data.nelement() + + packed_param_start = _pad_start_of_param(packed_param_start) + + # Finalize previous bucket if we've moved to a new one. + if bucket_id != cur_bucket_id: + # Record unpadded numel, then pad the bucket end. + nvfp4_packed_per_bucket_numel_unpadded.append( + packed_param_start - packed_bucket_start + ) + packed_bucket_end = _pad_end_of_bucket(packed_param_start) + self.nvfp4_packed_bucket_indices.append((packed_bucket_start, packed_bucket_end)) + packed_bucket_start = packed_bucket_end + packed_param_start = packed_bucket_start + cur_bucket_id = bucket_id + + # NVFP4 tensors use half the numel in the packed param buffer. + if is_nvfp4tensor(param): + assert ( + param_numel % 2 == 0 + ), f"NVFP4 requires even numel for packing, got {param_numel}" + packed_numel = param_numel // 2 + else: + packed_numel = param_numel + + packed_param_end = packed_param_start + packed_numel + self.nvfp4_packed_param_index_map[param] = ( + packed_param_start, + packed_param_end, + bucket_id, + ) + packed_param_start = packed_param_end + + # Finalize last bucket. + if packed_param_start > packed_bucket_start: + nvfp4_packed_per_bucket_numel_unpadded.append(packed_param_start - packed_bucket_start) + packed_bucket_end = _pad_end_of_bucket(packed_param_start) + self.nvfp4_packed_bucket_indices.append((packed_bucket_start, packed_bucket_end)) + + assert len(self.nvfp4_packed_bucket_indices) == len(self.bucket_indices), ( + f"Packed bucket count ({len(self.nvfp4_packed_bucket_indices)}) != " + f"primary bucket count ({len(self.bucket_indices)})" + ) + self.nvfp4_packed_numel_unpadded = sum(nvfp4_packed_per_bucket_numel_unpadded) + def scale_gradients(self, scaling_factor: float) -> None: """Scale the gradient data by `scaling_factor`.""" self.grad_data *= scaling_factor diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index b4d52e6b56b..95e8a0c407b 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -47,7 +47,11 @@ 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 @@ -55,6 +59,7 @@ 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, diff --git a/megatron/core/optimizer/param_layout.py b/megatron/core/optimizer/param_layout.py new file mode 100644 index 00000000000..6ebcc348f84 --- /dev/null +++ b/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) diff --git a/megatron/training/training.py b/megatron/training/training.py index c95d386c748..2b929f9067d 100644 --- a/megatron/training/training.py +++ b/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) diff --git a/tests/unit_tests/dist_checkpointing/test_optimizer.py b/tests/unit_tests/dist_checkpointing/test_optimizer.py index 1db844dd9a8..149323707de 100644 --- a/tests/unit_tests/dist_checkpointing/test_optimizer.py +++ b/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( diff --git a/tests/unit_tests/distributed/test_param_and_grad_buffer.py b/tests/unit_tests/distributed/test_param_and_grad_buffer.py index 223383bba64..ac1fdfe2ed6 100644 --- a/tests/unit_tests/distributed/test_param_and_grad_buffer.py +++ b/tests/unit_tests/distributed/test_param_and_grad_buffer.py @@ -11,10 +11,39 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig from megatron.core.distributed.param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets +from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer from megatron.core.transformer import TransformerConfig from tests.unit_tests.test_utilities import TestModel, Utils +class TestModelWithExperts(torch.nn.Module): + """Model with both dense and expert-parallel parameters. + + Dense layers have the default allreduce=True. Expert layers have + allreduce=False on their parameters, which routes them to a separate + buffer with a different data-parallel group. + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + num_dense_layers: int, + num_expert_layers: int, + bias: bool, + ): + super().__init__() + self.dense_layers = torch.nn.ModuleList( + [torch.nn.Linear(input_dim, output_dim, bias) for _ in range(num_dense_layers)] + ) + self.expert_layers = torch.nn.ModuleList( + [torch.nn.Linear(input_dim, output_dim, bias) for _ in range(num_expert_layers)] + ) + for layer in self.expert_layers: + for param in layer.parameters(): + param.allreduce = False + + def get_model_and_buffers( input_dim: int, output_dim: int, @@ -49,8 +78,18 @@ def get_model_and_buffers( # Wrap with DistributedDataParallel, and get underlying buffer. # Use dummy TransformerConfig with mostly default values. Avoid divide-by-zero # errors for num_attention_heads and num_layers. + # Pre-compute parameter layouts for the distributed optimizer. + full_param_layout = None + if use_distributed_optimizer: + all_params = [p for p in model.parameters() if p.requires_grad] + full_param_layout = DistributedOptimizer.compute_full_param_layout( + all_params, bucket_size, parallel_state.get_data_parallel_world_size(), ddp_config + ) model = DistributedDataParallel( - TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config=ddp_config, module=model + TransformerConfig(num_attention_heads=1, num_layers=1), + ddp_config=ddp_config, + module=model, + full_param_layout=full_param_layout, ) assert len(model.buffers) == 1 param_and_grad_buffer = model.buffers[0] @@ -733,6 +772,14 @@ def mock_packed_shape(shape): average_in_collective=False, ) + # Pre-compute layout for distributed optimizer (with padding); + # otherwise use default (no padding). + param_layout = None + if use_distributed_optimizer: + param_layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size, dp_world_size, ddp_config, list(range(len(params))) + ) + with ( mock.patch( 'megatron.core.distributed.param_and_grad_buffer.is_nvfp4tensor', @@ -760,6 +807,7 @@ def mock_packed_shape(shape): param_indices=list(range(len(params))), nccl_ub=False, pg_collection=mock_pg, + param_layout=param_layout, ) return buffer, params @@ -884,3 +932,101 @@ def test_nvfp4_varied_param_sizes(self): assert buffer.param_index_map[params[1]] == (large_unpacked_start, large_unpacked_end, 0) assert buffer.param_index_map[params[0]] == (small_unpacked_start, small_unpacked_end, 0) + + +@pytest.mark.parametrize("use_distributed_optimizer", [False, True]) +def test_expert_parallel_params_get_separate_buffers(use_distributed_optimizer: bool): + """Verify that expert-parallel params (allreduce=False) land in separate buffers + with correctly scoped layouts and independent param_index_maps.""" + Utils.initialize_model_parallel() + + input_dim = 95 + output_dim = 95 + num_dense_layers = 3 + num_expert_layers = 2 + bucket_size = None # Single bucket per buffer. + + ddp_config = DistributedDataParallelConfig( + grad_reduce_in_fp32=True, + use_distributed_optimizer=use_distributed_optimizer, + overlap_grad_reduce=True, + bucket_size=bucket_size, + average_in_collective=False, + ) + model = TestModelWithExperts( + input_dim=input_dim, + output_dim=output_dim, + num_dense_layers=num_dense_layers, + num_expert_layers=num_expert_layers, + bias=True, + ).bfloat16() + + full_param_layout = None + if use_distributed_optimizer: + all_params = [p for p in model.parameters() if p.requires_grad] + full_param_layout = DistributedOptimizer.compute_full_param_layout( + all_params, bucket_size, parallel_state.get_data_parallel_world_size(), ddp_config + ) + + ddp_model = DistributedDataParallel( + TransformerConfig(num_attention_heads=1, num_layers=1), + ddp_config=ddp_config, + module=model, + full_param_layout=full_param_layout, + ) + + # Should have exactly one dense buffer and one expert buffer. + assert len(ddp_model.buffers) == 1, f"Expected 1 dense buffer, got {len(ddp_model.buffers)}" + assert ( + len(ddp_model.expert_parallel_buffers) == 1 + ), f"Expected 1 expert buffer, got {len(ddp_model.expert_parallel_buffers)}" + + dense_buffer = ddp_model.buffers[0] + expert_buffer = ddp_model.expert_parallel_buffers[0] + + # Collect expected params for each buffer. + expected_dense_params = set() + expected_expert_params = set() + for param in model.parameters(): + if not param.requires_grad: + continue + if getattr(param, 'allreduce', True): + expected_dense_params.add(param) + else: + expected_expert_params.add(param) + + # Verify each buffer contains exactly the right params. + dense_buffer_params = set() + for bucket in dense_buffer.buckets: + dense_buffer_params.update(bucket.params) + assert ( + dense_buffer_params == expected_dense_params + ), "Dense buffer should contain exactly the dense params" + + expert_buffer_params = set() + for bucket in expert_buffer.buckets: + expert_buffer_params.update(bucket.params) + assert ( + expert_buffer_params == expected_expert_params + ), "Expert buffer should contain exactly the expert-parallel params" + + # Verify param_index_maps are scoped to their own buffer (no cross-contamination). + assert set(dense_buffer.param_index_map.keys()) == expected_dense_params + assert set(expert_buffer.param_index_map.keys()) == expected_expert_params + + # Verify both buffers have indices starting from 0 (independent index spaces). + dense_starts = [s for s, _, _ in dense_buffer.param_index_map.values()] + expert_starts = [s for s, _, _ in expert_buffer.param_index_map.values()] + assert min(dense_starts) == 0, "Dense buffer indices should start at 0" + assert min(expert_starts) == 0, "Expert buffer indices should start at 0" + + # Verify DP divisibility for distributed optimizer. + if use_distributed_optimizer: + dp_world_size = parallel_state.get_data_parallel_world_size() + for buffer_name, buffer in [("dense", dense_buffer), ("expert", expert_buffer)]: + assert buffer.numel % dp_world_size == 0, ( + f"{buffer_name} buffer numel ({buffer.numel}) should be " + f"divisible by dp_world_size ({dp_world_size})" + ) + + Utils.destroy_model_parallel() diff --git a/tests/unit_tests/distributed/test_param_layout.py b/tests/unit_tests/distributed/test_param_layout.py new file mode 100644 index 00000000000..99922e37970 --- /dev/null +++ b/tests/unit_tests/distributed/test_param_layout.py @@ -0,0 +1,478 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Tests for parameter layout computation functions. + +These tests verify the pure-computation layout functions without requiring +GPU or distributed setup: +- pad_to_divisor, pad_param_start, pad_bucket_end (shared padding utilities) +- group_params_for_buffers (parameter grouping by dtype/expert) +- _compute_default_per_buffer_param_layout (no-padding layout) +- DistributedOptimizer._compute_per_buffer_param_layout (padded layout) +- DistributedOptimizer.compute_full_param_layout (end-to-end layout) +""" + +import math +from unittest import mock + +import pytest +import torch + +from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.distributed.param_and_grad_buffer import ( + _compute_default_per_buffer_param_layout, + group_params_for_buffers, +) +from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer +from megatron.core.optimizer.param_layout import ( + BufferKey, + pad_bucket_end, + pad_param_start, + pad_to_divisor, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_params(*shapes, dtype=torch.bfloat16): + """Create a list of nn.Parameters with the given shapes.""" + return [torch.nn.Parameter(torch.randn(s, dtype=dtype)) for s in shapes] + + +def _make_param_with_attrs(shape, dtype=torch.bfloat16, **attrs): + """Create an nn.Parameter with extra attributes (e.g. allreduce, shared_embedding).""" + param = torch.nn.Parameter(torch.randn(shape, dtype=dtype)) + for k, v in attrs.items(): + setattr(param, k, v) + return param + + +# --------------------------------------------------------------------------- +# Tests for shared padding utilities +# --------------------------------------------------------------------------- + + +class TestPaddingUtilities: + + def test_pad_to_divisor_exact_multiple(self): + assert pad_to_divisor(128, 64) == 128 + + def test_pad_to_divisor_rounds_up(self): + assert pad_to_divisor(65, 64) == 128 + + def test_pad_to_divisor_zero(self): + assert pad_to_divisor(0, 64) == 0 + + def test_pad_param_start(self): + assert pad_param_start(0) == 0 + assert pad_param_start(1) == 64 + assert pad_param_start(63) == 64 + assert pad_param_start(64) == 64 + assert pad_param_start(65) == 128 + + def test_pad_bucket_end_basic(self): + dp_size = 4 + divisor = math.lcm(dp_size, 128) + result = pad_bucket_end(1, dp_size, pad_for_high_nccl_busbw=False) + assert result == divisor + assert result % dp_size == 0 + assert result % 128 == 0 + + def test_pad_bucket_end_high_busbw(self): + dp_size = 4 + divisor = math.lcm(dp_size, 128, 2**16) + result = pad_bucket_end(1, dp_size, pad_for_high_nccl_busbw=True) + assert result == divisor + assert result % (2**16) == 0 + + def test_pad_bucket_end_already_aligned(self): + dp_size = 2 + divisor = math.lcm(dp_size, 128) + result = pad_bucket_end(divisor, dp_size, pad_for_high_nccl_busbw=False) + assert result == divisor + + +# --------------------------------------------------------------------------- +# Tests for group_params_for_buffers +# --------------------------------------------------------------------------- + + +class TestGroupParamsForBuffers: + + def test_single_dtype_no_experts(self): + """All bf16 params with no expert-parallel should go in one group.""" + params = _make_params((100, 100), (50, 50)) + result = group_params_for_buffers(params, grad_reduce_in_fp32=True) + + assert len(result) == 1 + key = list(result.keys())[0] + assert key == BufferKey(torch.bfloat16, torch.float, False) + group_params, indices = result[key] + assert group_params == params + assert indices == [0, 1] + + def test_grad_reduce_not_fp32(self): + """When grad_reduce_in_fp32=False, grad_dtype matches param dtype.""" + params = _make_params((100,)) + result = group_params_for_buffers(params, grad_reduce_in_fp32=False) + + key = list(result.keys())[0] + assert key.grad_dtype == torch.bfloat16 + + def test_expert_parallel_separation(self): + """Params with allreduce=False should be in a separate group.""" + dense = _make_param_with_attrs((100,)) + expert = _make_param_with_attrs((100,), allreduce=False) + result = group_params_for_buffers([dense, expert], grad_reduce_in_fp32=True) + + assert len(result) == 2 + dense_key = BufferKey(torch.bfloat16, torch.float, False) + expert_key = BufferKey(torch.bfloat16, torch.float, True) + assert dense_key in result + assert expert_key in result + assert result[dense_key][0] == [dense] + assert result[expert_key][0] == [expert] + + def test_param_indices_independent_per_group(self): + """Expert and dense groups should have independent param_indices starting at 0.""" + dense_params = _make_params((100,), (200,)) + expert = _make_param_with_attrs((100,), allreduce=False) + result = group_params_for_buffers( + [dense_params[0], expert, dense_params[1]], grad_reduce_in_fp32=True + ) + + dense_key = BufferKey(torch.bfloat16, torch.float, False) + expert_key = BufferKey(torch.bfloat16, torch.float, True) + _, dense_indices = result[dense_key] + _, expert_indices = result[expert_key] + assert dense_indices == [0, 1] + assert expert_indices == [0] + + def test_mixed_dtypes(self): + """Params with different dtypes go in separate groups.""" + bf16_param = _make_params((100,), dtype=torch.bfloat16)[0] + fp32_param = _make_params((100,), dtype=torch.float32)[0] + result = group_params_for_buffers([bf16_param, fp32_param], grad_reduce_in_fp32=True) + + assert len(result) == 2 + bf16_key = BufferKey(torch.bfloat16, torch.float, False) + fp32_key = BufferKey(torch.float32, torch.float, False) + assert bf16_key in result + assert fp32_key in result + + +# --------------------------------------------------------------------------- +# Tests for _compute_default_per_buffer_param_layout +# --------------------------------------------------------------------------- + + +class TestDefaultParamLayout: + + def test_single_bucket_no_padding(self): + """With bucket_size=None, all params go in one bucket with no padding.""" + params = _make_params((100, 100), (50, 50)) + layout = _compute_default_per_buffer_param_layout(params, bucket_size=None) + + # Params iterated in reverse: params[1] first (2500 elems), params[0] second (10000). + assert layout.param_index_map[params[1]] == (0, 2500, 0) + assert layout.param_index_map[params[0]] == (2500, 12500, 0) + assert layout.bucket_indices == [(0, 12500)] + assert layout.per_bucket_numel_unpadded == [12500] + + def test_multiple_buckets(self): + """Params should split into buckets when exceeding bucket_size.""" + params = _make_params((100, 100), (100, 100), (100, 100)) + layout = _compute_default_per_buffer_param_layout(params, bucket_size=15000) + + # Reverse order: params[2], params[1], params[0]. Each is 10000 elems. + # After params[2]: 10000 < 15000, continue. + # After params[1]: 20000 >= 15000, finalize bucket. + # After params[0]: 10000 < 15000, finalize at end. + assert len(layout.bucket_indices) == 2 + assert layout.per_bucket_numel_unpadded[0] == 20000 + assert layout.per_bucket_numel_unpadded[1] == 10000 + + def test_no_padding_applied(self): + """Default layout should never add padding.""" + params = _make_params((97,), (103,)) + layout = _compute_default_per_buffer_param_layout(params, bucket_size=None) + + total_numel = 97 + 103 + assert layout.bucket_indices == [(0, total_numel)] + assert layout.per_bucket_numel_unpadded == [total_numel] + + def test_single_param(self): + params = _make_params((256,)) + layout = _compute_default_per_buffer_param_layout(params, bucket_size=None) + + assert layout.param_index_map[params[0]] == (0, 256, 0) + assert layout.bucket_indices == [(0, 256)] + + +# --------------------------------------------------------------------------- +# Tests for DistributedOptimizer._compute_per_buffer_param_layout +# --------------------------------------------------------------------------- + + +class TestDistOptParamLayout: + + @staticmethod + def _make_ddp_config(**overrides): + defaults = dict( + grad_reduce_in_fp32=False, + use_distributed_optimizer=True, + overlap_grad_reduce=True, + bucket_size=None, + average_in_collective=False, + ) + defaults.update(overrides) + return DistributedDataParallelConfig(**defaults) + + def test_param_start_64_alignment(self): + """Each param's start index should be 64-aligned.""" + # 97 is not a multiple of 64, so second param must be padded. + params = _make_params((97,), (103,)) + ddp_config = self._make_ddp_config() + layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=None, data_parallel_world_size=2, ddp_config=ddp_config + ) + + for param in params: + start, end, _ = layout.param_index_map[param] + assert start % 64 == 0, f"Start {start} should be 64-aligned" + assert end - start == param.numel() + + def test_bucket_end_dp_divisible(self): + """Each bucket end should be divisible by lcm(dp_size, 128).""" + params = _make_params((1000,), (1000,)) + dp_size = 4 + ddp_config = self._make_ddp_config() + layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=None, data_parallel_world_size=dp_size, ddp_config=ddp_config + ) + + divisor = math.lcm(dp_size, 128) + for start, end in layout.bucket_indices: + assert end % divisor == 0, f"Bucket end {end} should be divisible by {divisor}" + + def test_bucket_end_high_busbw_padding(self): + """With pad_buckets_for_high_nccl_busbw, bucket end should be divisible by 2^16.""" + params = _make_params((1000,)) + dp_size = 2 + ddp_config = self._make_ddp_config(pad_buckets_for_high_nccl_busbw=True) + layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=None, data_parallel_world_size=dp_size, ddp_config=ddp_config + ) + + divisor = math.lcm(dp_size, 128, 2**16) + for _, end in layout.bucket_indices: + assert end % divisor == 0 + + def test_shared_embedding_gets_separate_bucket(self): + """Params with shared_embedding=True should be placed in their own bucket.""" + regular = _make_param_with_attrs((1000,)) + shared = _make_param_with_attrs((1000,), shared_embedding=True) + # Reverse order: shared first (since it's last in list), then regular. + params = [regular, shared] + ddp_config = self._make_ddp_config() + layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=None, data_parallel_world_size=2, ddp_config=ddp_config + ) + + # shared_embedding param should be in its own bucket. + _, _, shared_bucket = layout.param_index_map[shared] + _, _, regular_bucket = layout.param_index_map[regular] + assert shared_bucket != regular_bucket + + def test_shared_embedding_as_first_reversed_param_no_extra_bucket(self): + """If shared_embedding param is the first in reversed order (last in list), + it should not create an empty extra bucket before it.""" + shared = _make_param_with_attrs((1000,), shared_embedding=True) + regular = _make_param_with_attrs((1000,)) + # Reverse order: regular first, then shared. + params = [shared, regular] + ddp_config = self._make_ddp_config() + layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=None, data_parallel_world_size=2, ddp_config=ddp_config + ) + + # Both params should be in their own buckets (shared splits after regular). + assert len(layout.bucket_indices) == 2 + + def test_multiple_buckets_with_bucket_size(self): + """Verify bucket splitting with an explicit bucket_size.""" + params = _make_params((5000,), (5000,), (5000,)) + dp_size = 2 + ddp_config = self._make_ddp_config() + layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=8000, data_parallel_world_size=dp_size, ddp_config=ddp_config + ) + + # Each param is 5000 elems. With 64-alignment: + # Bucket 0: params[2] starts at 0, ends at 5000 (< 8000); params[1] starts at 5000 + # (already 64-aligned since 5000 rounds to 5056), ends at 10056 >= 8000 → finalize. + # Bucket 1: params[0]. + assert len(layout.bucket_indices) == 2 + assert len(layout.per_bucket_numel_unpadded) == 2 + + def test_numel_unpadded_vs_padded(self): + """per_bucket_numel_unpadded should be <= padded bucket size.""" + params = _make_params((1000,)) + dp_size = 8 + ddp_config = self._make_ddp_config() + layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=None, data_parallel_world_size=dp_size, ddp_config=ddp_config + ) + + for i, (start, end) in enumerate(layout.bucket_indices): + padded_numel = end - start + assert layout.per_bucket_numel_unpadded[i] <= padded_numel + + def test_shared_embedding_with_bucket_size(self): + """Shared embedding that hits bucket_size threshold should still get its own bucket.""" + # 3 regular params + 1 shared embedding, each 5000 elements. + # With bucket_size=8000, regular params would form 2-param buckets, + # but shared embedding must always be isolated. + regulars = [_make_param_with_attrs((5000,)) for _ in range(3)] + shared = _make_param_with_attrs((5000,), shared_embedding=True) + # Order: regulars first, shared last → reversed: shared first. + params = regulars + [shared] + dp_size = 2 + ddp_config = self._make_ddp_config() + layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=8000, data_parallel_world_size=dp_size, ddp_config=ddp_config + ) + + # Shared embedding must be alone in its bucket. + _, _, shared_bucket = layout.param_index_map[shared] + shared_bucket_params = [ + p for p, (_, _, bid) in layout.param_index_map.items() if bid == shared_bucket + ] + assert len(shared_bucket_params) == 1 + assert shared_bucket_params[0] is shared + + def test_layout_matches_default_when_no_padding_needed(self): + """When params are 64-aligned and bucket_size=None, distributed optimizer layout + should produce the same param_index_map ordering as the default layout + (only bucket end padding differs).""" + # Use param sizes that are multiples of 64 to avoid start-of-param padding. + params = _make_params((64 * 10,), (64 * 20,), (64 * 15,)) + ddp_config = self._make_ddp_config() + dist_layout = DistributedOptimizer._compute_per_buffer_param_layout( + params, bucket_size=None, data_parallel_world_size=1, ddp_config=ddp_config + ) + default_layout = _compute_default_per_buffer_param_layout(params, bucket_size=None) + + # With dp_size=1, lcm(1, 128) = 128, so bucket end padding may differ. + # But param ordering and unpadded numel should match. + for param in params: + dist_start, dist_end, dist_bid = dist_layout.param_index_map[param] + def_start, def_end, def_bid = default_layout.param_index_map[param] + assert dist_start == def_start, f"Start mismatch for param: {dist_start} vs {def_start}" + assert dist_end == def_end, f"End mismatch for param: {dist_end} vs {def_end}" + assert dist_bid == def_bid, f"Bucket ID mismatch for param: {dist_bid} vs {def_bid}" + assert dist_layout.per_bucket_numel_unpadded == default_layout.per_bucket_numel_unpadded + + +# --------------------------------------------------------------------------- +# Tests for DistributedOptimizer.compute_full_param_layout +# --------------------------------------------------------------------------- + + +class TestComputeFullParamLayout: + + @staticmethod + def _make_ddp_config(**overrides): + defaults = dict( + grad_reduce_in_fp32=False, + use_distributed_optimizer=True, + overlap_grad_reduce=True, + bucket_size=None, + average_in_collective=False, + ) + defaults.update(overrides) + return DistributedDataParallelConfig(**defaults) + + def test_dense_only(self): + """With only dense params, should produce a single layout.""" + params = _make_params((100, 100), (50, 50)) + ddp_config = self._make_ddp_config() + full_layout = DistributedOptimizer.compute_full_param_layout( + params, bucket_size=None, data_parallel_world_size=2, ddp_config=ddp_config + ) + + assert len(full_layout.layouts) == 1 + key = list(full_layout.layouts.keys())[0] + assert key.is_expert_parallel is False + layout = full_layout.layouts[key] + assert set(layout.param_index_map.keys()) == set(params) + + def test_dense_and_expert_separate_layouts(self): + """Dense and expert-parallel params should get independent layouts.""" + dense = _make_param_with_attrs((100, 100)) + expert = _make_param_with_attrs((100, 100), allreduce=False) + ddp_config = self._make_ddp_config() + full_layout = DistributedOptimizer.compute_full_param_layout( + [dense, expert], bucket_size=None, data_parallel_world_size=2, ddp_config=ddp_config + ) + + assert len(full_layout.layouts) == 2 + dense_key = BufferKey(torch.bfloat16, torch.bfloat16, False) + expert_key = BufferKey(torch.bfloat16, torch.bfloat16, True) + assert dense_key in full_layout.layouts + assert expert_key in full_layout.layouts + + # Each layout should only contain its own params. + assert set(full_layout.layouts[dense_key].param_index_map.keys()) == {dense} + assert set(full_layout.layouts[expert_key].param_index_map.keys()) == {expert} + + # Both should start at index 0 (independent index spaces). + dense_starts = [s for s, _, _ in full_layout.layouts[dense_key].param_index_map.values()] + expert_starts = [s for s, _, _ in full_layout.layouts[expert_key].param_index_map.values()] + assert min(dense_starts) == 0 + assert min(expert_starts) == 0 + + def test_expert_uses_expert_dp_world_size(self): + """Expert-parallel layout should use expert_data_parallel_world_size for padding.""" + dense = _make_param_with_attrs((1000,)) + expert = _make_param_with_attrs((1000,), allreduce=False) + ddp_config = self._make_ddp_config() + + # Dense dp_size=3, expert dp_size=256. + # lcm(3, 128) = 384, lcm(256, 128) = 256 — different divisors. + full_layout = DistributedOptimizer.compute_full_param_layout( + [dense, expert], + bucket_size=None, + data_parallel_world_size=3, + ddp_config=ddp_config, + expert_data_parallel_world_size=256, + ) + + dense_key = BufferKey(torch.bfloat16, torch.bfloat16, False) + expert_key = BufferKey(torch.bfloat16, torch.bfloat16, True) + + # Expert bucket end should be divisible by lcm(256, 128) = 256. + expert_divisor = math.lcm(256, 128) + assert expert_divisor == 256 + for _, end in full_layout.layouts[expert_key].bucket_indices: + assert end % expert_divisor == 0 + + # Dense bucket end should be divisible by lcm(3, 128) = 384. + dense_divisor = math.lcm(3, 128) + assert dense_divisor == 384 + for _, end in full_layout.layouts[dense_key].bucket_indices: + assert end % dense_divisor == 0 + + def test_param_indices_populated(self): + """compute_full_param_layout should populate param_indices on each layout.""" + params = _make_params((100,), (200,), (300,)) + ddp_config = self._make_ddp_config() + full_layout = DistributedOptimizer.compute_full_param_layout( + params, bucket_size=None, data_parallel_world_size=2, ddp_config=ddp_config + ) + + layout = list(full_layout.layouts.values())[0] + assert len(layout.param_indices) == 3 + assert min(layout.param_indices) == 0 + assert max(layout.param_indices) == 2