From c2cc562f9e240558497da077c177a0e320716269 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Thu, 21 May 2026 10:14:52 +0800 Subject: [PATCH 1/2] fix(combined-1f1b): release loss-node input storage after combined backward In combined_forward_backward_step(), the loss_node retains references to its input tensors (notably the fp32 loss tensor) after the loss backward has materialized the gradient. Those tensors stay live in GPU memory through the remainder of the combined model backward, holding on to roughly O(seq_len) bytes of fp32 storage per micro-batch. Capture loss_node.inputs, release the autograd state immediately so the node is dropped from the graph, and free the input tensor storage with untyped_storage().resize_(0) once the combined model backward has consumed the loss gradient. record_stream() guards against lingering async users on the current stream. --- .../core/pipeline_parallel/combined_1f1b.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/megatron/core/pipeline_parallel/combined_1f1b.py b/megatron/core/pipeline_parallel/combined_1f1b.py index 1f45f2dd04f..c860c855314 100644 --- a/megatron/core/pipeline_parallel/combined_1f1b.py +++ b/megatron/core/pipeline_parallel/combined_1f1b.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import contextlib from contextlib import nullcontext @@ -21,6 +21,17 @@ Shape = Union[List[int], torch.Size] +def _release_tensor_storage(tensors): + """Release tensor storage after all backward users are done.""" + if tensors is None: + return + + for tensor in tensors: + if isinstance(tensor, torch.Tensor): + tensor.record_stream(torch.cuda.current_stream()) + tensor.untyped_storage().resize_(0) + + def combined_1f1b_schedule_for_no_pipelining( forward_step_func, data_iterator, @@ -405,6 +416,7 @@ def forward_backward_step(): # backward preprocess, the same as the backward_step() unwrap_input_tensor_grad = False b_schedule_plan = None + loss_node_inputs_to_release = None if b_model is not None: # Retain the grad on the input_tensor. if not isinstance(b_input_tensor, list): @@ -432,6 +444,8 @@ def forward_backward_step(): # Backward pass for loss function torch.autograd.backward(b_output_tensor[0], grad_tensors=b_output_tensor_grad[0]) b_output_tensor_grad[0] = loss_node.get_grad() + loss_node_inputs_to_release = loss_node.inputs + loss_node._release_state() # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(), # otherwise do nothing extra at the outer level @@ -454,6 +468,7 @@ def forward_backward_step(): post_forward=post_forward, post_backward=post_backward, ) + _release_tensor_storage(loss_node_inputs_to_release) # forward post process num_tokens = None From 25ff5cc551947c67912c4c729cadb85c37190edc Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 27 May 2026 09:55:25 +0800 Subject: [PATCH 2/2] fix(combined-1f1b): guard record_stream on cuda tensors only record_stream() raises on CPU tensors; add an is_cuda guard before calling it in _release_tensor_storage. Addresses review comment on PR #4908. --- megatron/core/pipeline_parallel/combined_1f1b.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/pipeline_parallel/combined_1f1b.py b/megatron/core/pipeline_parallel/combined_1f1b.py index c860c855314..3caaa900468 100644 --- a/megatron/core/pipeline_parallel/combined_1f1b.py +++ b/megatron/core/pipeline_parallel/combined_1f1b.py @@ -27,7 +27,7 @@ def _release_tensor_storage(tensors): return for tensor in tensors: - if isinstance(tensor, torch.Tensor): + if isinstance(tensor, torch.Tensor) and tensor.is_cuda: tensor.record_stream(torch.cuda.current_stream()) tensor.untyped_storage().resize_(0)