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 6c907bc25ec..a5ba81124f7 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,28 @@ 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_csf_base: int, + chunk_size_factor: int, + same_factor_params: List[torch.nn.Parameter], + ) -> bool: + """ + Split grouped expert (>=3D) tensors with heterogeneous CSF bases 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 CSF (always true for the first param + # after sort); no split needed. + if param_csf_base == 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 +1525,27 @@ def _does_param_require_new_bucket(param): remaining_params = [] for param in params: param_shape = to_local_if_dtensor(param).shape + param_csf_base = param_shape[1:].numel() + if _should_split_from_grouped_expert_bucket( + group.is_expert_param, + param, + param_csf_base, + chunk_size_factor, + same_factor_params, + ): + remaining_params.append(param) + continue if ( - param_shape[1:].numel() == chunk_size_factor + param_csf_base == chunk_size_factor or ( - chunk_size_factor % param_shape[1:].numel() == 0 + chunk_size_factor % param_csf_base == 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_csf_base) chunk_size_factor = lcm_chunk_size_factor same_factor_params.append(param) # Create a new parameter group with the same chunk size factor.