diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py index d501c11a0a9..7ea84bffb25 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 @@ -40,6 +41,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 @@ -55,6 +57,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={}): @@ -137,6 +140,7 @@ def create_node(stream, module, name): moe_dispatch_module, mlp_module, moe_combine_module, + post_combine_module, mtp_post_process_module, ) = fwd_callables @@ -148,10 +152,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(comp_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( @@ -174,19 +180,46 @@ 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, + fine_grained_overlap=False, + post_forward=None, + post_backward=None, + f_schedule_plan=None, + b_schedule_plan=None, + is_last_layer=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 (fine_grained_overlap = 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 (fine_grained_overlap = 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) @@ -194,13 +227,25 @@ 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. - + fine_grained_overlap (bool): + Enable fine-grained communication / computation overlap + post_forward (callable or None): + The function to call after the forward pass + post_backward (callable or None): + The function to call after the backward pass + f_schedule_plan (TransformerModelChunkSchedulePlan): + The forward schedule plan + b_schedule_plan (TransformerModelChunkSchedulePlan): + The backward schedule plan + is_last_layer (bool): + Whether the current layer is the overlap boundary layer of the current chunk 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: @@ -208,34 +253,79 @@ 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 b_layer is not None: - b_grad = b_layer.mlp.backward(b_grad) + if fine_grained_overlap: + 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 b_layer is not None: + b_grad = b_layer.moe_dispatch.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) + + if b_layer is not None: + b_layer.mlp.backward_dw() + b_grad = b_layer.post_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) + + if is_last_layer: + if f_schedule_plan is not None and post_forward is not None: + f_schedule_plan.wait_current_stream() + post_forward(f_input, f_schedule_plan.vp_stage) + + if b_layer is not None: + b_grad = b_layer.attn.backward(b_grad) + + if is_last_layer: + if b_schedule_plan is not None and post_backward is not None: + b_schedule_plan.wait_current_stream() + post_backward(b_grad, b_schedule_plan.vp_stage) + + if b_layer is not None: + b_layer.attn.backward_dw() + else: + 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 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: + b_layer.mlp.backward_dw() + b_grad = b_layer.moe_dispatch.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.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 f_layer is not None: + with f_layer.get_fp8_context(): + f_input = f_layer.moe_combine.forward(f_input) + f_input = f_layer.post_combine.forward(f_input) + f_input = f_layer.mtp_post_process.forward(f_input) - if b_layer is not None: - b_grad = b_layer.post_attn.backward(b_grad) - b_grad = b_layer.attn.backward(b_grad) + if b_layer is not None: + b_grad = b_layer.post_attn.backward(b_grad) + 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 - if b_layer is not None and not is_last_layer_in_bwd: - b_layer.attn.backward_dw() + # Delay the last attn_dw in backward pass (attn_dw of the first layer) + # for overlapping with the p2p comm + if b_layer is not None and not is_last_layer_in_bwd: + b_layer.attn.backward_dw() return f_input, b_grad @@ -444,6 +534,9 @@ 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) + equal_layers = (f_num_layers == b_num_layers) + fine_grained_overlap = os.environ.get("CUDA_DEVICE_MAX_CONNECTIONS") == "1" + # combined forward and backward pass for overlapped layers for i in range(overlapped_layers): @@ -456,6 +549,12 @@ def run( f_input=f_input, b_grad=b_grad, is_last_layer_in_bwd=(i == b_num_layers - 1), + fine_grained_overlap = fine_grained_overlap, + post_forward=post_forward if fine_grained_overlap else None, + post_backward=post_backward if fine_grained_overlap else None, + f_schedule_plan=f_schedule_plan if fine_grained_overlap else None, + b_schedule_plan=b_schedule_plan if fine_grained_overlap else None, + is_last_layer=(i == overlapped_layers - 1 and equal_layers) if fine_grained_overlap else False, ) torch.cuda.nvtx.range_pop() @@ -464,7 +563,11 @@ def run( b_layer = b_schedule_plan.get_layer(b_num_layers - 1 - i) torch.cuda.nvtx.range_push(f"layer_{b_num_layers - 1 - i}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), + fine_grained_overlap = fine_grained_overlap, ) torch.cuda.nvtx.range_pop() @@ -472,26 +575,33 @@ def run( for i in range(overlapped_layers, f_num_layers): f_layer = f_schedule_plan.get_layer(i) torch.cuda.nvtx.range_push(f"layer_{i}f") - f_input, _ = TransformerLayerSchedulePlan.run(f_layer, None, f_input=f_input) + f_input, _ = TransformerLayerSchedulePlan.run( + f_layer, + None, + f_input=f_input, + fine_grained_overlap = fine_grained_overlap, + ) torch.cuda.nvtx.range_pop() - if f_schedule_plan is not None and post_forward is not None: - # 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() - 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() - post_backward(b_grad, b_schedule_plan.vp_stage) - - # Delay the last attn_dw in backward pass (attn_dw of the first layer) - # for overlapping with the p2p comm - if b_num_layers > 0: - b_schedule_plan.get_layer(0).attn.backward_dw() + if (not equal_layers and fine_grained_overlap) or not fine_grained_overlap: + if f_schedule_plan is not None and post_forward is not None: + # 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() + 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() + post_backward(b_grad, b_schedule_plan.vp_stage) + + if not fine_grained_overlap: + # Delay the last attn_dw in backward pass (attn_dw of the first layer) + # for overlapping with the p2p comm + if b_num_layers > 0: + b_schedule_plan.get_layer(0).attn.backward_dw() # post process forward if f_schedule_plan is not None and f_schedule_plan.post_process is not None: diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index fd1cc3d33c6..e5dc93a653e 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -384,7 +384,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 @@ -410,27 +412,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(): @@ -444,8 +452,11 @@ 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 return output def mlp_wrapper(node: ScheduleNode, *args, **kwargs): @@ -462,8 +473,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 @@ -476,7 +488,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) @@ -551,6 +563,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 = [ @@ -559,6 +572,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 0b0ee9b6850..ca9c8955b0c 100644 --- a/megatron/core/transformer/moe/moe_layer.py +++ b/megatron/core/transformer/moe/moe_layer.py @@ -268,14 +268,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) # Project the output back from latent dimension to hidden dimension after combine # in latent dimension. @@ -306,14 +313,15 @@ 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): shared_expert_output = self.shared_experts_compute(hidden_states) hidden_states, probs, residual = self.router_and_preprocess(hidden_states) dispatched_input, probs = self.dispatch(hidden_states, probs) output, mlp_bias = self.routed_experts_compute(dispatched_input, probs, residual) assert mlp_bias is None, f"mlp_bias is not supported for {type(self.token_dispatcher)}" - output = self.combine(output, shared_expert_output) + output = self.combine(output) + output = self.post_combine(output, shared_expert_output) return output, mlp_bias