MoE Overlap: Add 2 CUDA Events to Synchronize Computation & Communication Operations for Forward and Backward Respectively

megatron/core/models/common/model_chunk_schedule_plan.py

@@ -57,14 +57,16 @@ class TransformerLayerSchedulePlan:
     moe_combine = None
     mtp_post_process = None
 
-    def __init__(self, layer, event, chunk_state, comp_stream, comm_stream, extra_args={}):
+    def __init__(self, layer, fwd_event, back_event, chunk_state, comp_stream, comm_stream, extra_args={}):
         """Initializes a transformer layer schedule plan.
 
         Args:
             layer (TransformerLayer):
                 split a transformer layer into multiple nodes for fine-grained scheduling.
-            event (torch.cuda.Event):
-                record CUDA event across multiple nodes on different streams for synchronization.
+            fwd_event (torch.cuda.Event):
+                record CUDA event across multiple nodes on different streams for synchronization with regard to forward pass.
+            fwd_event (torch.cuda.Event):
+                record CUDA event across multiple nodes on different streams for synchronization with regard to backward pass.
             chunk_state (ModelChunkState): model state shared in the model chunk.
             comp_stream (torch.cuda.Stream): CUDA stream for computation.
             comm_stream (torch.cuda.Stream): CUDA stream for communication.
@@ -78,14 +80,15 @@ def __init__(self, layer, event, chunk_state, comp_stream, comm_stream, extra_ar
         self.layer_state = TransformerLayerState()
         self.chunk_state = chunk_state
         self.layer = layer
-        self.event = event
+        self.fwd_event = fwd_event
+        self.back_event = back_event
         self.comp_stream = comp_stream
         self.comm_stream = comm_stream
 
         # get callable nodes for transformer/mtp layer
-        self._build_callable_nodes(event, comp_stream, comm_stream, extra_args)
+        self._build_callable_nodes(fwd_event, back_event, comp_stream, comm_stream, extra_args)
 
-    def _build_callable_nodes(self, event, comp_stream, comm_stream, extra_args):
+    def _build_callable_nodes(self, fwd_event, back_event, comp_stream, comm_stream, extra_args):
         """
         Builds the callable nodes for the transformer/mtp layer:
             attn, post_attn, mlp, moe_dispatch and moe_combine, and mtp_post_process.
@@ -122,7 +125,8 @@ def create_node(stream, module, name):
             bwd_dw_callables = bwd_dw_callable_map.get(name, None)
             return TransformerLayerNode(
                 stream,
-                event,
+                fwd_event,
+                back_event,
                 self.layer_state,
                 self.chunk_state,
                 module,
@@ -292,7 +296,8 @@ def __init__(
 
         self._model_chunk_state = ModelChunkState()
         self._transformer_layers = []
-        self._event = torch.cuda.Event()
+        self._fwd_event = torch.cuda.Event()
+        self._back_event = torch.cuda.Event()
         self.pre_process = None
         self.post_process = None
         self.vp_stage = model.vp_stage
@@ -320,12 +325,12 @@ def __init__(
         mtp_num_layers = get_mtp_num_layers_to_build(model.config, vp_stage=self.vp_stage)
 
         # build preprocess
-        self.pre_process = PreProcessNode(model, self._model_chunk_state, self._event, comp_stream)
+        self.pre_process = PreProcessNode(model, self._model_chunk_state, self._fwd_event, self._back_event, comp_stream)
         # build layer schedule plan for each layer
         for layer_idx in range(transformer_num_layers):
             layer = model.decoder._get_layer(layer_idx)
             layer_plan = TransformerLayerSchedulePlan(
-                layer, self._event, self._model_chunk_state, comp_stream, comm_stream
+                layer, self._fwd_event, self._back_event, self._model_chunk_state, comp_stream, comm_stream
             )
             self._transformer_layers.append(layer_plan)
 
@@ -337,31 +342,56 @@ def __init__(
             }
             layer = model.mtp.layers[layer_idx]
             layer_plan = TransformerLayerSchedulePlan(
-                layer, self.event, self.state, comp_stream, comm_stream, extra_args
+                layer, self._fwd_event, self._back_event, self.state, comp_stream, comm_stream, extra_args
             )
             self._transformer_layers.append(layer_plan)
 
         # build post process
         if model.post_process:
             self.post_process = PostProcessNode(
-                model, self._model_chunk_state, self._event, comp_stream
+                model, self._model_chunk_state, self._fwd_event, self._back_event, comp_stream
             )
 
     @property
-    def event(self):
-        """Gets the CUDA event for synchronization."""
-        return self._event
+    def event_fwd(self):
+        """Gets the CUDA event for forward synchronization."""
+        return self._fwd_event
+
+    @property
+    def event_back(self):
+        """Gets the CUDA event for backward synchronization."""
+        return self._back_event
 
     def record_current_stream(self):
         """Records the current CUDA stream in the event."""
         stream = torch.cuda.current_stream()
         self.event.record(stream)
 
+    def record_current_stream_fwd(self):
+        """Records the current CUDA stream in the event for forward synchronization."""
+        stream = torch.cuda.current_stream()
+        self.event_fwd.record(stream)
+
+    def record_current_stream_back(self):
+        """Records the current CUDA stream in the event for backward synchronization."""
+        stream = torch.cuda.current_stream()
+        self.event_back.record(stream)
+
     def wait_current_stream(self):
         """Waits for the event to complete on the current CUDA stream."""
         stream = torch.cuda.current_stream()
         self.event.wait(stream)
 
+    def wait_current_stream_fwd(self):
+        """Waits for the event to complete on the current CUDA stream in the forward pass."""
+        stream = torch.cuda.current_stream()
+        self.event_fwd.wait(stream)
+
+    def wait_current_stream_back(self):
+        """Waits for the event to complete on the current CUDA stream in the backward pass."""
+        stream = torch.cuda.current_stream()
+        self.event_back.wait(stream)
+
     def get_layer(self, i):
         """Gets the transformer layer at the specified index."""
         assert i < self.num_layers()
@@ -428,15 +458,15 @@ def run(
             # pp output send/receive sync
             if pre_forward is not None:
                 pre_forward(f_schedule_plan.vp_stage)
-            f_schedule_plan.record_current_stream()
+            f_schedule_plan.record_current_stream_fwd()
             f_input = f_schedule_plan.pre_process.forward()
 
         if b_schedule_plan:
-            b_schedule_plan.record_current_stream()
+            b_schedule_plan.record_current_stream_back()
             assert b_grad is not None
             if pre_backward is not None:
                 pre_backward(b_schedule_plan.vp_stage)
-                b_schedule_plan.record_current_stream()
+                b_schedule_plan.record_current_stream_back()
 
             if b_schedule_plan.post_process is not None:
                 b_grad = b_schedule_plan.post_process.backward(b_grad)
@@ -479,13 +509,13 @@ def run(
             # post_forward()/send_forward_recv_forward() is running in the communication stream,
             # so the p2p comm could be overlapped with the attn backward
             with torch.cuda.stream(get_comm_stream()):
-                f_schedule_plan.wait_current_stream()
+                f_schedule_plan.wait_current_stream_fwd()
                 post_forward(f_input, f_schedule_plan.vp_stage)
 
         # 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:
-            b_schedule_plan.wait_current_stream()
+            b_schedule_plan.wait_current_stream_back()
             post_backward(b_grad, b_schedule_plan.vp_stage)
 
         # Delay the last attn_dw in backward pass (attn_dw of the first layer)
@@ -501,12 +531,12 @@ def run(
             b_schedule_plan.pre_process.backward(b_grad)
 
         if f_schedule_plan:
-            f_schedule_plan.wait_current_stream()
+            f_schedule_plan.wait_current_stream_fwd()
         if b_schedule_plan:
             b_schedule_plan.wait_current_stream()
 
         # Release reference as early as possible, this helps avoid memory leak.
         if b_schedule_plan is not None:
-            b_schedule_plan.release_state()
+            b_schedule_plan.wait_current_stream_back()
 
         return f_input

megatron/core/models/gpt/fine_grained_callables.py

@@ -83,16 +83,17 @@ class PreProcessNode(ScheduleNode):
     before the main transformer layers.
     """
 
-    def __init__(self, gpt_model, chunk_state, event, stream):
+    def __init__(self, gpt_model, chunk_state, fwd_event, back_event, stream):
         """Initializes a preprocessing node.
 
         Args:
             gpt_model: The GPT model instance.
             chunk_state (TransformerChunkState): State shared within a chunk
-            event: CUDA event for synchronization.
+            fwd_event: CUDA event for forward synchronization.
+            back_event: CUDA event for backward synchronization.
             stream: CUDA stream for execution.
         """
-        super().__init__(weak_method(self.forward_impl), stream, event, name="pre_process")
+        super().__init__(weak_method(self.forward_impl), stream, fwd_event, back_event, name="pre_process")
         self.gpt_model = gpt_model
         self.chunk_state = chunk_state
 
@@ -136,16 +137,17 @@ class PostProcessNode(ScheduleNode):
     after the main transformer layers.
     """
 
-    def __init__(self, gpt_model, chunk_state, event, stream):
+    def __init__(self, gpt_model, chunk_state, fwd_event, back_event, stream):
         """Initializes a postprocessing node.
 
         Args:
             gpt_model: The GPT model instance.
             chunk_state (TransformerChunkState): State shared within a chunk
-            event: CUDA event for synchronization.
+            fwd_event: CUDA event for forward synchronization.
+            back_event: CUDA event for backward synchronization.
             stream: CUDA stream for execution.
         """
-        super().__init__(weak_method(self.forward_impl), stream, event, name="post_process")
+        super().__init__(weak_method(self.forward_impl), stream, fwd_event, back_event, name="post_process")
         self.gpt_model = gpt_model
         self.chunk_state = chunk_state
 
@@ -206,7 +208,8 @@ class TransformerLayerNode(ScheduleNode):
     def __init__(
         self,
         stream,
-        event,
+        fwd_event,
+        back_event,
         layer_state,
         chunk_state,
         submodule,
@@ -218,7 +221,8 @@ def __init__(
 
         Args:
             stream (torch.cuda.Stream): CUDA stream for execution
-            event (torch.cuda.Event): Synchronization event
+            fwd_event (torch.cuda.Event): Forward synchronization event
+            back_event (torch.cuda.Event): Backward synchronization event
             layer_state (TransformerLayerState): State shared within a layer
             chunk_state (TransformerChunkState): State shared within a chunk
             submodule (function): The submodule contain forward and dw function
@@ -236,7 +240,8 @@ def __init__(
         super().__init__(
             weak_method(self.forward_impl),
             stream,
-            event,
+            fwd_event,
+            back_event,
             weak_method(self.backward_impl),
             free_input=free_input,
             name=name,

megatron/core/pipeline_parallel/combined_1f1b.py

@@ -405,7 +405,7 @@ def forward_backward_step():
         from megatron.core.pipeline_parallel.schedules import forward_step_calc_loss
 
         loss_node = ScheduleNode(
-            loss_func, torch.cuda.current_stream(), f_schedule_plan.event, name="loss_func"
+            loss_func, torch.cuda.current_stream(), f_schedule_plan.event_fwd, f_schedule_plan.event_back, name="loss_func"
         )
         loss_func = loss_node.forward
         output_tensor, num_tokens = forward_step_calc_loss(

megatron/core/pipeline_parallel/utils.py

@@ -119,7 +119,8 @@ def __init__(
         self,
         forward_func: Callable,
         stream: torch.cuda.Stream,
-        event: torch.cuda.Event,
+        fwd_event: torch.cuda.Event,
+        back_event: torch.cuda.Event,
         backward_func: Optional[Callable] = None,
         free_input: bool = False,
         name: str = "schedule_node",
@@ -133,9 +134,12 @@ def __init__(
                 - 'compute' stream: Used for computational nodes like attention and experts.
                 - 'communicate' stream: Used for nodes that handle token communication,
                   such as token dispatch and combine operations in MoE layers.
-            event (torch.cuda.Event): The CUDA event used for synchronization. Each
-                microbatch within a model chunk shares the same event, which is used
-                to manage dependencies between nodes operating on different streams.
+            fwd_event (torch.cuda.Event): The CUDA event used for forward pass synchronization.
+                 It is used to manage dependencies between nodes operating on different streams 
+                 in forward pass. 
+            back_event (torch.cuda.Event): The CUDA event used for backward pass synchronization.
+                 It is used to manage dependencies between nodes operating on different streams 
+                 in backward pass. 
             backward_func (callable, optional): Function for the backward pass.
             free_input (bool): Flag to indicate if the input should be freed after the
                 forward pass.
@@ -145,7 +149,8 @@ def __init__(
         self.forward_func = forward_func
         self.backward_func = backward_func if backward_func else self.default_backward_func
         self.stream = stream
-        self.event = event
+        self.fwd_event = fwd_event
+        self.back_event = back_event
         self.free_input = free_input
         self.inputs = None
         self.outputs = None
@@ -170,7 +175,7 @@ def forward(self, inputs=()):
         return self._forward(*inputs)
 
     def _forward(self, *inputs):
-        with stream_acquire_context(self.stream, self.event):
+        with stream_acquire_context(self.stream, self.fwd_event):
             torch.cuda.nvtx.range_push(f"{self.name} forward")
             with torch.cuda.stream(self.stream):
                 self.inputs = [make_viewless(e).detach() if e is not None else None for e in inputs]
@@ -212,7 +217,7 @@ def backward(self, output_grad):
         return self._backward(*output_grad)
 
     def _backward(self, *output_grad):
-        with stream_acquire_context(self.stream, self.event):
+        with stream_acquire_context(self.stream, self.back_event):
             torch.cuda.nvtx.range_push(f"{self.name} backward")
             with torch.cuda.stream(self.stream):
                 outputs = self.output