[Main][feat] Support A2A Overlap for Megatron-FSDP

megatron/core/distributed/fsdp/mcore_fsdp_adapter.py

@@ -40,6 +40,7 @@
 from megatron.core.distributed.data_parallel_base import _BaseDataParallel
 from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig
 from megatron.core.process_groups_config import ProcessGroupCollection
+from megatron.core.transformer.enums import CudaGraphScope
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.transformer.transformer_layer import TransformerLayer
 from megatron.core.utils import is_te_min_version, log_single_rank
@@ -157,6 +158,31 @@ def __init__(
 
         self._annotate_tensor_parallelism(module)
 
+        if config.overlap_moe_expert_parallel_comm:
+            assert not ddp_config.fsdp_double_buffer, (
+                "1F1B overlap with FSDP does not support double buffer. "
+                "Please set fsdp_double_buffer=False in the ddp config."
+            )
+            partial_cuda_graph_scopes = [
+                scope
+                for scope in config.cuda_graph_scope
+                if scope
+                not in (CudaGraphScope.full_iteration, CudaGraphScope.full_iteration_inference)
+            ]
+            assert not partial_cuda_graph_scopes, (
+                "1F1B overlap with FSDP does not support partial CUDA graph scopes "
+                f"({partial_cuda_graph_scopes}). "
+                "Please use cuda_graph_scope='full' or disable CUDA graphs."
+            )
+
+        if (
+            config.overlap_moe_expert_parallel_comm
+            and ddp_config.data_parallel_sharding_strategy == "optim_grads_params"
+        ):
+            assert self.fsdp_unit_modules == [TransformerLayer], (
+                "EP overlap with FSDP currently requires fsdp_unit_modules "
+                f"to be [TransformerLayer], got {self.fsdp_unit_modules}."
+            )
         super().__init__(
             config=config,
             module=MegatronFSDP(
@@ -169,8 +195,20 @@ def __init__(
                 dist_index=self.megatron_fsdp_dist_index,
                 calculate_per_token_loss=config.calculate_per_token_loss,
                 init_model_with_meta_device=config.init_model_with_meta_device,
+                # EP overlap schedule calls sub-modules directly instead of
+                # TransformerLayer.forward(), so fine-grained hooks are needed
+                # to manage _training_state and all-gather each sub-module's
+                # parameters individually.  This applies to all sharding
+                # strategies (not only optim_grads_params) because the hooks
+                # also maintain per-module training-state bookkeeping that the
+                # gradient-reduction pipeline relies on.
                 enable_fine_grained_param_gather_hook=(
-                    config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather
+                    (config.fp8_recipe == "mxfp8" and ddp_config.fp8_param_gather)
+                    or config.overlap_moe_expert_parallel_comm
+                ),
+                enable_fine_grained_param_gather_backward_hook=(
+                    config.overlap_moe_expert_parallel_comm
+                    and ddp_config.data_parallel_sharding_strategy == "optim_grads_params"
                 ),
             ),
         )

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

@@ -17,6 +17,7 @@
 import logging
 from contextlib import contextmanager
 from enum import Enum, auto
+from functools import partial
 from typing import Any, Dict, List, Optional, Tuple
 
 import torch
@@ -215,6 +216,7 @@ def __init__(
         fsdp_db_use_persist_buf_on_alloc_fail: bool = False,
         disable_symmetric_registration: bool = False,
         enable_fine_grained_param_gather_hook: bool = False,
+        enable_fine_grained_param_gather_backward_hook: bool = False,
         report_nan_in_param_grad: bool = False,
     ):
         super().__init__()
@@ -267,6 +269,9 @@ def __init__(
         self.calculate_per_token_loss = calculate_per_token_loss
         self.init_model_with_meta_device = init_model_with_meta_device
         self.enable_fine_grained_param_gather_hook = enable_fine_grained_param_gather_hook
+        self.enable_fine_grained_param_gather_backward_hook = (
+            enable_fine_grained_param_gather_backward_hook
+        )
         self.report_nan_in_param_grad = report_nan_in_param_grad
 
         # FSDPDistributedIndex stores the process groups and meshes used by Megatron-FSDP.
@@ -694,18 +699,12 @@ def _process_post_backward_gradients(param_list):
             # Filter out shared parameters whose gradients are handled by the root hook.
             param_list = [p for p in param_list if not getattr(p, "_is_shared", False)]
 
-            # Filter out parameters whose gradient processing is deferred to a delayed
-            # wgrad accumulation hook (post_wgrad_grad_acc_hook).  If skip_backward_post_hook
-            # is set but the delayed hook was never installed, process the parameter
-            # immediately as a safety fallback to avoid silently dropping gradients.
-            param_list = [
-                p
-                for p in param_list
-                if not (
-                    getattr(p, 'skip_backward_post_hook', False)
-                    and hasattr(p, 'post_wgrad_grad_acc_hook')
-                )
-            ]
+            # Make sure for delayed wgrad params, the grad_acc_hooks are registered.
+            for p in param_list:
+                if getattr(p, 'skip_backward_post_hook', False):
+                    assert hasattr(
+                        p, 'post_wgrad_grad_acc_hook'
+                    ), "Missing grad accumulation hook for delayed_wgrad_compute param."
 
             if not param_list:
                 return
@@ -885,7 +884,7 @@ def _pre_backward_param_unshard(module: nn.Module, *unused):
 
         self._root_pre_backward_hook_issued = False
 
-        def _root_pre_backward(module: nn.Module, *unused):
+        def _root_pre_backward(module: nn.Module, *unused, skip_backward_hook: bool = False):
             """Marks the module's training state as PRE_BACKWARD before the
             backprop, this function is registered on the root module.
 
@@ -921,6 +920,8 @@ def _root_pre_backward(module: nn.Module, *unused):
                     param.grad_added_to_main_grad = False
             # Queue the root post-backward hook to reduce leftover gradients after
             # the backward pass.
+            if skip_backward_hook:
+                return
             torch.autograd.Variable._execution_engine.queue_callback(_root_post_backward)
 
         @torch.compiler.disable
@@ -1008,12 +1009,24 @@ def _register_pre_backward_param_unshard_hook(module):
                 create_custom_backward_hook(module, _pre_backward_param_unshard)
             )
 
+        # These hooks need to be exposed for manual management by 1F1B Overlapping
+        # and triggered by 1F1B Overlapped execution pipeline, except for
+        # `param_unshard` hook that needs to be installed at param level,
+        # such that non-overlapped params like embedding layer are also correctly
+        # unsharded.
+        self.post_forward_release_module = partial(_post_forward, input=None, output=None)
+        self.post_backward_release_module = _post_backward_release_module
+        self.pre_backward = partial(_root_pre_backward, module=None, skip_backward_hook=True)
+        self.post_backward = _root_post_backward
+
         fsdp_modules = []
         for name, module in root_module.named_modules():
             # Set post backward hook for TE grouped gemm if enabled comm overlap
             setup_delayed_wgrad_acc_hook(module, _process_post_backward_gradients)
             if self.enable_fine_grained_param_gather_hook:
                 _register_pre_forward_param_unshard_hook(module)
+            if self.enable_fine_grained_param_gather_backward_hook:
+                _register_pre_backward_param_unshard_hook(module)
 
             # Skip if the module is already registered in fsdp_modules.
             if any(is_submodule(module, fsdp_module) for fsdp_module in fsdp_modules):
@@ -1068,7 +1081,11 @@ def _register_pre_backward_param_unshard_hook(module):
                     continue
                 self.grad_acc_hooks[f"grad_acc and reduce for {self.param_to_name[param]}"] = (
                     param.register_post_accumulate_grad_hook(
-                        lambda p: _process_post_backward_gradients([p])
+                        lambda p: (
+                            None
+                            if getattr(p, 'skip_backward_post_hook', False)
+                            else _process_post_backward_gradients([p])
+                        )
                     )
                 )
 

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

@@ -91,6 +91,23 @@ def is_submodule(module, parent_module, strict=True):
     return False
 
 
+def find_megatron_fsdp(model):
+    """Walk the model wrapper chain to find a MegatronFSDP instance, if any."""
+    # Lazy import to avoid a circular import: megatron_fsdp.py transitively imports
+    # this module during its own initialization, so a top-level import of
+    # MegatronFSDP here would fail with a partially-initialized module error.
+    try:
+        from megatron.core.distributed.fsdp.src.megatron_fsdp.megatron_fsdp import MegatronFSDP
+    except (ImportError, ModuleNotFoundError):
+        return None
+    m = model
+    while m is not None:
+        if isinstance(m, MegatronFSDP):
+            return m
+        m = getattr(m, 'module', None)
+    return None
+
+
 def get_mesh_names(
     device_mesh: Optional[DeviceMesh] = None, only_submesh_dims: bool = False
 ) -> list[str]:

megatron/core/models/common/model_chunk_schedule_plan.py

@@ -173,6 +173,46 @@ def create_node(stream, module, name):
         else:
             self.mtp_post_process = NoopScheduleNode()
 
+    def set_fsdp_reshard_hooks(self, post_forward_hook, post_backward_hook):
+        """Wire FSDP parameter release callbacks for the fine-grained overlap schedule.
+
+        The EP overlap schedule bypasses the normal FSDP forward/backward hooks
+        (registered on the FSDP unit module) because it calls sub-modules directly
+        instead of going through TransformerLayer.forward(). This method attaches
+        explicit release hooks to individual schedule nodes so that all-gathered
+        parameters are freed at the right time.
+
+        Args:
+            post_forward_hook: Callable(module) that releases forward-pass params
+                (bwd=False). Typically ``fsdp_wrapper.post_forward_release_module``.
+            post_backward_hook: Callable(module) that releases backward-pass params
+                (bwd=True). Typically ``fsdp_wrapper.post_backward_release_module``.
+        """
+        from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer
+        from megatron.core.transformer.transformer_layer import TransformerLayer
+
+        assert isinstance(self.layer, (TransformerLayer, MultiTokenPredictionLayer)), (
+            f"Megatron FSDP with EP Overlap only supports TransformerLayer, "
+            f"but got {type(self.layer).__name__}."
+        )
+
+        if isinstance(self.layer, TransformerLayer):
+            hook_module = self.layer
+        else:
+            hook_module = self.layer.mtp_model_layer
+
+        # After the last backward op (attn), release backward-pass params.
+        self.attn.set_post_backward_hook(lambda: post_backward_hook(hook_module))
+
+        # Determine the last node in forward order.
+        if isinstance(self.moe_combine, NoopScheduleNode):
+            last_fwd_node = self.mlp
+        else:
+            last_fwd_node = self.moe_combine
+
+        # After the last forward op, release forward-pass params.
+        last_fwd_node.set_post_forward_hook(lambda: post_forward_hook(hook_module))
+
     def get_fp8_context(self):
         """
         Get the fp8 context for the transformer layer.

megatron/core/models/gpt/fine_grained_callables.py

@@ -267,14 +267,17 @@ def __init__(
             bwd_dw_callables (list): List of weight gradient functions for the layer.
             extra_args (dict): Extra arguments for the node: is_moe, config.
         """
-        # determine whether to free input memory
+        # Determine whether to free input memory
         config = extra_args.get("config", None)
         assert config is not None, "model config must be passed to TransformerLayerNode."
         is_moe = extra_args.get("is_moe", False)
         num_local_experts = extra_args.get("num_local_experts", None)
         free_input = should_free_input(name, is_moe, config, num_local_experts)
         self.delay_wgrad_compute = extra_args.get("delay_wgrad_compute", False)
 
+        self.is_layer_first_node = None
+        self.is_layer_last_node = None
+
         super().__init__(
             weak_method(self.forward_impl),
             stream,
@@ -289,6 +292,7 @@ def __init__(
         self.detached = tuple()
         self.before_detached = tuple()
         self.is_mtp = extra_args.get("is_mtp", False)
+        self.post_wgrad_grad_acc_hooks = None
 
         # Create flags to indicate first and last layer
         self.is_first_layer = extra_args.get("is_first_layer", False)
@@ -322,6 +326,20 @@ def backward_impl(self, outputs, output_grad):
         # return grads for record stream
         return grads
 
+    def forward(self, *inputs):
+        """Execute forward pass and corresponding hooks."""
+        output = super().forward(*inputs)
+        if self.is_layer_last_node:
+            self._post_forward_hook()
+        return output
+
+    def backward(self, *output_grad):
+        """Execute backward pass and corresponding hooks."""
+        grads = super().backward(*output_grad)
+        if not self.delay_wgrad_compute and self.is_layer_first_node:
+            self._post_backward_hook()
+        return grads
+
     def backward_dw(self):
         """Computes the weight gradients for the transformer layer node."""
         if not self.delay_wgrad_compute:
@@ -335,8 +353,45 @@ def backward_dw(self):
                 module.backward_dw()
             nvtx_range_pop(nvtx_msg)
 
+        # Collecting gradient acc hooks if there is `post_wgrad_grad_acc_hook`
+        # attribute attached to param, o.w. the wgrad hook wouldn't be fired.
+        if self.post_wgrad_grad_acc_hooks is None:
+            self.post_wgrad_grad_acc_hooks = []
+            for module in self.bwd_dw_callables:
+                for param in module.parameters():
+                    # Collect hook only if the gradient is generated in current
+                    # TransformerLayerNode, because the grad_acc hook needs
+                    # to be executed right after `backward_dw` finishes.
+                    # For example: Shared expert's hook should be collected in
+                    # `attn` Node, even if the param belongs to `mlp` Node.
+                    if (
+                        getattr(param, "post_wgrad_grad_acc_hook", False)
+                        and param.requires_grad
+                        and param.grad is not None
+                    ):
+                        self.post_wgrad_grad_acc_hooks.append(param.post_wgrad_grad_acc_hook)
+
+        # Execute gradient accumulation hooks after wgrad compute.
+        if self.post_wgrad_grad_acc_hooks:
+            with torch.cuda.stream(self.stream):
+                for hook in self.post_wgrad_grad_acc_hooks:
+                    hook()
+
+        # Execute TransformerLayer backward hook.
+        if self.is_layer_first_node:
+            self._post_backward_hook()
         self.bwd_dw_callables = None
 
+    def set_post_forward_hook(self, hook):
+        """Register post_forward_hook at TransformerLayer level."""
+        self.is_layer_last_node = True
+        self._post_forward_hook = hook
+
+    def set_post_backward_hook(self, hook):
+        """Register post_backward_hook at TransformerLayer level."""
+        self.is_layer_first_node = True
+        self._post_backward_hook = hook
+
     def __del__(self):
         # Release reference as early as possible, this helps avoid memory leak.
         self.before_detached = None
@@ -369,10 +424,13 @@ def __init__(self, layer):
         self.layer = layer
         self.graphed_backward_dw_callable = None
         self.attn_dw_callable = layer.self_attention.backward_dw
+        self.submodules = [layer.self_attention]
         if layer.is_moe_layer:
             self.shared_expert_dw_callable = partial(
                 layer.mlp.backward_dw, routed_experts=False, shared_experts=True
             )
+            if layer.mlp.use_shared_expert:
+                self.submodules.append(layer.mlp.shared_experts)
         else:
             self.shared_expert_dw_callable = None
         self.cuda_graph_scope = layer.config.cuda_graph_scope
@@ -394,6 +452,17 @@ def set_graphed_backward_dw_callable(self, graphed_backward_dw_callable):
         """Store the CUDA graphed backward weight gradient callable."""
         self.graphed_backward_dw_callable = graphed_backward_dw_callable
 
+    def parameters(self):
+        """Returns an iterator over module parameters.
+
+        This method mimics the behavior of torch.nn.Module.parameters() by yielding
+        all parameters from the submodules managed by this wrapper. It is used to
+        collect parameters that require gradient computation during the backward pass.
+        """
+        for module in self.submodules:
+            for param in module.parameters():
+                yield param
+
 
 def build_transformer_layer_callables(layer: TransformerLayer):
     """Create callables for transformer layer nodes.