diff --git a/megatron/core/models/common/model_chunk_schedule_plan.py b/megatron/core/models/common/model_chunk_schedule_plan.py index d11e53d7fc2..f451942ffc2 100644 --- a/megatron/core/models/common/model_chunk_schedule_plan.py +++ b/megatron/core/models/common/model_chunk_schedule_plan.py @@ -123,14 +123,13 @@ def _build_callable_nodes(self, event, comp_stream, comm_stream, extra_args): # get flags for latter use is_mtp = isinstance(self.layer, MultiTokenPredictionLayer) - is_moe = ( - isinstance(self.layer.transformer_layer.mlp, MoELayer) - if is_mtp - else isinstance(self.layer.mlp, MoELayer) - ) + transformer_layer = self.layer.transformer_layer if is_mtp else self.layer + is_moe = isinstance(transformer_layer.mlp, MoELayer) + num_local_experts = transformer_layer.mlp.num_local_experts if is_moe else None extra_args["config"] = self.layer.config extra_args["is_moe"] = is_moe + extra_args["num_local_experts"] = num_local_experts extra_args["delay_wgrad_compute"] = self.layer.config.delay_wgrad_compute extra_args["is_mtp"] = is_mtp diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index 7cee9d2973c..b4fe64ee9bb 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -43,13 +43,14 @@ def wrapped_func(*args, **kwarg): @internal_api -def should_free_input(name, is_moe, config): +def should_free_input(name, is_moe, config, num_local_experts): """Determine if the node should free its input memory. Args: name: Node name is_moe: Whether it's a MoE model config: TransformerConfig object + num_local_experts: Number of local experts in MoE module Returns: bool: Whether to free input memory @@ -70,8 +71,19 @@ def should_free_input(name, is_moe, config): # when and how to free the input memory. # The input and output of A2A are not needed anymore after the forward pass, # so we can free the input memory after the forward pass. + + # When low precision fp8/4 is enabled, the casted tensors are saved and the + # original bf16 tensors are safe to be freed. + free_mlp = config.fp8 is not None or config.fp4 is not None + if not free_mlp: + # AlltoAll dispatcher with local_num_experts=1 and HybridEP both use identity + # operation for `dispatch_postprocess`, hence the mlp inputs will be directly + # passed to GroupedGemm and should be saved for backward pass. + free_mlp = num_local_experts > 1 or config.moe_token_dispatcher_type != "alltoall" + free_mlp = free_mlp and not enable_hybridep + free_input_nodes = { - "mlp": not enable_hybridep, + "mlp": free_mlp, "moe_combine": True, # 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 @@ -256,7 +268,8 @@ def __init__( config = extra_args.get("config", None) assert config is not None, "model config must be passed to TransformerLayerNode." is_moe = extra_args.get("is_moe", False) - free_input = should_free_input(name, is_moe, config) + num_local_experts = extra_args.get("num_local_experts", None) + free_input = should_free_input(name, is_moe, config, num_local_experts) self.delay_wgrad_compute = extra_args.get("delay_wgrad_compute", False) super().__init__( @@ -316,7 +329,7 @@ def backward_dw(self): """Computes the weight gradients for the transformer layer node.""" if not self.delay_wgrad_compute: return - with torch.cuda.nvtx.range(f"{self.name} wgrad"): + with self.stream_acquire_context(f"{self.name} wgrad"): for module in self.bwd_dw_callables: module.backward_dw() @@ -514,15 +527,15 @@ def submodule_dispatch_forward( token_dispatcher._comm_manager.token_probs = probs dispatched_tokens, dispatched_probs = layer.mlp.dispatch(local_tokens, probs) - node.layer_state.dispatched_probs = node.detach(dispatched_probs) - return dispatched_tokens + return dispatched_tokens, dispatched_probs - def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor): + def submodule_moe_forward( + node: ScheduleNode, dispatched_tokens: torch.Tensor, dispatched_probs: torch.Tensor + ): """ Run forward pass for computations between dispatch and combine: post dispatch->experts->combine preprocess """ - dispatched_probs = node.layer_state.dispatched_probs token_dispatcher = layer.mlp.token_dispatcher if enable_deepep or enable_hybridep: # update dispatched_probs to be detached version, prevents @@ -531,13 +544,16 @@ def submodule_moe_forward(node: ScheduleNode, dispatched_tokens: torch.Tensor): expert_output, _ = layer.mlp.routed_experts_compute(dispatched_tokens, dispatched_probs) + # For HybridEP, tokens_per_expert is generated on comm stream, as the input to + # `routed_experts_compute`, it needs to be recorded to comp stream. + if enable_hybridep: + tokens_per_expert = token_dispatcher._comm_manager.get_number_of_tokens_per_expert() + tokens_per_expert.record_stream(torch.cuda.current_stream()) + if layer.recompute_pre_mlp_layernorm: # discard the output of the pre-mlp layernorm and register the recompute # as a gradient hook of expert_output layer.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(expert_output) - # release tensor reference after use - node.layer_state.dispatched_probs = None - node.layer_state.pre_mlp_layernorm_output = None return expert_output diff --git a/megatron/core/pipeline_parallel/utils.py b/megatron/core/pipeline_parallel/utils.py index 03c5f01f443..695968e2443 100644 --- a/megatron/core/pipeline_parallel/utils.py +++ b/megatron/core/pipeline_parallel/utils.py @@ -116,16 +116,6 @@ def set_ideal_affinity_for_current_gpu(): ) -@contextmanager -def stream_acquire_context(stream, event): - """Stream acquire context""" - event.wait(stream) - try: - yield - finally: - event.record(stream) - - class NoopScheduleNode: """A placeholder node in the computation graph that simply passes through inputs and outputs. @@ -208,26 +198,21 @@ def forward(self, inputs=()): return self._forward(*inputs) def _forward(self, *inputs): - with stream_acquire_context(self.stream, self.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] - for i, input in enumerate(self.inputs): - if input is not None: - input.requires_grad = inputs[i].requires_grad + with self.stream_acquire_context(f"{self.name} forward"): + self.inputs = [make_viewless(e).detach() if e is not None else None for e in inputs] + for i, input in enumerate(self.inputs): + if input is not None: + input.requires_grad = inputs[i].requires_grad - data = tuple(self.inputs) - data = self.forward_func(*data) + data = tuple(self.inputs) + data = self.forward_func(*data) - if not isinstance(data, tuple): - data = make_viewless(data) - else: - data = tuple( - [make_viewless(e) if isinstance(e, torch.Tensor) else e for e in data] - ) + if not isinstance(data, tuple): + data = make_viewless(data) + else: + data = tuple([make_viewless(e) if isinstance(e, torch.Tensor) else e for e in data]) - self.output = data - torch.cuda.nvtx.range_pop() + self.output = data # Immediately frees input tensors after they are used for nodes # where inputs are no longer needed after computation. @@ -250,18 +235,15 @@ def backward(self, output_grad): return self._backward(*output_grad) def _backward(self, *output_grad): - with stream_acquire_context(self.stream, self.event): - torch.cuda.nvtx.range_push(f"{self.name} backward") - with torch.cuda.stream(self.stream): - outputs = self.output - if not isinstance(outputs, tuple): - outputs = (outputs,) - assert len(outputs) == len(output_grad), ( - f"{len(outputs)} of {type(outputs[0])} is not equal to " - f"{len(output_grad)} of {type(output_grad[0])}" - ) - output_grad = self.backward_func(outputs, output_grad) - torch.cuda.nvtx.range_pop() + with self.stream_acquire_context(f"{self.name} backward"): + outputs = self.output + if not isinstance(outputs, tuple): + outputs = (outputs,) + assert len(outputs) == len(output_grad), ( + f"{len(outputs)} of {type(outputs[0])} is not equal to " + f"{len(output_grad)} of {type(output_grad[0])}" + ) + output_grad = self.backward_func(outputs, output_grad) # output_grad maybe from another stream if output_grad: @@ -288,6 +270,32 @@ def get_grad(self): grad = grad[0] return grad + @contextmanager + def stream_acquire_context(self, name=None): + """Stream acquire context that handles event synchronization, + NVTX profiling, and stream context. + + This context manager consolidates: + 1. Event wait/record for synchronization between streams + 2. NVTX range for profiling (if name is provided) + 3. torch.cuda.stream context for execution on the specified stream + + Args: + stream: The CUDA stream to execute on + event: The CUDA event for synchronization + name: Optional name for NVTX range profiling + """ + self.event.wait(self.stream) + if name: + torch.cuda.nvtx.range_push(name) + try: + with torch.cuda.stream(self.stream): + yield + finally: + if name: + torch.cuda.nvtx.range_pop() + self.event.record(self.stream) + def _release_state(self): """Clear the state of the node""" self.inputs = None diff --git a/tests/unit_tests/transformer/test_submodule_callables.py b/tests/unit_tests/transformer/test_submodule_callables.py index 73059495c06..31bd3d18b80 100644 --- a/tests/unit_tests/transformer/test_submodule_callables.py +++ b/tests/unit_tests/transformer/test_submodule_callables.py @@ -79,10 +79,10 @@ def run_model_submodules_with_capture(model, input_tensors, microbatches): local_tokens, probs = attn(node, input_tensors[i]) # dispatch fwd - dispatched_tokens = dispatch(node, local_tokens, probs) + dispatched_tokens, dispatched_probs = dispatch(node, local_tokens, probs) # moe fwd - expert_output = moe(node, dispatched_tokens) + expert_output = moe(node, dispatched_tokens, dispatched_probs) # combine fwd hidden_states = combine(node, expert_output)