diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py index d501c11a0a9..a04bf7e1aa1 100644 --- a/megatron/core/models/common/model_chunk_schedule_plan.py +++ b/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 diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index fd1cc3d33c6..fcb1c8a3fa2 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/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, diff --git a/megatron/core/pipeline_parallel/combined_1f1b.py b/megatron/core/pipeline_parallel/combined_1f1b.py index 232d9c8cd70..4c3442cee68 100644 --- a/megatron/core/pipeline_parallel/combined_1f1b.py +++ b/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( diff --git a/megatron/core/pipeline_parallel/utils.py b/megatron/core/pipeline_parallel/utils.py index fae8e5466da..677b11bbc21 100644 --- a/megatron/core/pipeline_parallel/utils.py +++ b/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