diff --git a/README.md b/README.md index fb74c9420e8..12893b47ac5 100644 --- a/README.md +++ b/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 diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index e2cbccf4356..f4bc476a137 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/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]) 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 04ea09970f4..e9be17b0c2e 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 @@ -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] diff --git a/megatron/core/transformer/moe/grouped_gemm_util.py b/megatron/core/transformer/moe/grouped_gemm_util.py index 5dd344816bd..be096c817fd 100644 --- a/megatron/core/transformer/moe/grouped_gemm_util.py +++ b/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 diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 91d443fd9ec..a7077593c80 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/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.""" diff --git a/megatron/core/utils.py b/megatron/core/utils.py index 8bbce518096..6e81e101a92 100644 --- a/megatron/core/utils.py +++ b/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: