[Dev](moe):Refine A2A overlap under CUDA_DEVICE_MAX_CONNECTIONS=1

megatron/core/models/common/model_chunk_schedule_plan.py

@@ -3,6 +3,7 @@
 from contextlib import nullcontext
 from typing import Optional
 
+import os
 import torch
 from torch import Tensor
 
@@ -42,6 +43,7 @@ class TransformerLayerSchedulePlan:
     ├── moe_dispatch (TransformerLayerNode): dispatch All2All
     ├── mlp (TransformerLayerNode): mlp module
     ├── moe_combine (TransformerLayerNode): combine All2All
+    ├── post_combine (TransformerLayerNode): combine post process -> mlp_bda
     └── mtp_post_process (PostProcessNode): mtp post process
 
     Note that MTP layer has the same operation and execution order with TransformerLayer regarding
@@ -57,6 +59,7 @@ class TransformerLayerSchedulePlan:
     moe_dispatch = None
     mlp = None
     moe_combine = None
+    post_combine = None
     mtp_post_process = None
 
     def __init__(self, layer, event, chunk_state, comp_stream, comm_stream, extra_args={}):
@@ -105,6 +108,9 @@ def release_state(self):
         if hasattr(self, 'moe_combine') and self.moe_combine is not None:
             del self.moe_combine
             self.moe_combine = None
+        if hasattr(self, 'post_combine') and self.post_combine is not None:
+            del self.post_combine
+            self.post_combine = None
         if hasattr(self, 'mtp_post_process') and self.mtp_post_process is not None:
             del self.mtp_post_process
             self.mtp_post_process = None
@@ -171,6 +177,7 @@ def create_node(stream, module, name):
             moe_dispatch_module,
             mlp_module,
             moe_combine_module,
+            post_combine_module,
             mtp_post_process_module,
         ) = fwd_callables
 
@@ -182,10 +189,12 @@ def create_node(stream, module, name):
             self.post_attn = create_node(comp_stream, post_attn_module, "post_attn")
             self.moe_dispatch = create_node(comm_stream, moe_dispatch_module, "moe_dispatch")
             self.moe_combine = create_node(comm_stream, moe_combine_module, "moe_combine")
+            self.post_combine = create_node(comm_stream, post_combine_module, "post_combine")
         else:
             self.post_attn = NoopScheduleNode()
             self.moe_dispatch = NoopScheduleNode()
             self.moe_combine = NoopScheduleNode()
+            self.post_combine = NoopScheduleNode()
 
         if is_mtp:
             self.mtp_post_process = create_node(
@@ -208,67 +217,101 @@ def get_fp8_context(self):
         )
 
     @staticmethod
-    def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False):
+    def run(
+        f_layer, 
+        b_layer, 
+        f_input=None, 
+        b_grad=None, 
+        is_last_layer_in_bwd=False,
+        early_comm_launch=False,
+    ):
         """Schedule one-forward-one-backward operations for a single transformer layer.
 
         This function interleaves forward and backward operations, overlapping the communications
         (dispatch or combine) of one with the computations (att or mlp) of the other
         to maximize parallelism and efficiency.
 
+        Two execution modes are supported:
+
+        1) Coarse-grained overlap (early_comm_launch = False):
+
         When f_layer and b_layer are not None, forward and backward pass are overlapped as follows:
         comm_stream: combine_bwd            | dispatch_fwd->dispatch_bwd  | combine_fwd
         comp_stream: attn_fwd->post_attn_fwd| mlp_bwd->mlp_bwd_dw->mlp_fwd| post_attn_bwd->attn_bwd
         For MTP, mtp_post_process_fwd is executed after the combine_fwd in the comp_stream,
         and mtp_post_process_bwd is executed before the combine_bwd in the comp_stream.
 
+        2) Fine-grained overlap (early_comm_launch = True):
+
+        This mode further decomposes communication and computation into smaller
+        stages and interleaves them more aggressively, including post-processing
+        steps (e.g., MTP / post-combine hooks).
+
+        The execution timeline becomes:
+
+        comm_stream: combine_bwd                                 | dispatch_fwd | dispatch_bwd | combine_fwd                                         | PP_fwd     | PP_bwd        |
+        comp_stream: post_combine_bwd → attn_fwd → post_attn_fwd | mlp_bwd      | mlp_fwd      | mlp_bwd_dw → post_attn_bwd → post_combine_fwd       | attn_bwd   | attn_bwd_dw   |
+
         Args:
             f_layer (TransformerLayerSchedulePlan): Forward layer (for current microbatch)
             b_layer (TransformerLayerSchedulePlan): Backward layer (for previous microbatch)
             f_input (Tensor): Input for forward computation
             b_grad (Tensor): Gradient for backward computation
             is_last_layer_in_bwd (bool):
                 Whether the current layer is the last layer in the backward pass.
-
+            early_comm_launch (bool):
+                Enable fine-grained communication / computation overlap
         Returns:
             Functions or values for next iteration's computation
         """
 
         if b_layer is not None:
             b_grad = b_layer.mtp_post_process.backward(b_grad)
+            b_grad = b_layer.post_combine.backward(b_grad)
             b_grad = b_layer.moe_combine.backward(b_grad)
 
         if f_layer is not None:
             with f_layer.get_fp8_context():
                 f_input = f_layer.attn.forward(f_input)
                 f_input = f_layer.post_attn.forward(f_input)
 
+        if early_comm_launch:
+            if f_layer is not None:
+                with f_layer.get_fp8_context():
+                    f_input = f_layer.moe_dispatch.forward(f_input)
+
         if b_layer is not None:
             b_grad = b_layer.mlp.backward(b_grad)
 
-        if f_layer is not None:
-            with f_layer.get_fp8_context():
-                f_input = f_layer.moe_dispatch.forward(f_input)
+        if not early_comm_launch:
+            if f_layer is not None:
+                with f_layer.get_fp8_context():
+                    f_input = f_layer.moe_dispatch.forward(f_input)
 
         if b_layer is not None:
-            b_layer.mlp.backward_dw()
             b_grad = b_layer.moe_dispatch.backward(b_grad)
 
-        if b_layer is not None and b_layer.config.ep_overlap_early_attn_memory_release:
-            b_grad = b_layer.post_attn.backward(b_grad)
-            b_grad = b_layer.attn.backward(b_grad)
-
         if f_layer is not None:
             with f_layer.get_fp8_context():
                 f_input = f_layer.mlp.forward(f_input)
 
         if f_layer is not None:
             with f_layer.get_fp8_context():
                 f_input = f_layer.moe_combine.forward(f_input)
-                f_input = f_layer.mtp_post_process.forward(f_input)
 
-        if b_layer is not None and not b_layer.config.ep_overlap_early_attn_memory_release:
+        if b_layer is not None:
+            b_layer.mlp.backward_dw()
             b_grad = b_layer.post_attn.backward(b_grad)
-            b_grad = b_layer.attn.backward(b_grad)
+
+        if f_layer is not None:
+            with f_layer.get_fp8_context():
+                f_input = f_layer.post_combine.forward(f_input)
+                f_input = f_layer.mtp_post_process.forward(f_input)
+
+        # Delay the last attn_bwd in backward pass
+        # for overlapping with the p2p comm            
+        if b_layer is not None and not is_last_layer_in_bwd:
+                b_grad = b_layer.attn.backward(b_grad)
 
         # Delay the last attn_dw in backward pass (attn_dw of the first layer)
         # for overlapping with the p2p comm
@@ -489,6 +532,7 @@ def run(
         f_num_layers = f_schedule_plan.num_layers() if f_schedule_plan is not None else 0
         b_num_layers = b_schedule_plan.num_layers() if b_schedule_plan is not None else 0
         overlapped_layers = min(f_num_layers, b_num_layers)
+        early_comm_launch = os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1"
 
         f_layer = b_layer = None
         # combined forward and backward pass for overlapped layers
@@ -504,6 +548,7 @@ def run(
                 f_input=f_input,
                 b_grad=b_grad,
                 is_last_layer_in_bwd=(i == b_num_layers - 1),
+                early_comm_launch = early_comm_launch,
             )
             if i < b_num_layers - 1:
                 b_layer.release_state()
@@ -514,7 +559,11 @@ def run(
             b_layer = b_schedule_plan.pop_layer()
             torch.cuda.nvtx.range_push(f"layer_{b_schedule_plan.num_layers()}b")
             _, b_grad = TransformerLayerSchedulePlan.run(
-                None, b_layer, b_grad=b_grad, is_last_layer_in_bwd=(i == b_num_layers - 1)
+                None,
+                b_layer,
+                b_grad=b_grad,
+                is_last_layer_in_bwd=(i == b_num_layers - 1),
+                early_comm_launch = early_comm_launch,
             )
             if i < b_num_layers - 1:
                 b_layer.release_state()
@@ -526,7 +575,12 @@ def run(
             torch.cuda.nvtx.range_push(f"layer_{i}f")
             if f_layer.layer.config.fine_grained_activation_offloading:
                 fine_grained_offloading_set_last_layer(i == f_num_layers - 1)
-            f_input, _ = TransformerLayerSchedulePlan.run(f_layer, None, f_input=f_input)
+            f_input, _ = TransformerLayerSchedulePlan.run(
+                f_layer,
+                None,
+                f_input=f_input,
+                early_comm_launch = early_comm_launch,
+            )
             torch.cuda.nvtx.range_pop()
 
         if f_schedule_plan is not None and post_forward is not None:
@@ -536,6 +590,12 @@ def run(
                 f_schedule_plan.wait_current_stream()
                 post_forward(f_input, f_schedule_plan.vp_stage)
 
+        # Delay the last attn_bwd in backward pass
+        # for overlapping with the p2p comm
+        if b_num_layers > 0:
+            assert b_layer is not None
+            b_grad = b_layer.attn.backward(b_grad)
+
         # post_backward()/send_backward_recv_backward() is running in the computation stream,
         # so the p2p comm could be overlapped with the wgrad of attn backward
         if b_schedule_plan is not None and post_backward is not None:

megatron/core/models/gpt/fine_grained_callables.py

@@ -65,6 +65,7 @@ def should_free_input(name, is_moe, enable_deepep, enable_hybridep):
     free_input_nodes = {
         "mlp": not enable_hybridep,
         "moe_combine": True,
+        "post_combine": enable_deepep,
         # For non-DeepEP and non-HybridEP dispatcher mode, the input is the un-dispatched tokens
         # and probs before dispatch A2A and it's not needed anymore after the forward pass
         # For DeepEP and HybridEP dispatcher mode, they are both needed in backward pass
@@ -410,7 +411,9 @@ def submodule_dispatch_forward(
         node.layer_state.dispatched_probs = node.detach(dispatched_probs)
         return dispatched_tokens
 
-    def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor):
+    def submodule_moe_forward(
+        node: ScheduleNode, dispatched_tokens: torch.Tensor
+    ):
         """
         Run forward pass for computations between dispatch and combine:
             post dispatch->experts->combine preprocess
@@ -436,27 +439,33 @@ def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor):
         # release tensor reference after use
         node.layer_state.dispatched_probs = None
         node.layer_state.pre_mlp_layernorm_output = None
-        if shared_expert_output is None:
-            # Return only expert_output, since shared_expert_output causes backward on None
-            return expert_output
-        return expert_output, shared_expert_output
+        if shared_expert_output is not None:
+        # Save shared_expert_output to layer state for later use in post_combine
+            node.layer_state.shared_expert_output = node.detach(shared_expert_output)
+        return expert_output
 
     def submodule_combine_forward(
-        node: ScheduleNode,
-        output: torch.Tensor,
-        shared_expert_output: Optional[torch.Tensor] = None,
+        node: ScheduleNode, output: torch.Tensor
     ):
         """
-        # Triggers token combine and the remaining computation in the transformer layer.
-        # The `mlp_bda` computation is placed after `mlp.combine` due to data dependency.
-        # This ordering is also critical for pipeline performance. Starting the `mlp.combine`
-        # communication at first allows it to be overlapped with computation from another
-        # microbatch. If `mlp_bda` were to run first, it would compete for SM resources
-        # with another microbatch's computation and expose the communication.
+        Triggers token combine communication.
+        This communication can be overlapped with computation from another microbatch.
         """
-        residual = node.layer_state.residual
+        output = layer.mlp.combine(output)
+        return output
 
-        output = layer.mlp.combine(output, shared_expert_output)
+    def submodule_post_combine_forward(
+        node: ScheduleNode, output: torch.Tensor
+    ):
+        """
+        Post-processes combined output and completes the transformer layer computation.
+        Adds shared expert output and performs bias-dropout-add operation.
+        """
+        residual = node.layer_state.residual
+        shared_expert_output = getattr(node.layer_state, 'shared_expert_output', None)
+
+        # Post-process combine and add shared expert output
+        output = layer.mlp.post_combine(output, shared_expert_output)
         mlp_output_with_bias = (output, None)
 
         with layer.bias_dropout_add_exec_handler():
@@ -474,8 +483,12 @@ def submodule_combine_forward(
         # Need to record residual to comm stream, since it's created on comp stream
         node.layer_state.residual.record_stream(torch.cuda.current_stream())
 
-        # release tensor reference after use
+        # release tensor references after use
+        if shared_expert_output is not None:
+            shared_expert_output.untyped_storage().resize_(0)
         node.layer_state.residual = None
+        node.layer_state.shared_expert_output = None
+
 
         # final layer norm from decoder
         final_layernorm = node.chunk_state.model.decoder.final_layernorm
@@ -498,8 +511,9 @@ def raise_not_implemented(*args):
     dispatch_func = submodule_dispatch_forward if is_moe else raise_not_implemented
     mlp_func = submodule_moe_forward if is_moe else mlp_wrapper
     combine_func = submodule_combine_forward if is_moe else raise_not_implemented
+    post_combine_func = submodule_post_combine_forward if is_moe else raise_not_implemented
 
-    forward_funcs = [attn_func, post_attn_func, dispatch_func, mlp_func, combine_func, None]
+    forward_funcs = [attn_func, post_attn_func, dispatch_func, mlp_func, combine_func, post_combine_func, None]
     backward_dw = {"attn": layer.self_attention, "mlp": layer.mlp}
     return forward_funcs, backward_dw
 
@@ -512,7 +526,7 @@ def build_mtp_layer_callables(layer):
     """
 
     forward_funcs, backward_dw = build_transformer_layer_callables(layer.transformer_layer)
-    attn_forward, post_attn_forward, dispatch_forward, mlp_forward, combine_forward, _ = (
+    attn_forward, post_attn_forward, dispatch_forward, mlp_forward, combine_forward, post_combine_forward, _ = (
         forward_funcs
     )
     is_moe = isinstance(layer.transformer_layer.mlp, MoELayer)
@@ -579,6 +593,7 @@ def rng_context_wrapper(func, *args, **kwargs):
     dispatch_func = partial(rng_context_wrapper, dispatch_forward)
     mlp_func = partial(rng_context_wrapper, mlp_forward)
     combine_func = partial(rng_context_wrapper, combine_forward)
+    post_combine_func = partial(rng_context_wrapper, post_combine_forward)
     mtp_post_process_func = submodule_mtp_postprocess_forward
 
     forward_funcs = [
@@ -587,6 +602,7 @@ def rng_context_wrapper(func, *args, **kwargs):
         dispatch_func,
         mlp_func,
         combine_func,
+        post_combine_func,
         mtp_post_process_func,
     ]
     backward_dw = {

megatron/core/transformer/moe/moe_layer.py

@@ -257,14 +257,21 @@ def routed_experts_compute(
 
         return output, mlp_bias
 
-    def combine(self, output: torch.Tensor, shared_expert_output: Optional[torch.Tensor]):
-        """Combines expert outputs via communication and adds shared expert output.
+    def combine(self, output: torch.Tensor):
+        """Combines expert outputs via communication.
 
         This method uses the token dispatcher to combine the outputs from different
-        experts (e.g., via an All-to-All communication). It then adds the output
-        from the shared expert if it exists.
+        experts (e.g., via an All-to-All communication).
         """
         output = self.token_dispatcher.token_combine(output)
+        return output
+
+    def post_combine(self, output: torch.Tensor, shared_expert_output: Optional[torch.Tensor]):
+        """Post-processes combined output and adds shared expert output.
+
+        This method applies post-processing to the combined expert outputs and
+        adds the output from the shared expert if it exists.
+        """
         output = self.token_dispatcher.combine_postprocess(output)
         if shared_expert_output is not None:
             output = output + shared_expert_output
@@ -291,7 +298,7 @@ def forward(self, hidden_states: torch.Tensor):
                 "are enabled without also enabling sequence parallelism."
             )
 
-        # MoE forward: route -> dispatch -> compute -> combine
+        # MoE forward: route -> dispatch -> compute -> combine -> post-combine
         def custom_forward(hidden_states):
             try:
                 shared_expert_output = self.shared_experts_compute(hidden_states)
@@ -307,7 +314,8 @@ def custom_forward(hidden_states):
 
             dispatched_input, probs = self.dispatch(hidden_states, probs)
             output, mlp_bias = self.routed_experts_compute(dispatched_input, probs, residual)
-            output = self.combine(output, shared_expert_output)
+            output = self.combine(output)
+            output = self.post_combine(output, shared_expert_output)
             return output, mlp_bias
 
         if self.moe_layer_recompute: