diff --git a/megatron/core/transformer/cuda_graphs.py b/megatron/core/transformer/cuda_graphs.py index 782fd6bf14f..1520c5cb8e1 100644 --- a/megatron/core/transformer/cuda_graphs.py +++ b/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) diff --git a/tests/unit_tests/transformer/test_local_cudagraph_pipeline.py b/tests/unit_tests/transformer/test_local_cudagraph_pipeline.py new file mode 100644 index 00000000000..f9d32120b4a --- /dev/null +++ b/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()