From 07e17d347608e16085910e9e7b54079f4ab84a8c Mon Sep 17 00:00:00 2001 From: jianbinc Date: Wed, 25 Feb 2026 15:49:06 +0800 Subject: [PATCH 1/2] Fix edge case in `build_data_parallel_buffer_index` and improve clarity - Resolve a bug in handling edge cases with a large number of fragments. - Refactor `build_data_parallel_buffer_index` for clearer structure and better maintainability, adding inline documentation. - Add comprehensive unit tests for parameter splitting and buffer index construction to verify correctness. --- .../megatron_fsdp/param_and_grad_buffer.py | 263 ++++++++++------ .../test_mfsdp_param_and_grad_buffer.py | 280 ++++++++++++++++++ 2 files changed, 455 insertions(+), 88 deletions(-) create mode 100644 tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py index 38e904ca947..aaa0d114fc3 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py @@ -263,130 +263,217 @@ def _pad(number_to_be_padded: int, divisor: int) -> int: def build_data_parallel_buffer_index( - elements: List[torch.Size], + param_shapes: List[torch.Size], data_parallel_rank: int, data_parallel_world_size: int, is_data_distributed: bool, ddp_config: DistributedDataParallelConfig, bucket_id: int = 0, chunk_size_factor: int = 1, -) -> Tuple[List[tuple], BucketIndex, ShardBucketIndex]: +) -> Tuple[Dict[int, TensorItemIndex], BucketIndex, ShardBucketIndex]: """ - Assuming that all input tensor elements contiguously compose a global - buffer, give the index range of every tensor, the bucket in the buffer, - and the (distributed) shard within the bucket. Note that the global bucket - buffer is only temporarily allocated, but is abstractly tracked via indices - deduced from the number of raw parameters assigned to this buffer / bucket. + Build indexing metadata for parameters packed into a single data-parallel + buffer, including bucket and shard indices. + + All tensors in `param_shapes` are assumed to be laid out contiguously in a + global buffer for this bucket. The function computes: + - where each parameter lives in that buffer (global offset and size), + - bucket-level metadata, + - and, if sharding is enabled, per-rank shard metadata. + + `chunk_size_factor` is used to ensure that splitting only happens along the + first tensor dimension: the product of all remaining dimensions is treated + as the minimal indivisible unit when packing. In other words, the buffer + is aligned in multiples of `chunk_size_factor`, and leftover space may be + filled with smaller “fragment” tensors so that total padding is minimized. Args: - elements (List[torch.Size]): List of input tensor. - data_parallel_rank (int): Rank of the current process in the data parallel group. - data_parallel_world_size (int): World size of the data parallel group. - bucket_id (int, optional): The id of the bucket. Defaults to 0. + param_shapes: + List of parameter shapes to be packed into the global buffer. + data_parallel_rank: + Rank of the current process in the data-parallel group. + data_parallel_world_size: + World size of the data-parallel group. + is_data_distributed: + Whether parameters/gradients are distributed across ranks. + ddp_config: + Configuration object controlling DDP and sharding behavior. + bucket_id: + Identifier of the bucket to which these parameters belong. + chunk_size_factor: + Alignment unit along the flattened buffer. The first dimension is + the only allowed split dimension; the product of remaining + dimensions must be a multiple of this factor to be split. Returns: - Tuple[Dict[int, TensorItemIndex], BucketIndex, ShardBucketIndex]: The index - range of every tensor, every bucket and every in bucket local buffer. + item_index_map: + A mapping from parameter id to its buffer index metadata. + bucket_index: + Index metadata for the full (possibly padded) bucket buffer. + shard_bucket_index: + Sharded bucket metadata for the current data-parallel rank. """ - def _pad_if_needed(data_index: int) -> int: + def _pad_if_needed(global_index: int) -> int: + """Optionally pad the buffer length for sharding alignment.""" if ddp_config.data_parallel_sharding_strategy != "no_shard": - return _pad(data_index, data_parallel_world_size * chunk_size_factor) - return data_index - - def add_item(item_id, item, offset, item_index_map): - # The item index map contains information on where each parameter item will - # be stored in the tensor data buffer in a bucket. - item_index_map[item_id] = TensorItemIndex( - # Global data index of the starting idx of this parameter - # = running global data index + updated bucket size - the parameter size. + return _pad(global_index, data_parallel_world_size * chunk_size_factor) + return global_index + + def add_item( + param_id: int, shape: torch.Size, offset: int, item_index_map: Dict[int, TensorItemIndex] + ) -> None: + """Register a single parameter’s index information in the global buffer.""" + item_index_map[param_id] = TensorItemIndex( + # Global flat buffer offset of this parameter. global_data_index=offset, - # Number of tensor elements in the parameter. - size=item.numel(), - # Index of the parameter to be buffered in the list of parameter shapes. - item_id=item_id, - # ID of the bucket that this parameter belongs to. + # Number of elements in the parameter tensor. + size=shape.numel(), + # Index of this parameter in the original `param_shapes` list. + item_id=param_id, + # Bucket this parameter belongs to. bucket_id=bucket_id, - # Shape of the parameter. - shape=item, + # Original tensor shape. + shape=shape, ) - fragment_items = [] - regular_items = [] - for item_id, item in enumerate(elements): - if item.numel() < chunk_size_factor: - fragment_items.append((item_id, item)) + # Separate “regular” items (at least one full alignment unit) from small fragments. + fragment_params: List[Tuple[int, torch.Size]] = [] + regular_params: List[Tuple[int, torch.Size]] = [] + + for param_id, shape in enumerate(param_shapes): + if shape.numel() < chunk_size_factor: + fragment_params.append((param_id, shape)) else: - item[1:].numel() - regular_items.append((item_id, item)) - - # Sort the fragments so that items with larger sizes come first. - # When filling the remaining space, prioritize placing the larger fragments first. - sorted(fragment_items, key=lambda id_item: -id_item[1].numel()) - - # For all bucket parameters, add information on the parameter to the item index map, - # and add the size of the parameter to the bucket. - item_index_map = {} - data_index = 0 - while len(regular_items) > 0: - item_id, item = regular_items.pop(0) - add_item(item_id, item, data_index, item_index_map) - if item.numel() % chunk_size_factor == 0: - data_index += item.numel() + regular_params.append((param_id, shape)) + + # Sort fragments so larger ones are placed first when filling gaps. + fragment_params.sort(key=lambda pair: -pair[1].numel()) + + item_index_map: Dict[int, TensorItemIndex] = {} + global_data_index = 0 + + # First pass: place all regular parameters, trying to pack their remainders together. + while regular_params: + param_id, shape = regular_params.pop(0) + numel = shape.numel() + + # Place the main body of the parameter. + add_item(param_id, shape, global_data_index, item_index_map) + + # If perfectly aligned, just advance. + if numel % chunk_size_factor == 0: + global_data_index += numel continue - gap_offset = data_index + item.numel() - data_index += (item.numel() // chunk_size_factor + 1) * chunk_size_factor - remain = item.numel() % chunk_size_factor - space = chunk_size_factor - remain - found_rhs = False - for id_rhs in regular_items[:]: - rhs_id, rhs = id_rhs - if rhs.numel() % chunk_size_factor == 0: + # Otherwise, this parameter creates a partial “grid” at the tail. + gap_offset = global_data_index + numel + full_aligned_size = (numel // chunk_size_factor + 1) * chunk_size_factor + global_data_index += full_aligned_size + + remainder = numel % chunk_size_factor + remaining_space = chunk_size_factor - remainder + + # Try to pair another misaligned regular parameter to fill the same grid. + rhs_found = False + rhs_param_id = None + rhs_shape = None + + for candidate in list(regular_params): + cand_id, cand_shape = candidate + cand_numel = cand_shape.numel() + if cand_numel % chunk_size_factor == 0: + # Already aligned, skip for pairing. continue - rhs_remain = rhs.numel() % chunk_size_factor - if remain + rhs_remain <= chunk_size_factor: - found_rhs = True - regular_items.remove(id_rhs) + cand_remainder = cand_numel % chunk_size_factor + if remainder + cand_remainder <= chunk_size_factor: + rhs_found = True + rhs_param_id, rhs_shape = cand_id, cand_shape + regular_params.remove(candidate) break - # If a item is found to have remnants, then the remnants of the two - # items are placed in one "grid". - if found_rhs: - add_item(rhs_id, rhs, data_index - rhs_remain, item_index_map) - space -= rhs_remain - data_index += rhs.numel() // chunk_size_factor * chunk_size_factor - - # Try adding the fragments into the gaps - for id_frag in fragment_items[:]: - frag_id, frag = id_frag - if frag.numel() > space: + # If we find a partner, place its remainder into the same grid. + if rhs_found and rhs_param_id is not None and rhs_shape is not None: + rhs_numel = rhs_shape.numel() + rhs_remainder = rhs_numel % chunk_size_factor + + # Place the full parameter; its remainder lands at the end of this grid. + add_item(rhs_param_id, rhs_shape, global_data_index - rhs_numel, item_index_map) + + remaining_space -= rhs_remainder + # Advance only by the aligned part; the remainder is in the current grid. + global_data_index += (rhs_numel // chunk_size_factor) * chunk_size_factor + + # Fill any remaining space in this grid with fragment parameters. + for frag in list(fragment_params): + frag_id, frag_shape = frag + frag_numel = frag_shape.numel() + if frag_numel > remaining_space: continue - add_item(frag_id, frag, gap_offset, item_index_map) - space -= frag.numel() - gap_offset += frag.numel() - fragment_items.remove(id_frag) + add_item(frag_id, frag_shape, gap_offset, item_index_map) + remaining_space -= frag_numel + gap_offset += frag_numel + fragment_params.remove(frag) + + # Helper to bin-pack leftover fragments into grids of size `chunk_size_factor`. + def pack_fragments( + fragments: List[Tuple[int, torch.Size]], capacity: int + ) -> List[List[Tuple[int, torch.Size]]]: + """ + Pack remaining fragment parameters into fixed-capacity slots. + + Each slot corresponds to one alignment grid of size `capacity` + (equal to `chunk_size_factor`). A slot contains a list of + (param_id, shape) pairs whose total numel does not exceed `capacity`. + """ + # Largest fragments first improves packing efficiency. + sorted_frags = sorted(fragments, key=lambda pair: -pair[1].numel()) + + slots: List[List[Tuple[int, torch.Size]]] = [] + + for param in sorted_frags: + param_id, shape = param + param_size = shape.numel() + placed = False + + for slot in slots: + used = sum(p[1].numel() for p in slot) + if used + param_size <= capacity: + slot.append(param) + placed = True + break + + if not placed: + slots.append([param]) + + return slots - for frag_id, frag in fragment_items: - add_item(frag_id, frag, data_index, item_index_map) - data_index += frag.numel() + # Second pass: any fragments that were not used to fill gaps get their own grids. + if fragment_params: + fragment_slots = pack_fragments(fragment_params, chunk_size_factor) + for slot in fragment_slots: + offset_within_grid = 0 + for param_id, shape in slot: + add_item(param_id, shape, global_data_index + offset_within_grid, item_index_map) + offset_within_grid += shape.numel() + global_data_index += chunk_size_factor - # Bucket index contains information on what tensor items are in this bucket. + # Build bucket-level index. bucket_index = BucketIndex( bucket_id=bucket_id, global_data_index=0, - size=_pad_if_needed(data_index), + size=_pad_if_needed(global_data_index), items=list(item_index_map.values()), ) - # Sharded bucket index contains local bucket shard information. + # Build sharded bucket index for this DP rank. shard_bucket_index = _get_dp_buffer_shard_bucket_index( - bucket_index, is_data_distributed, data_parallel_world_size, data_parallel_rank + bucket_index=bucket_index, + is_data_distributed=is_data_distributed, + data_parallel_world_size=data_parallel_world_size, + data_parallel_rank=data_parallel_rank, ) - # Return the tensor item index map in the buffer, - # the bucket index with information on what items this bucket contains, - # and the sharded bucket index. return item_index_map, bucket_index, shard_bucket_index diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py b/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py new file mode 100644 index 00000000000..189baf115b0 --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py @@ -0,0 +1,280 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +from unittest.mock import Mock + +import pytest +import torch +import torch.nn as nn +from torch.testing._internal.distributed.fake_pg import FakeStore + +from megatron.core.distributed.fsdp.src.megatron_fsdp.param_and_grad_buffer import ( + BucketingPolicy, + DataParallelBuffer, + _get_parameter_groups, +) + +FACTORY_META = {"device": "meta", "dtype": torch.bfloat16} + + +class AllInOneDummyFSDPModel(nn.Module): + """ + One dummy module that recreates all parameter shapes and counts + from the original paligemma_with_expert + heads, on meta device. + Forward is not implemented; this is only for FSDP sharding tests. + """ + + def __init__(self): + super().__init__() + + # + # Language model parts + # + # embed_tokens.weight (257152, 2048) + self.lm_embed_tokens = nn.Embedding( + num_embeddings=257152, embedding_dim=2048, **FACTORY_META + ) + + # 18 language layers: self-attn + MLP + norms + self.lm_layers = nn.ModuleList([LanguageLayerFlat() for _ in range(18)]) + + # final norm.weight (2048,) + self.lm_final_norm = nn.LayerNorm( + normalized_shape=2048, device="meta", dtype=torch.bfloat16 + ) + + # + # Vision tower parts + # + # patch_embedding.weight (1152, 3, 14, 14) + self.vision_patch_embedding = nn.Conv2d( + in_channels=3, + out_channels=1152, + kernel_size=14, + stride=14, + device="meta", + dtype=torch.bfloat16, + ) + # position_embedding.weight (256, 1152) + self.vision_position_embedding = nn.Embedding( + num_embeddings=256, embedding_dim=1152, **FACTORY_META + ) + + # 27 encoder layers: self-attn + MLP + norms + self.vision_layers = nn.ModuleList([VisionLayerFlat() for _ in range(27)]) + + # post_layer_norm.weight/bias (1152,) + self.vision_post_layernorm = nn.LayerNorm( + normalized_shape=1152, device="meta", dtype=torch.bfloat16 + ) + + # + # Multi-modal projector + # + # linear.weight (2048, 1152) + self.mm_projector = nn.Linear(1152, 2048, **FACTORY_META) + + # + # Gemma expert LM + # + # 18 expert layers + self.gemma_layers = nn.ModuleList([GemmaLayerFlat() for _ in range(18)]) + # norm.dense (3072, 1024) + bias (3072,) + self.gemma_norm_dense = nn.Linear(1024, 3072, **FACTORY_META) + # lm_head.weight (257152, 1024) + self.gemma_lm_head = nn.Linear(1024, 257152, bias=False, **FACTORY_META) + + # + # Action head + time MLP + # + # action_in_proj.weight (1024, 32), bias (1024,) + self.action_in_proj = nn.Linear(32, 1024, **FACTORY_META) + # action_out_proj.weight (32, 1024), bias (32,) + self.action_out_proj = nn.Linear(1024, 32, **FACTORY_META) + + # time_mlp_in.weight (1024, 1024), bias (1024,) + self.time_mlp_in = nn.Linear(1024, 1024, **FACTORY_META) + # time_mlp_out.weight (1024, 1024), bias (1024,) + self.time_mlp_out = nn.Linear(1024, 1024, **FACTORY_META) + + def forward(self, *args, **kwargs): + raise NotImplementedError("Dummy model for FSDP tests only.") + + +class LanguageLayerFlat(nn.Module): + """ + One language layer matching: + - self_attn.{q,k,v,o}_proj + - mlp.{gate,up,down}_proj + - input_layernorm, post_attention_layernorm (2048,) + """ + + def __init__(self): + super().__init__() + + # Self-attention projections + # q_proj.weight (2048, 2048) + self.self_attn_q_proj = nn.Linear(2048, 2048, **FACTORY_META) + # k_proj.weight (256, 2048) + self.self_attn_k_proj = nn.Linear(2048, 256, **FACTORY_META) + # v_proj.weight (256, 2048) + self.self_attn_v_proj = nn.Linear(2048, 256, **FACTORY_META) + # o_proj.weight (2048, 2048) + self.self_attn_o_proj = nn.Linear(2048, 2048, **FACTORY_META) + + # MLP + # gate_proj.weight (16384, 2048) + self.mlp_gate_proj = nn.Linear(2048, 16384, **FACTORY_META) + # up_proj.weight (16384, 2048) + self.mlp_up_proj = nn.Linear(2048, 16384, **FACTORY_META) + # down_proj.weight (2048, 16384) + self.mlp_down_proj = nn.Linear(16384, 2048, **FACTORY_META) + + # input_layernorm.weight/bias (2048,) + self.input_layernorm = nn.LayerNorm( + normalized_shape=2048, device="meta", dtype=torch.bfloat16 + ) + # post_attention_layernorm.weight/bias (2048,) + self.post_attention_layernorm = nn.LayerNorm( + normalized_shape=2048, device="meta", dtype=torch.bfloat16 + ) + + +class VisionLayerFlat(nn.Module): + """ + One vision encoder layer matching: + - self_attn.{k,v,q,out}_proj (1152, 1152) + - mlp.fc1 (4304, 1152), mlp.fc2 (1152, 4304) + - layer_norm1, layer_norm2 (1152,) + """ + + def __init__(self): + super().__init__() + + # Self-attention projections (1152, 1152) + self.self_attn_k_proj = nn.Linear(1152, 1152, **FACTORY_META) + self.self_attn_v_proj = nn.Linear(1152, 1152, **FACTORY_META) + self.self_attn_q_proj = nn.Linear(1152, 1152, **FACTORY_META) + self.self_attn_out_proj = nn.Linear(1152, 1152, **FACTORY_META) + + # MLP + # fc1.weight (4304, 1152) + self.mlp_fc1 = nn.Linear(1152, 4304, **FACTORY_META) + # fc2.weight (1152, 4304) + self.mlp_fc2 = nn.Linear(4304, 1152, **FACTORY_META) + + # layer_norm1.weight/bias (1152,) + self.layer_norm1 = nn.LayerNorm(normalized_shape=1152, device="meta", dtype=torch.bfloat16) + # layer_norm2.weight/bias (1152,) + self.layer_norm2 = nn.LayerNorm(normalized_shape=1152, device="meta", dtype=torch.bfloat16) + + +class GemmaLayerFlat(nn.Module): + """ + One gemma_expert layer matching: + - self_attn.{q,k,v,o}_proj + - mlp.{gate,up,down}_proj + - input_layernorm.dense, post_attention_layernorm.dense (3072, 1024) + """ + + def __init__(self): + super().__init__() + + # Self-attention + # q_proj.weight (2048, 1024) + self.self_attn_q_proj = nn.Linear(1024, 2048, **FACTORY_META) + # k_proj.weight (256, 1024) + self.self_attn_k_proj = nn.Linear(1024, 256, **FACTORY_META) + # v_proj.weight (256, 1024) + self.self_attn_v_proj = nn.Linear(1024, 256, **FACTORY_META) + # o_proj.weight (1024, 2048) + self.self_attn_o_proj = nn.Linear(2048, 1024, **FACTORY_META) + + # MLP + # gate_proj.weight (4096, 1024) + self.mlp_gate_proj = nn.Linear(1024, 4096, **FACTORY_META) + # up_proj.weight (4096, 1024) + self.mlp_up_proj = nn.Linear(1024, 4096, **FACTORY_META) + # down_proj.weight (1024, 4096) + self.mlp_down_proj = nn.Linear(4096, 1024, **FACTORY_META) + + # input_layernorm.dense (3072, 1024) + self.input_layernorm_dense = nn.Linear(1024, 3072, **FACTORY_META) + # post_attention_layernorm.dense (3072, 1024) + self.post_attention_layernorm_dense = nn.Linear(1024, 3072, **FACTORY_META) + + +@pytest.mark.parametrize("dp_world_size", [1, 3, 8, 17]) +@pytest.mark.parametrize( + ("data_parallel_sharding_strategy", "suggested_bucket_size", "fsdp_unit_modules"), + [ + ("no_shard", 40_000_000, [LanguageLayerFlat, VisionLayerFlat]), + ("optim", None, []), + ("optim_grads", 40_000_000, [LanguageLayerFlat]), + ("optim_grads_params", None, [LanguageLayerFlat, VisionLayerFlat, GemmaLayerFlat]), + ], +) +def test_parameter_splitting( + dp_world_size, suggested_bucket_size, fsdp_unit_modules, data_parallel_sharding_strategy +): + model = AllInOneDummyFSDPModel() + param_to_name = {param: name for name, param in model.named_parameters()} + bucketing_policy = BucketingPolicy( + suggested_bucket_size=suggested_bucket_size, + fsdp_unit_modules=fsdp_unit_modules, + data_parallel_sharding_strategy=data_parallel_sharding_strategy, + ) + + ddp_config = Mock() + ddp_config.data_parallel_sharding_strategy = data_parallel_sharding_strategy + + bucket_groups, _, _ = _get_parameter_groups(model, bucketing_policy, {}) + for rank in range(dp_world_size): + try: + store = FakeStore() + init_process_group_kwargs = dict( + backend="fake", store=store, world_size=dp_world_size, rank=rank + ) + torch.distributed.init_process_group(**init_process_group_kwargs) + + for param_group in bucket_groups: + param_list = param_group.params + dp_buf = DataParallelBuffer( + bucket_id=0, + ddp_config=ddp_config, + params=param_list, + chunk_size_factor=param_group.chunk_size_factor, + is_data_distributed=(data_parallel_sharding_strategy != "no_shard"), + ) + + for item_id, param in enumerate(param_list): + elem_per_slice = param.shape[1:].numel() + + # Tensor shard + if data_parallel_sharding_strategy != "no_shard": + start, end = dp_buf._get_item_local_shard_index(item_id) + local_numel = end - start + assert local_numel % elem_per_slice == 0, ( + f"[local_shard] param_name={param_to_name[param]}, " + f"world_size={dp_world_size}, " + f"rank={rank}, item_id={item_id}, " + f"param_shape={tuple(param.shape)}, " + f"local_numel={local_numel}, " + f"elem_per_slice={elem_per_slice}" + ) + + # Full tensor + start, end = dp_buf._get_item_local_index(item_id) + local_numel = end - start + assert local_numel % elem_per_slice == 0, ( + f"[full_tensor] param_name={param_to_name[param]}, " + f"world_size={dp_world_size}, " + f"rank={rank}, item_id={item_id}, " + f"param_shape={tuple(param.shape)}, " + f"local_numel={local_numel}, " + f"elem_per_slice={elem_per_slice}" + ) + finally: + # Clean up process group if created + try: + torch.distributed.destroy_process_group() + except Exception: + pass From 2c98a292b0ddef6ac7104d339bb9561a61eafdd1 Mon Sep 17 00:00:00 2001 From: jianbinc Date: Wed, 25 Feb 2026 16:27:38 +0800 Subject: [PATCH 2/2] Fix compatibility issues of UT in older versions of PyTorch. --- .../megatron_fsdp/test_mfsdp_param_and_grad_buffer.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py b/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py index 189baf115b0..cc59e7ab29a 100644 --- a/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py +++ b/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py @@ -246,7 +246,10 @@ def test_parameter_splitting( ) for item_id, param in enumerate(param_list): - elem_per_slice = param.shape[1:].numel() + if len(param.shape) == 1: + elem_per_slice = 1 + else: + elem_per_slice = param.shape[1:].numel() # Tensor shard if data_parallel_sharding_strategy != "no_shard":