fix: make local CUDA graph outputs pipeline-safe

megatron/core/transformer/cuda_graphs.py

@@ -37,6 +37,7 @@
     is_te_min_version,
     log_on_each_pipeline_stage,
     log_single_rank,
+    make_viewless_tensor,
 )
 
 try:
@@ -1564,6 +1565,21 @@ def apply_cudagraph_record_metadata(self, args, kwargs, outputs):
         for t in self.get_tensors(outputs):
             _apply_cudagraph_buffer_metadata(t, is_output=True)
 
+    def _make_pipeline_output_viewless(self, output):
+        """Make a last-stage record or replay output safe for pipeline deallocation."""
+
+        if not (self.is_last_layer and self.deallocate_pipeline_outputs):
+            return output
+
+        return tree_map(
+            lambda value: (
+                make_viewless_tensor(inp=value, requires_grad=value.requires_grad, keep_graph=True)
+                if torch.is_tensor(value)
+                else value
+            ),
+            output,
+        )
+
     def record_graph_capture(self, args, kwargs):
         """Records the data needed to create this runner's forward cudagraph.
         The first pass records a graph and appends the runner to _CudagraphGlobalRecord.
@@ -1587,6 +1603,9 @@ def record_graph_capture(self, args, kwargs):
                 for i, o in enumerate(out)
             ]
         )
+        # Custom autograd Function outputs are views. Pipeline schedules may pseudo-deallocate
+        # this first-pass output before create_cudagraphs() switches the runner to replay mode.
+        out = self._make_pipeline_output_viewless(out)
 
         if not self.fwd_graph_recorded:
             logger.debug(f"Recording forward graph creation...")
@@ -1639,6 +1658,8 @@ def replay_graph_capture(self, is_first_microbatch, args, kwargs):
             func_args = inp_tensors
 
         out = _CudagraphReplayNode.apply(self, is_first_microbatch, *func_args)
+        # The replay node has the same custom-autograd output-view behavior as the record node.
+        out = self._make_pipeline_output_viewless(out)
 
         out_iter = iter(self.to_list(out))
         fwd_outputs = self.to_list(self.fwd_graph_outputs)

tests/unit_tests/transformer/test_local_cudagraph_pipeline.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+import pytest
+import torch
+
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec
+from megatron.core.pipeline_parallel.schedules import custom_backward, deallocate_output_tensor
+from megatron.core.tensor_parallel.random import (
+    HAVE_TE,
+    initialize_rng_tracker,
+    model_parallel_cuda_manual_seed,
+)
+from megatron.core.transformer.cuda_graphs import (
+    CudaGraphManager,
+    _CudagraphGlobalRecord,
+    create_cudagraphs,
+)
+from megatron.core.transformer.enums import CudaGraphModule
+from megatron.core.transformer.transformer_block import TransformerBlock
+from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.utils import is_te_min_version
+from tests.unit_tests.test_utilities import Utils
+
+
+@pytest.mark.skipif(
+    not (HAVE_TE and is_te_min_version("1.5.0")),
+    reason="use_te_rng_tracker requires TransformerEngine version >= 1.5",
+)
+class TestLocalCudagraphPipelineOutput:
+    def setup_method(self, method):
+        initialize_rng_tracker(use_te_rng_tracker=True, force_reset=True)
+        Utils.initialize_model_parallel(
+            tensor_model_parallel_size=2, pipeline_model_parallel_size=2
+        )
+        model_parallel_cuda_manual_seed(123)
+
+    def teardown_method(self, method):
+        Utils.destroy_model_parallel()
+        _CudagraphGlobalRecord.cudagraph_created = False
+        _CudagraphGlobalRecord.cudagraph_record = []
+        CudaGraphManager.global_mempool = None
+
+    def test_record_and_replay_outputs_support_pipeline_deallocation(self):
+        config = TransformerConfig(
+            num_layers=4,
+            hidden_size=64,
+            num_attention_heads=4,
+            attention_dropout=0.0,
+            hidden_dropout=0.0,
+            cuda_graph_impl="local",
+            cuda_graph_modules=[CudaGraphModule.attn],
+            cuda_graph_warmup_steps=1,
+            deallocate_pipeline_outputs=True,
+            use_cpu_initialization=True,
+        )
+        block = TransformerBlock(config, get_gpt_layer_with_transformer_engine_spec()).cuda()
+        block.train()
+        for param in block.parameters():
+            param.main_grad = torch.zeros_like(param)
+
+        sequence_length = 32
+        hidden_states = torch.randn(
+            (sequence_length, 1, config.hidden_size), device="cuda", requires_grad=True
+        )
+        attention_mask = torch.ones(
+            (1, 1, sequence_length, sequence_length), dtype=bool, device="cuda"
+        )
+
+        record_out = block(hidden_states=hidden_states, attention_mask=attention_mask)
+        expected_shape = record_out.shape
+        assert torch.isfinite(record_out).all()
+        assert record_out._base is None
+
+        record_grad = torch.ones_like(record_out)
+        deallocate_output_tensor(record_out, deallocate_pipeline_outputs=True)
+        custom_backward(record_out, record_grad)
+        assert hidden_states.grad is not None
+        assert torch.isfinite(hidden_states.grad).all()
+        hidden_states.grad = None
+        create_cudagraphs()
+
+        replay_out = block(hidden_states=hidden_states, attention_mask=attention_mask)
+        assert replay_out.shape == expected_shape
+        assert torch.isfinite(replay_out).all()
+        assert replay_out._base is None
+
+        replay_grad = torch.ones_like(replay_out)
+        deallocate_output_tensor(replay_out, deallocate_pipeline_outputs=True)
+        custom_backward(replay_out, replay_grad)
+        assert hidden_states.grad is not None
+        assert torch.isfinite(hidden_states.grad).all()
+
+        for layer in block.layers:
+            for runner in layer.cudagraph_manager.cudagraph_runners:
+                if hasattr(runner, "fwd_graph"):
+                    del runner.fwd_graph
+                if hasattr(runner, "bwd_graph"):
+                    del runner.bwd_graph
+        torch.cuda.synchronize()