fsdp: avoid double sharding of MoE experts when EP is enabled

README.md

@@ -232,6 +232,21 @@ For a version of Megatron Core with only torch, run:
 pip install megatron-core
 ```
 
+### Optional MoE Dependencies
+
+For Mixture of Experts (MoE) training with Grouped GEMM support:
+
+```bash
+pip install --no-build-isolation megatron-core[moe]
+```
+
+**Note:** The `nv-grouped-gemm` package requires:
+- CUDA toolkit (nvcc) with CUTLASS headers
+- On Ubuntu/Debian: `apt-get install libcutlass-dev`
+- GPU with compute capability >= 8.0
+
+If you encounter build errors, you can skip this optional dependency and use MoE without Grouped GEMM optimization.
+
 ## System Requirements
 
 ### Hardware Requirements

megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py

@@ -349,8 +349,15 @@ def _init_fsdp_param_and_grad_buffer(self):
             device=self.device,
             reset_parameters_for_meta_device_init_module=self.init_model_with_meta_device,
         )
-        self.param_to_name = {p: name for name, p in self.module.named_parameters()}
-        self.raw_param = dict(self.module.named_parameters())
+        # Exclude expert-parallel params from FSDP bookkeeping
+        self.param_to_name = {}
+        self.raw_param = {}
+        from megatron.core.utils import is_ep_owned_param
+        for name, param in self.module.named_parameters():
+            if is_ep_owned_param(self.module, name):
+                continue
+            self.param_to_name[param] = name
+            self.raw_param[name] = param
 
         # Initialize a gradient buffer and accumulation stream for the GradReducePipeline.
         self.side_stream_for_buffer_copy_and_grad_accum = torch.cuda.Stream()
@@ -1135,7 +1142,13 @@ def _replace_param_with_distributed_if_needed(self):
 
         pg_buffer = self.param_and_grad_buffer
         fsdp_params = dict(pg_buffer.optimizer_named_parameters)
-        for name, _ in self.module.named_parameters():
+        from megatron.core.utils import is_ep_owned_param
+        from megatron.core.utils import assert_not_fsdp_wrapped_ep_param
+        for name, param in self.module.named_parameters():
+            # Skip FSDP replacement for expert-parallel-owned parameters
+            if is_ep_owned_param(self.module, name):
+                assert_not_fsdp_wrapped_ep_param(self.module, name)
+                continue
             assert name in fsdp_params, f"Parameter {name} not found in FSDP parameters."
             dist_param = fsdp_params[name]
             # Set the __fsdp_param__ attribute to True to indicate that this
@@ -1152,7 +1165,11 @@ def _replace_param_with_raw_if_needed(self):
             return
         self.is_param_fsdp_distributed = False
 
-        for name, _ in self.module.named_parameters():
+        from megatron.core.utils import is_ep_owned_param
+        for name, param in self.module.named_parameters():
+            # Skip FSDP replacement for expert-parallel-owned parameters
+            if is_ep_owned_param(self.module, name):
+                continue
             assert name in self.raw_param, f"Raw parameter {name} not found in module."
             _replace_module_parameter(self.module, name, self.raw_param[name])
 

megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py

@@ -1317,6 +1317,13 @@ def _does_param_require_new_bucket(param):
     # All parameters in the module are assigned a parameter group, even non-FSDP modules.
     parameter_groups = []
     for name, param in module.named_parameters():
+        # Skip FSDP param/grad buffer construction for expert-parallel-owned parameters
+        parent_module = module
+        for submodule_name in name.split(".")[:-1]:
+            parent_module = getattr(parent_module, submodule_name, parent_module)
+        if getattr(parent_module, "expert_parallel_enabled", False):
+            continue
+
         # We need this information to correctly dynamically allocate Tensors!
         is_fp8 = is_float8tensor(param)
         is_fp8_meta_device_init = meta_device_init_fp8_params.get(name, (False, False))[0]

megatron/core/transformer/moe/grouped_gemm_util.py

@@ -13,10 +13,23 @@ def grouped_gemm_is_available():
 
 def assert_grouped_gemm_is_available():
     """Assert that grouped_gemm is available."""
-    assert grouped_gemm_is_available(), (
-        "Grouped GEMM is not available. Please run "
-        "`pip install git+https://github.com/fanshiqing/grouped_gemm@v1.1.4`."
+    error_msg = (
+        "Grouped GEMM is not available. To use MoE with grouped GEMM, you need to install "
+        "nv-grouped-gemm.\n\n"
+        "Installation options:\n"
+        "1. Install from PyPI (requires CUDA toolkit and CUTLASS headers):\n"
+        "   pip install 'megatron-core[moe]'\n"
+        "   or\n"
+        "   pip install nv-grouped-gemm\n\n"
+        "2. Build from source:\n"
+        "   pip install git+https://github.com/fanshiqing/grouped_gemm@v1.1.4\n\n"
+        "Note: Building from source requires:\n"
+        "- CUDA toolkit (nvcc)\n"
+        "- CUTLASS headers (can be installed via 'apt-get install libcutlass-dev' on Ubuntu)\n"
+        "- Compatible GPU with compute capability >= 8.0\n\n"
+        "If you don't need MoE functionality, you can continue without this package."
     )
+    assert grouped_gemm_is_available(), error_msg
 
 
 ops = grouped_gemm.ops if grouped_gemm_is_available() else None

megatron/core/transformer/moe/moe_layer.py

@@ -84,6 +84,14 @@ def __init__(
         self.token_dispatcher: Optional[MoETokenDispatcher] = None
         self.layer_number = layer_number
 
+        # Used by FSDP to avoid double sharding
+        # Flag indicating if this MoE layer's expert parameters are managed by Expert Parallelism
+        self.expert_parallel_enabled = ep_size > 1
+
+        # Invariant:
+        # When expert_parallel_enabled=True, these parameters must never be wrapped,
+        # replaced, or sharded by FSDP. All FSDP code paths must explicitly skip them.
+
     @abstractmethod
     def forward(self, hidden_states):
         """Forward method for the MoE layer."""

megatron/core/utils.py

@@ -869,6 +869,41 @@ def log_on_each_pipeline_stage(
         logger.log(*args, **kwargs)
 
 
+def assert_not_fsdp_wrapped_ep_param(module, param_name: str):
+    """
+    EP-owned parameters must never be wrapped or replaced by FSDP.
+    This assertion exists to prevent silent double-sharding regressions.
+    """
+    if getattr(module, "expert_parallel_enabled", False):
+        raise AssertionError(
+            f"FSDP attempted to manage EP-owned parameter: {param_name}. "
+            "This indicates a missing EP exclusion in the FSDP code path."
+        )
+
+
+def is_ep_owned_param(module, name):
+    """
+    Check if a parameter is owned by an expert-parallel-enabled submodule.
+
+    Traverses the module hierarchy using the parameter name to find the parent
+    module and checks if any parent has expert_parallel_enabled set to True.
+
+    Args:
+        module: The root module to start traversal from
+        name: The fully qualified parameter name (e.g., "layers.0.mlp.weight")
+
+    Returns:
+        bool: True if the parameter is owned by an expert-parallel-enabled submodule,
+              False otherwise
+    """
+    parent = module
+    for sub in name.split(".")[:-1]:
+        parent = getattr(parent, sub, parent)
+        if getattr(parent, "expert_parallel_enabled", False):
+            return True
+    return False
+
+
 def check_param_hashes_across_dp_replicas(
     model: List[torch.nn.Module], cross_check: bool = False
 ) -> bool: