diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py index 486a498dd73..cf330394eb5 100644 --- a/megatron/core/models/common/model_chunk_schedule_plan.py +++ b/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,19 +217,41 @@ 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) @@ -228,13 +259,15 @@ def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False) 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: @@ -242,21 +275,22 @@ def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False) 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) @@ -264,11 +298,20 @@ def run(f_layer, b_layer, f_input=None, b_grad=None, is_last_layer_in_bwd=False) 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: diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 60094976a9a..a9e454ab978 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/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 = { diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py index 10d10f667fe..a673747435c 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/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: