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 9190ac04167..2d7f573afdb 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 @@ -1404,6 +1404,29 @@ def _does_param_require_new_bucket(param): is_expert_parameter = lambda n, p: ".experts." in n + def _should_split_from_grouped_expert_bucket( + is_expert_param: bool, + param: torch.nn.Parameter, + param_chunk_size_factor: int, + chunk_size_factor: int, + same_factor_params: List[torch.nn.Parameter], + ) -> bool: + """ + Split grouped expert (>=3D) tensors with heterogeneous chunk size + factors into separate buckets to avoid LCM-inflated bucket alignment + padding. + """ + # Non-expert groups keep the original LCM/fragment merge. + if not is_expert_param: + return False + # Param already aligns with bucket chunk size factor (always true for + # the first param after sort); no split needed. + if param_chunk_size_factor == chunk_size_factor: + return False + return to_local_if_dtensor(param).dim() >= 3 or any( + to_local_if_dtensor(p).dim() >= 3 for p in same_factor_params + ) + # Step 1: Group the parameters according to their execution order and attributes. # FSDP unit module parameters are split into multiple parameter sub-groups. # All parameters in the module are assigned a parameter group, even non-FSDP modules. @@ -1503,17 +1526,27 @@ def _does_param_require_new_bucket(param): remaining_params = [] for param in params: param_shape = to_local_if_dtensor(param).shape + param_chunk_size_factor = param_shape[1:].numel() + if _should_split_from_grouped_expert_bucket( + group.is_expert_param, + param, + param_chunk_size_factor, + chunk_size_factor, + same_factor_params, + ): + remaining_params.append(param) + continue if ( - param_shape[1:].numel() == chunk_size_factor + param_chunk_size_factor == chunk_size_factor or ( - chunk_size_factor % param_shape[1:].numel() == 0 + chunk_size_factor % param_chunk_size_factor == 0 and param_shape.numel() % chunk_size_factor == 0 ) or (param_shape.numel() < chunk_size_factor) ): same_factor_params.append(param) else: - lcm_chunk_size_factor = math.lcm(chunk_size_factor, param_shape[1:].numel()) + lcm_chunk_size_factor = math.lcm(chunk_size_factor, param_chunk_size_factor) chunk_size_factor = lcm_chunk_size_factor same_factor_params.append(param) # Create a new parameter group with the same chunk size factor. 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..8ad3789b449 --- /dev/null +++ b/tests/unit_tests/distributed/megatron_fsdp/test_mfsdp_param_and_grad_buffer.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import math + +import torch + +from megatron.core.distributed.fsdp.src.megatron_fsdp.param_and_grad_buffer import ( + BucketingPolicy, + _get_parameter_groups, +) + + +class _ExpertTestModule(torch.nn.Module): + """ + Mock module whose params are routed under `.experts.` to trigger + is_expert_param=True. The outer `layer` attribute puts a dot before + `experts` in the parameter path (e.g. `layer.experts.linear_fc1`). + """ + + def __init__(self, shapes): + super().__init__() + self.layer = torch.nn.Module() + self.layer.experts = torch.nn.ParameterDict( + {name: torch.nn.Parameter(torch.empty(shape)) for name, shape in shapes.items()} + ) + + +def _get_bucket_signatures(module): + bucket_groups, _, _ = _get_parameter_groups( + module, BucketingPolicy(suggested_bucket_size=None), meta_device_init_fp8_params={} + ) + param_to_name = {param: name for name, param in module.named_parameters()} + return [ + { + "chunk_size_factor": group.chunk_size_factor, + "params": [(param_to_name[param], tuple(param.shape)) for param in group.params], + } + for group in bucket_groups + ] + + +def test_grouped_expert_weights_split_when_chunk_size_factors_differ(): + """Grouped expert weights with mismatched chunk size factors get routed to separate buckets.""" + num_local_experts = 4 + hidden_size = 12 + moe_ffn_hidden_size = 8 + shapes = { + "linear_fc1": (num_local_experts, 2 * moe_ffn_hidden_size, hidden_size), + "linear_fc2": (num_local_experts, hidden_size, moe_ffn_hidden_size), + } + module = _ExpertTestModule(shapes) + + assert _get_bucket_signatures(module) == [ + { + "chunk_size_factor": torch.Size(shapes["linear_fc1"])[1:].numel(), + "params": [("layer.experts.linear_fc1", shapes["linear_fc1"])], + }, + { + "chunk_size_factor": torch.Size(shapes["linear_fc2"])[1:].numel(), + "params": [("layer.experts.linear_fc2", shapes["linear_fc2"])], + }, + ] + + +def test_per_expert_2d_weights_merge_via_lcm(): + """Per-expert 2D weights merge into a single bucket via LCM chunk size factor.""" + hidden_size = 12 + moe_ffn_hidden_size = 8 + shapes = { + "linear_fc1": (2 * moe_ffn_hidden_size, hidden_size), + "linear_fc2": (hidden_size, moe_ffn_hidden_size), + } + module = _ExpertTestModule(shapes) + + assert _get_bucket_signatures(module) == [ + { + "chunk_size_factor": math.lcm( + torch.Size(shapes["linear_fc1"])[1:].numel(), + torch.Size(shapes["linear_fc2"])[1:].numel(), + ), + "params": [ + ("layer.experts.linear_fc1", shapes["linear_fc1"]), + ("layer.experts.linear_fc2", shapes["linear_fc2"]), + ], + } + ]