fix(combined-1f1b): release loss-node input storage after combined backward

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) and tensor.is_cuda:
+            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