[DEV] fix(megatron-fsdp): reduce padding for grouped expert weights

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.

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"]),
+            ],
+        }
+    ]