[feat] FSDP support for HybridStack EP-overlap (3/4 of #4798)

megatron/core/distributed/fsdp/mcore_fsdp_adapter.py

@@ -181,7 +181,14 @@ def __init__(
             config.overlap_moe_expert_parallel_comm
             and ddp_config.data_parallel_sharding_strategy == "optim_grads_params"
         ):
-            supported_fsdp_unit_modules = [TransformerLayer, MoETransformerLayer, MambaLayer]
+            from megatron.core.models.hybrid.hybrid_block import HybridStack
+
+            supported_fsdp_unit_modules = [
+                TransformerLayer,
+                MoETransformerLayer,
+                MambaLayer,
+                HybridStack,
+            ]
             assert self.fsdp_unit_modules and all(
                 module in supported_fsdp_unit_modules for module in self.fsdp_unit_modules
             ), (
@@ -190,13 +197,27 @@ def __init__(
                 f"{supported_fsdp_unit_modules}, "
                 f"got {self.fsdp_unit_modules}."
             )
+
+        # HybridStack-specific filter: when bracketed hybrid patterns are used,
+        # the model has a nested layout -- an outer HybridStack root
+        # (``is_layer_group_stack=False``) whose ``layers`` are inner
+        # bracket-group HybridStacks (``is_layer_group_stack=True``).
+        # ``named_modules()`` walks root-first, so with ``[HybridStack]`` the
+        # outer matches first, the inner ones get skipped as its submodules,
+        # and the whole decoder becomes a single FSDP unit. We exclude the
+        # outer so each bracket group is its own unit. Modules without the
+        # attribute (TransformerLayer, etc.) keep the default ``True``.
+        def _fsdp_unit_filter(m):
+            return getattr(m, "is_layer_group_stack", True)
+
         super().__init__(
             config=config,
             module=MegatronFSDP(
                 ddp_config=ddp_config,
                 mixed_precision_policy=self.mp_policy,
                 module=module,
                 fsdp_unit_modules=self.fsdp_unit_modules,
+                fsdp_unit_filter=_fsdp_unit_filter,
                 disable_bucketing=disable_bucketing,
                 device=self.device,
                 dist_index=self.megatron_fsdp_dist_index,

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

@@ -18,7 +18,7 @@
 from contextlib import contextmanager
 from enum import Enum, auto
 from functools import partial
-from typing import Any, Dict, List, Literal, Optional, Tuple, Type
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Type
 
 import torch
 import torch.nn as nn
@@ -211,6 +211,7 @@ def __init__(
         ddp_config: DistributedDataParallelConfig = None,
         mixed_precision_policy: MixedPrecisionPolicy = MixedPrecisionPolicy(),
         fsdp_unit_modules: Optional[List[torch.nn.Module] | List[str]] = None,
+        fsdp_unit_filter: Optional[Callable[[torch.nn.Module], bool]] = None,
         disable_bucketing: bool = False,
         device: Optional[torch.device] = None,
         calculate_per_token_loss: bool = False,
@@ -322,6 +323,11 @@ def __init__(
             if fsdp_unit_modules is not None
             else []
         )
+        # Optional caller-supplied filter run after the isinstance check; lets the
+        # adapter exclude specific class instances (e.g. an outer wrapper that
+        # shares its class with the actual FSDP-unit instances) without leaking
+        # model knowledge into this library.
+        self.fsdp_unit_filter = fsdp_unit_filter
 
         # Determine if we should delay the gradient reduction.
         self.is_delay_grad_reduce = self.data_parallel_sharding_strategy in ["no_shard", "optim"]
@@ -423,7 +429,7 @@ def _init_fsdp_param_and_grad_buffer(self):
                 total_param_elements = 0
                 total_fsdp_module = 0
                 for module in self.module.modules():
-                    if isinstance(module, tuple(self.fsdp_unit_modules)):
+                    if self._is_fsdp_unit_module(module):
                         total_fsdp_module += 1
                         total_param_elements += sum(p.numel() for p in module.parameters())
                 # The suggested size is twice the number of elements in the FSDP modules.
@@ -455,6 +461,21 @@ def _import_class_from_path(self, class_path: str):
         cls = getattr(module, class_name)
         return cls
 
+    def _is_fsdp_unit_module(self, module: nn.Module) -> bool:
+        """Whether ``module`` should be treated as an FSDP unit.
+
+        Default: ``isinstance(module, tuple(fsdp_unit_modules))``. When the
+        caller provides ``fsdp_unit_filter``, the filter runs after the
+        isinstance check and can exclude specific instances -- useful when a
+        wrapping container shares its class with the actual unit instances
+        (so a pure class-based match would register the wrong layer).
+        """
+        if not isinstance(module, tuple(self.fsdp_unit_modules)):
+            return False
+        if self.fsdp_unit_filter is not None:
+            return self.fsdp_unit_filter(module)
+        return True
+
     def all_gather_and_wait_parameters_ready(
         self,
         params,
@@ -553,7 +574,6 @@ def _register_fsdp_hooks(self, root_module):
         `optim` and `optim_grads` do not require FSDP units because they do not
         shard model parameters.
         """
-        fsdp_unit_modules = self.fsdp_unit_modules
 
         def _param_list_for_submodule_unshard(
             module: nn.Module, pass_direction: Literal["forward", "backward"]
@@ -590,7 +610,7 @@ def _param_list_for_submodule_unshard(
                     # recomputation on individual submodules.
                     return list(module.parameters(recurse=False))
             else:
-                if isinstance(module, tuple(fsdp_unit_modules)):
+                if self._is_fsdp_unit_module(module):
                     # FSDP unit modules should be unsharded and communicated together.
                     return list(module.parameters())
                 else:
@@ -690,7 +710,7 @@ def _post_backward_release_module(module, *unused):
             - Releases the module's parameters for the backward phase to free memory.
             - Marks the module as IDLE in the training state machine.
             """
-            assert isinstance(module, tuple(fsdp_unit_modules))
+            assert self._is_fsdp_unit_module(module)
             assert self.data_parallel_sharding_strategy == "optim_grads_params"
 
             # Release parameters for this module after backward.
@@ -967,8 +987,8 @@ def _post_forward(module: nn.Module, input: Any, output: Any):
                 lazy_release = False
                 module._training_state = TrainingState.IDLE
 
-            assert isinstance(
-                module, tuple(fsdp_unit_modules)
+            assert self._is_fsdp_unit_module(
+                module
             ), "_post_forward hook should only be registered on FSDP unit modules."
 
             # Release the module parameters after the forward pass to save memory.
@@ -1063,7 +1083,7 @@ def _register_pre_backward_param_unshard_hook(module):
             if not self.enable_fine_grained_param_gather_hook:
                 _register_pre_forward_param_unshard_hook(module)
 
-            if isinstance(module, tuple(fsdp_unit_modules)):
+            if self._is_fsdp_unit_module(module):
                 fsdp_modules.append(module)
                 # Register the forward post-hook to reshard FSDP unit module parameters
                 # after the forward pass, except when recomputing forward activations,
@@ -1087,7 +1107,7 @@ def _register_pre_backward_param_unshard_hook(module):
 
             # Register the post-backward hook to deallocate model parameters
             # and reduce-scatter gradients after the backward pass.
-            if isinstance(module, tuple(fsdp_unit_modules)):
+            if self._is_fsdp_unit_module(module):
                 if self.ddp_config.data_parallel_sharding_strategy == "optim_grads_params":
                     self.forward_pre_hooks[f"module {name} register post-backward hook"] = (
                         module.register_forward_pre_hook(

megatron/core/models/common/fine_grained_callables.py

@@ -64,8 +64,27 @@ def submodule_mtp_pre_dispatch_forward(node, hidden_states):
                     )
 
             offset = get_mtp_layer_offset(layer.config, node.chunk_state.model.vp_stage)
-            node.chunk_state.mtp_hidden_states = list(torch.chunk(hidden_states, 1 + offset, dim=0))
-            hidden_states = node.chunk_state.mtp_hidden_states[offset]
+            chunks = list(torch.chunk(hidden_states, 1 + offset, dim=0))
+            # Store DETACHED chunks in chunk_state. mtp_hidden_states is a
+            # ``chunk_state``-level Python list that crosses slot boundaries
+            # (set here in MTP's pre_dispatch slot, later torch.cat'd in MTP's
+            # mtp_post_process slot before feeding the LM head). Without an
+            # explicit detach the chunks keep their grad_fn from torch.chunk →
+            # whatever upstream node produced ``hidden_states`` (e.g. the
+            # final_norm we apply above for the HybridModel empty-decoder case),
+            # which means mtp_post_process.backward and pre_dispatch.backward
+            # both traverse that same grad_fn — for a TENorm-backed final_norm
+            # (an OpFuser op) the second traversal hits ``ctx.tensor_objects is
+            # None`` and raises ``ctx must have .tensor_objects to restore
+            # saved tensors``. Using ``node.detach`` records the originals in
+            # before_detached so pre_dispatch's backward_impl still pulls the
+            # LM-head-side grad (accumulated on the detached leaves by the
+            # post_process / mtp_post_process backward chain) back into the
+            # outputs+before_detached run_backward — i.e. the gradient flow
+            # remains mathematically equivalent, just no longer shared across
+            # slots.
+            node.chunk_state.mtp_hidden_states = [node.detach(c) for c in chunks]
+            hidden_states = chunks[offset]
 
         input_ids, position_ids, padding_mask, decoder_input, hidden_states = layer._get_embeddings(
             input_ids=node.chunk_state.input_ids,
@@ -144,9 +163,14 @@ def rng_context_wrapper(func, *args, **kwargs):
 
 def get_layer_moe_metadata(layer):
     """Return ``(is_moe, num_local_experts)`` for schedule-node construction."""
+    from megatron.core.models.hybrid.hybrid_block import HybridStack
 
     if isinstance(layer, MultiTokenPredictionLayer):
         return get_layer_moe_metadata(layer.mtp_model_layer)
+    if isinstance(layer, HybridStack):
+        from megatron.core.models.hybrid.fine_grained_callables import get_hybrid_stack_moe_metadata
+
+        return get_hybrid_stack_moe_metadata(layer)
     if isinstance(layer, TransformerLayer):
         is_moe = isinstance(layer.mlp, MoELayer)
         num_local_experts = layer.mlp.num_local_experts if is_moe else None
@@ -160,9 +184,15 @@ def build_layer_callables(layer):
 
     Returns ``(forward_funcs, backward_dw)``.
     """
+    from megatron.core.models.hybrid.hybrid_block import HybridStack
 
     if isinstance(layer, MultiTokenPredictionLayer):
         return build_mtp_layer_callables(layer)
+    if isinstance(layer, HybridStack):
+        from megatron.core.models.hybrid.fine_grained_callables import build_hybrid_stack_callables
+
+        forward_funcs, backward_dw, _, _ = build_hybrid_stack_callables(layer)
+        return forward_funcs, backward_dw
     if isinstance(layer, TransformerLayer):
         return build_transformer_layer_callables(layer)