From a9be8ee25ed53e9bfd72cdcc5c7dd604d474a45b Mon Sep 17 00:00:00 2001 From: xiaoyao0115 <1804647152@qq.com> Date: Wed, 12 Aug 2026 02:06:31 -0700 Subject: [PATCH 01/16] Share CUDA graph memory across dynamic CP variants Signed-off-by: xiaoyao0115 <1804647152@qq.com> --- tests/pytorch/test_cuda_graphs.py | 681 ++++++ .../dot_product_attention/context_parallel.py | 43 +- .../pytorch/csrc/extensions/pybind.cpp | 10 + transformer_engine/pytorch/graph.py | 1908 ++++++++++++++++- 4 files changed, 2609 insertions(+), 33 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 1b9e11792e..7ad712f5a5 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -2,6 +2,8 @@ # # See LICENSE for license information. +import gc +import weakref from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest @@ -22,6 +24,10 @@ is_bf16_available, ) from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import ( + _get_cp_p2p_transport_group, + set_cp_p2p_transport_group, +) import transformer_engine.pytorch.ops as te_ops from transformer_engine.common import recipe from utils import ModelConfig, reset_rng_states @@ -39,6 +45,33 @@ } +def test_cp_p2p_transport_group_override(): + class Group: + pass + + logical_group = Group() + transport_group = Group() + + assert _get_cp_p2p_transport_group(logical_group) == (logical_group, False) + set_cp_p2p_transport_group(logical_group, transport_group) + assert _get_cp_p2p_transport_group(logical_group) == (transport_group, True) + set_cp_p2p_transport_group(logical_group, None) + assert _get_cp_p2p_transport_group(logical_group) == (logical_group, False) + + set_cp_p2p_transport_group(logical_group, transport_group) + logical_group_ref = weakref.ref(logical_group) + del logical_group + gc.collect() + assert logical_group_ref() is None + + self_transport_group = Group() + self_transport_group_ref = weakref.ref(self_transport_group) + set_cp_p2p_transport_group(self_transport_group, self_transport_group) + del self_transport_group + gc.collect() + assert self_transport_group_ref() is None + + def nvfp4_vanilla(): nvfp4_recipe = recipe.NVFP4BlockScaling() nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() @@ -740,3 +773,651 @@ def test_make_graphed_callables_with_interleaved_pipeline_parallelism( **kwargs, ) assert_all_equal(outputs, graph_outputs) + + +def _slot(saved_arena, branch, io_arena, overlap=0, frame=0, warmup=0): + """Build one private graph-memory slot used by the focused tests below.""" + return (saved_arena, branch, io_arena, branch, overlap, frame, warmup, 0, 0) + + +@pytest.mark.parametrize("elements", (0, 4096), ids=("empty", "nonempty")) +def test_slot_memory_variants_share_one_backing(elements: int) -> None: + """Mutually exclusive variants must use identical slot storage and output addresses.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square().sum().unsqueeze(0) * 3.0 + + variants = 5 + module = Module().cuda() + samples = tuple( + (torch.ones(elements, device="cuda", requires_grad=True),) for _ in range(variants) + ) + order = [ + value for variant in reversed(range(variants)) for value in (variant + 1, -variant - 1) + ] + graphed = make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=order, + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _slot(0, variant, 1, warmup=variant) for variant in range(variants) + ), + ) + + try: + pool = graphed[0]._te_cuda_graph_allocator_pool + assert all(graph._te_cuda_graph_allocator_pool is pool for graph in graphed) + output_ptrs = [] + for graph in graphed: + # A physical slot is replayed by later logical microbatches after its matching + # backward has drained, so exercise two complete lifetimes per callable. + for _ in range(2): + inp = torch.randn(elements, device="cuda", requires_grad=True) + output = graph(inp) + output_ptrs.append(output.data_ptr()) + output.sum().backward() + torch.testing.assert_close(inp.grad, 6.0 * inp.detach()) + assert len(set(output_ptrs)) == 1 + finally: + reset_graphs(graphed) + + +def test_slot_memory_input_staging_respects_overlapping_liveness() -> None: + """Live microbatches must not overwrite inputs still needed by backward.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() + + module = Module().cuda() + samples = tuple((torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(2)) + graphed = make_graphed_callables( + (module,), + samples, + num_warmup_iters=2, + _order=[1, 1, -1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=( + _slot(0, 0, 0, warmup=0), + _slot(1, 1, 1, warmup=0), + ), + ) + + try: + inp0 = torch.full((4096,), 2.0, device="cuda", requires_grad=True) + inp1 = torch.full((4096,), 3.0, device="cuda", requires_grad=True) + out0 = graphed[0](inp0) + out1 = graphed[1](inp1) + out0.sum().backward() + out1.sum().backward() + torch.testing.assert_close(inp0.grad, 2.0 * inp0.detach()) + torch.testing.assert_close(inp1.grad, 2.0 * inp1.detach()) + finally: + reset_graphs(graphed) + + +@pytest.mark.parametrize("reverse_replay", (False, True), ids=("forward", "reverse")) +def test_slot_memory_checkpoint_reuses_lockstep_branches(reverse_replay) -> None: + """Lockstep CP branches must restore one live slot boundary between captures.""" + + class Module(torch.nn.Module): + def forward(self, inp): + transient = torch.cat((inp.square(), inp.sin(), inp.cos()), dim=0) + return transient[: inp.numel()] * 3.0 + + variants = 5 + elements = 4096 + module = Module().cuda() + samples = tuple( + (torch.ones(elements, device="cuda", requires_grad=True),) for _ in range(variants) + ) + graphed = make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[ + *(variant + 1 for variant in range(variants)), + *(-variant - 1 for variant in range(variants)), + ], + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _slot(0, variant, 1, warmup=variant) for variant in range(variants) + ), + ) + + try: + pool = graphed[0]._te_cuda_graph_allocator_pool + assert all(graph._te_cuda_graph_allocator_pool is pool for graph in graphed) + replay_order = [1, 0, 2, 3, 4] if reverse_replay else range(variants) + for variant in replay_order: + graph = graphed[variant] + inp = torch.randn(elements, device="cuda", requires_grad=True) + output = graph(inp) + output.sum().backward() + expected = 6.0 * inp.detach() + if not torch.allclose(inp.grad, expected): + print( + "CHECKPOINT_GRAD_MISMATCH", + { + "reverse_replay": reverse_replay, + "variant": variant, + "input_ptr": inp.data_ptr(), + "output_ptr": output.data_ptr(), + "grad_ptr": inp.grad.data_ptr(), + "grad_head": inp.grad[:4].tolist(), + "expected_head": expected[:4].tolist(), + "grad_head_i64": inp.grad[:4].view(torch.int64).tolist(), + }, + flush=True, + ) + torch.testing.assert_close( + inp.grad, + expected, + msg=lambda message: f"variant={variant}: {message}", + ) + finally: + reset_graphs(graphed) + + +def test_slot_memory_checkpoint_reclaims_retained_branch_outputs() -> None: + """A module-held source output must not pin a mutually exclusive branch allocation.""" + + class Module(torch.nn.Module): + def __init__(self): + super().__init__() + self.retained_output = None + + def forward(self, inp): + self.retained_output = inp.square() * 3.0 + return self.retained_output + + variants = 5 + elements = 4096 + module = Module().cuda() + samples = tuple( + (torch.ones(elements, device="cuda", requires_grad=True),) for _ in range(variants) + ) + graphed = make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[ + *(variant + 1 for variant in range(variants)), + *(-variant - 1 for variant in range(variants)), + ], + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _slot(0, variant, 1, warmup=variant) for variant in range(variants) + ), + ) + + try: + output_ptrs = [] + for graph in graphed: + inp = torch.randn(elements, device="cuda", requires_grad=True) + output = graph(inp) + output_ptrs.append(output.data_ptr()) + output.sum().backward() + torch.testing.assert_close(inp.grad, 6.0 * inp.detach()) + assert len(set(output_ptrs)) == 1 + finally: + module.retained_output = None + reset_graphs(graphed) + + +def test_slot_memory_checkpoint_tracks_native_saved_storage() -> None: + """Checkpoint branches must retain allocator owners hidden in autograd saved tensors.""" + + class Module(torch.nn.Module): + def forward(self, inp): + hidden = inp * 2.0 + return hidden.square() + + variants = 5 + module = Module().cuda() + samples = tuple((torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(variants)) + graphed = make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + _order=[ + *(variant + 1 for variant in range(variants)), + *(-variant - 1 for variant in range(variants)), + ], + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _slot(0, variant, 1, warmup=variant) for variant in range(variants) + ), + ) + + try: + for graph in graphed: + inp = torch.randn(4096, device="cuda", requires_grad=True) + graph(inp).sum().backward() + torch.testing.assert_close(inp.grad, 8.0 * inp.detach()) + finally: + reset_graphs(graphed) + + +def test_slot_memory_fork_reuses_native_saved_allocations() -> None: + """Forked branches must reuse native saved-tensor allocations, not grow the pool.""" + + class Module(torch.nn.Module): + def __init__(self, canonical): + super().__init__() + self.canonical = canonical + + def forward(self, inp): + if self.canonical: + return inp.square() + hidden = inp.sin() + return hidden.square() + + def capture(variants): + microbatches = 2 + modules = tuple(Module(variant == 0).cuda() for variant in range(variants)) + samples = tuple( + (torch.ones(4096, device="cuda", requires_grad=True),) + for _ in range(variants * microbatches) + ) + variant_group = [variant + 1 for variant in range(variants)] + graphed = make_graphed_callables( + modules, + samples, + num_warmup_iters=2, + _order=[ + *variant_group, + *(-value for value in variant_group), + *variant_group, + *(-value for value in variant_group), + ], + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _slot( + microbatch, + variant * microbatches + microbatch, + microbatch, + warmup=variant, + ) + for variant in range(variants) + for microbatch in range(microbatches) + ), + ) + pool = graphed[0]._te_cuda_graph_allocator_pool + snapshot = pool.snapshot(include_traces=False) + segments = snapshot["segments"] if isinstance(snapshot, dict) else snapshot + return graphed, sum(segment["total_size"] for segment in segments) + + baseline, baseline_pool_bytes = capture(2) + try: + for graph_idx, graph in enumerate(baseline): + inp = torch.randn(4096, device="cuda", requires_grad=True) + graph(inp).sum().backward() + expected = ( + 2.0 * inp.detach() + if graph_idx // 2 == 0 + else 2.0 * inp.detach().sin() * inp.detach().cos() + ) + torch.testing.assert_close(inp.grad, expected) + finally: + reset_graphs(baseline) + + graphed, forked_pool_bytes = capture(5) + try: + for graph_idx, graph in enumerate(graphed): + inp = torch.randn(4096, device="cuda", requires_grad=True) + graph(inp).sum().backward() + expected = ( + 2.0 * inp.detach() + if graph_idx // 2 == 0 + else 2.0 * inp.detach().sin() * inp.detach().cos() + ) + torch.testing.assert_close(inp.grad, expected) + assert forked_pool_bytes == baseline_pool_bytes + finally: + reset_graphs(graphed) + + +def test_slot_memory_checkpoint_aliases_parameter_gradients() -> None: + """Mutually exclusive branches must return parameter grads from one slot address.""" + + class Module(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(4096, device="cuda")) + + def forward(self, inp): + return inp * self.weight + + variants = 5 + module = Module() + samples = tuple((torch.ones_like(module.weight, requires_grad=True),) for _ in range(variants)) + graphed = make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[ + *(variant + 1 for variant in range(variants)), + *(-variant - 1 for variant in range(variants)), + ], + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _slot(0, variant, 1, warmup=variant) for variant in range(variants) + ), + ) + + try: + for graph in graphed: + module.weight.grad = None + inp = torch.randn_like(module.weight, requires_grad=True) + graph(inp).sum().backward() + torch.testing.assert_close(inp.grad, module.weight.detach()) + torch.testing.assert_close(module.weight.grad, inp.detach()) + finally: + module.weight.grad = None + reset_graphs(graphed) + + +def test_slot_memory_native_io_aliases_graph_pool_storage() -> None: + """Nine-field DCP slots default to forked graph-pool I/O aliases.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() * 3.0 + + variants = 5 + elements = 4096 + module = Module().cuda() + samples = tuple( + (torch.ones(elements, device="cuda", requires_grad=True),) for _ in range(variants) + ) + graphed = make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[ + *(variant + 1 for variant in range(variants)), + *(-variant - 1 for variant in range(variants)), + ], + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _slot(0, variant, 1, warmup=variant) for variant in range(variants) + ), + ) + + try: + assert not hasattr(graphed[0], "_te_cuda_graph_slot_memory_pool") + output_ptrs = [] + for graph in graphed: + inp = torch.randn(elements, device="cuda", requires_grad=True) + output = graph(inp) + output_ptrs.append(output.data_ptr()) + output.sum().backward() + torch.testing.assert_close(inp.grad, 6.0 * inp.detach()) + assert len(set(output_ptrs)) == 1 + finally: + reset_graphs(graphed) + + +def test_slot_memory_releases_outputs_before_next_forward_group() -> None: + """A completed backward group must not pin frame-local outputs into the next forward.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() * 3.0 + + variants = 3 + microbatches = 2 + elements = 4096 + module = Module().cuda() + samples = tuple( + (torch.ones(elements, device="cuda", requires_grad=True),) + for _ in range(variants * microbatches) + ) + variant_group = [variant + 1 for variant in range(variants)] + order = [ + *variant_group, + *(-value for value in variant_group), + *variant_group, + *(-value for value in variant_group), + ] + slots = tuple( + _slot( + microbatch, + variant * microbatches + microbatch, + microbatch, + warmup=variant, + ) + for variant in range(variants) + for microbatch in range(microbatches) + ) + graphed = make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=order, + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=slots, + ) + + try: + for graph in graphed: + inp = torch.randn(elements, device="cuda", requires_grad=True) + graph(inp).sum().backward() + torch.testing.assert_close(inp.grad, 6.0 * inp.detach()) + finally: + reset_graphs(graphed) + + +def test_slot_memory_releases_transients_across_vpp_tail_backward() -> None: + """PP/VPP tail backward groups must not inherit transient owners from prior events.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() * 3.0 + + variants = 3 + model_chunks = 2 + microbatches = 4 + elements = 4096 + module = Module().cuda() + samples = tuple( + (torch.ones(elements, device="cuda", requires_grad=True),) + for _ in range(variants * model_chunks * microbatches) + ) + slots = [] + for variant in range(variants): + for model_chunk in range(model_chunks): + logical_chunk = variant * model_chunks + model_chunk + for microbatch in range(microbatches): + frame = model_chunk * microbatches + microbatch + branch = variant * model_chunks * microbatches + frame + slots.append( + _slot( + frame, + branch, + microbatch, + overlap=model_chunk, + frame=0, + warmup=logical_chunk, + ) + ) + + # PP=2, VPP=2, rank 0, four-microbatch schedule with lockstep CP branches. + base_order = [1, 1, 2, 2, 1, -2, 1, -2, 2, -1, 2, -1, -2, -2, -1, -1] + order = [] + for chunk_id in base_order: + for variant in range(variants): + remapped = abs(chunk_id) + variant * model_chunks + order.append(remapped if chunk_id > 0 else -remapped) + + graphed = make_graphed_callables( + (module,) * (variants * model_chunks), + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=order, + _num_layers_per_chunk=[1] * (variants * model_chunks), + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple(slots), + ) + + try: + inp = torch.randn(elements, device="cuda", requires_grad=True) + graphed[0](inp).sum().backward() + torch.testing.assert_close(inp.grad, 6.0 * inp.detach()) + finally: + reset_graphs(graphed) + + +def test_slot_memory_snapshots_live_inputs_across_slot_wrap() -> None: + """A drained slot can wrap while another slot's forward remains live.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() * 3.0 + + module = Module().cuda() + samples = tuple((torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(2)) + graphed = make_graphed_callables( + (module,), + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[1, 1, -1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 2), _slot(1, 0, 3)), + ) + + try: + inp0 = torch.randn(4096, device="cuda", requires_grad=True) + inp1 = torch.randn(4096, device="cuda", requires_grad=True) + out0 = graphed[0](inp0) + out1 = graphed[1](inp1) + out0.sum().backward() + # Returned input-grad surfaces are valid until their physical slot wraps. + torch.testing.assert_close(inp0.grad, 6.0 * inp0.detach()) + + inp2 = torch.randn(4096, device="cuda", requires_grad=True) + out2 = graphed[0](inp2) + out1.sum().backward() + torch.testing.assert_close(inp1.grad, 6.0 * inp1.detach()) + out2.sum().backward() + torch.testing.assert_close(inp2.grad, 6.0 * inp2.detach()) + finally: + reset_graphs(graphed) + + +def test_slot_memory_coalesces_overlapping_saved_views() -> None: + """Saved views of one storage should occupy only their byte union in each live slot.""" + + class OverlappingSaves(torch.autograd.Function): + @staticmethod + def forward(ctx, inp): + backing = torch.cat(tuple(inp + value for value in (1.0, 2.0, 3.0, 4.0))) + ctx.save_for_backward(backing[: 3 * inp.numel()], backing[inp.numel() :]) + ctx.input_elements = inp.numel() + return inp + 0.25 + + @staticmethod + def backward(ctx, grad_output): + first, second = ctx.saved_tensors + saved = (first[: ctx.input_elements] + second[: ctx.input_elements]) / 2.0 + return grad_output * saved + + class Module(torch.nn.Module): + def forward(self, inp): + return OverlappingSaves.apply(inp) + + elements = 4096 + module = Module().cuda() + samples = tuple((torch.ones(elements, device="cuda", requires_grad=True),) for _ in range(2)) + graphed = make_graphed_callables( + (module,), + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[1, 1, -1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 2), _slot(1, 0, 3)), + ) + + try: + inp0 = torch.randn(elements, device="cuda", requires_grad=True) + inp1 = torch.randn(elements, device="cuda", requires_grad=True) + out0 = graphed[0](inp0) + out1 = graphed[1](inp1) + out1.sum().backward() + out0.sum().backward() + torch.testing.assert_close(inp0.grad, inp0.detach() + 1.5) + torch.testing.assert_close(inp1.grad, inp1.detach() + 1.5) + finally: + reset_graphs(graphed) + + +def test_slot_memory_preserves_fused_wgrad_hook() -> None: + """Fused wgrad must retain the parameter's autograd edge during replay.""" + dtype = torch.bfloat16 + module = Linear( + 32, + 32, + params_dtype=dtype, + fuse_wgrad_accumulation=True, + device="cuda", + ) + module.weight.main_grad = torch.zeros_like(module.weight) + module.weight.grad_added_to_main_grad = False + samples = tuple( + (torch.randn(8, 32, device="cuda", dtype=dtype, requires_grad=True),) for _ in range(2) + ) + graphed = make_graphed_callables( + (module, module), + samples, + allow_unused_input=True, + _order=[2, -2, 1, -1], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 1), _slot(0, 1, 1, warmup=1)), + ) + + hook_calls = 0 + + def count_hook(grad): + nonlocal hook_calls + hook_calls += 1 + return grad + + hook = module.weight.register_hook(count_hook) + try: + for graph in graphed: + hook_calls = 0 + module.weight.grad = None + module.weight.main_grad.zero_() + inp = torch.randn(8, 32, device="cuda", dtype=dtype, requires_grad=True) + graph(inp).sum().backward() + torch.cuda.synchronize() + assert hook_calls == 1 + assert torch.count_nonzero(module.weight.main_grad) > 0 + finally: + hook.remove() + reset_graphs(graphed) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 030b1d9cdc..177f3b38a0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3,9 +3,10 @@ # See LICENSE for license information. """Context Parallelism.""" -import os import itertools -from typing import List, Union, Tuple +import os +import weakref +from typing import List, Tuple, Union import torch import transformer_engine_torch as tex @@ -54,6 +55,25 @@ _seq_chunk_ids_cache_for_reordering_before_attn = {} _seq_chunk_ids_cache_for_reordering_after_attn = {} _softmax_offset_chunk_ids_cache = {} +_cp_p2p_transport_groups = weakref.WeakKeyDictionary() + + +def set_cp_p2p_transport_group(cp_group, transport_group): + """Override only the P2P transport group for a logical CP group.""" + if transport_group is None: + _cp_p2p_transport_groups.pop(cp_group, None) + return + _cp_p2p_transport_groups[cp_group] = weakref.ref(transport_group) + + +def _get_cp_p2p_transport_group(cp_group): + transport_group_ref = _cp_p2p_transport_groups.get(cp_group) + if transport_group_ref is not None: + transport_group = transport_group_ref() + if transport_group is not None: + return transport_group, True + _cp_p2p_transport_groups.pop(cp_group, None) + return cp_group, False # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" @@ -64,36 +84,39 @@ def flash_attn_p2p_communicate( ): """Point-to-point communications of KV and dKV in Attention with context parallelism""" send_recv_ops = [] + transport_group, transport_overridden = _get_cp_p2p_transport_group(cp_group) + if transport_overridden: + batch_p2p_comm = True if batch_p2p_comm: if rank % 2 == 0: send_op = torch.distributed.P2POp( - torch.distributed.isend, send_tensor, send_dst, cp_group + torch.distributed.isend, send_tensor, send_dst, transport_group ) recv_op = torch.distributed.P2POp( - torch.distributed.irecv, recv_tensor, recv_src, cp_group + torch.distributed.irecv, recv_tensor, recv_src, transport_group ) send_recv_ops.append(send_op) send_recv_ops.append(recv_op) else: recv_op = torch.distributed.P2POp( - torch.distributed.irecv, recv_tensor, recv_src, cp_group + torch.distributed.irecv, recv_tensor, recv_src, transport_group ) send_op = torch.distributed.P2POp( - torch.distributed.isend, send_tensor, send_dst, cp_group + torch.distributed.isend, send_tensor, send_dst, transport_group ) send_recv_ops.append(recv_op) send_recv_ops.append(send_op) send_recv_reqs = torch.distributed.batch_isend_irecv(send_recv_ops) else: if rank % 2 == 0: - send_op = torch.distributed.isend(send_tensor, send_dst, cp_group) - recv_op = torch.distributed.irecv(recv_tensor, recv_src, cp_group) + send_op = torch.distributed.isend(send_tensor, send_dst, transport_group) + recv_op = torch.distributed.irecv(recv_tensor, recv_src, transport_group) send_recv_ops.append(send_op) send_recv_ops.append(recv_op) else: - recv_op = torch.distributed.irecv(recv_tensor, recv_src, cp_group) - send_op = torch.distributed.isend(send_tensor, send_dst, cp_group) + recv_op = torch.distributed.irecv(recv_tensor, recv_src, transport_group) + send_op = torch.distributed.isend(send_tensor, send_dst, transport_group) send_recv_ops.append(recv_op) send_recv_ops.append(send_op) send_recv_reqs = send_recv_ops diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 18da5d0e9f..cac5fec488 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -12,6 +12,9 @@ #include #include +#include +#include + #include #include #include @@ -135,6 +138,13 @@ void init_extension() { PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { NVTE_DECLARE_COMMON_PYBIND11_HANDLES(m) + m.def("_graph_checkpoint_detach_storage", [](uintptr_t storage_impl_ptr) { + auto *storage = reinterpret_cast(storage_impl_ptr); + const auto &data_ptr = storage->data_ptr(); + NVTE_CHECK(data_ptr.get_deleter() == &c10::detail::deleteNothing, + "CUDA graph checkpoint storage must have a no-op deleter before detaching"); + storage->set_data_ptr_noswap(at::DataPtr(data_ptr.get(), data_ptr.device())); + }); m.def("quantize", transformer_engine::pytorch::quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("output") = py::none(), py::arg("noop") = py::none()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 86b8a4acf4..4d99d12215 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -3,14 +3,17 @@ # See LICENSE for license information. """Functions for CUDA Graphs support in FP8""" + from collections.abc import Iterable import contextlib import gc +import os import warnings from math import ceil -from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union import torch +import transformer_engine_torch as tex from torch.utils._pytree import tree_flatten as _tree_flatten from torch.utils._pytree import tree_unflatten as _tree_unflatten from torch._C import _graph_pool_handle @@ -23,7 +26,7 @@ get_default_fp8_recipe, ) from .distributed import get_all_rng_states, graph_safe_rng_available -from .module.base import TransformerEngineBaseModule +from .module.base import TransformerEngineBaseModule, get_dummy_wgrad from .ops.op import BasicOperation from .ops import Sequential from .ops.fuser import OperationFuser @@ -38,6 +41,78 @@ SingleOrTuple = Union[_T, Tuple[_T, ...]] +def _tensor_storage_ptr(tensor: torch.Tensor) -> int: + """Return the base storage pointer used to recognize static graph inputs.""" + return tensor.untyped_storage().data_ptr() + + +def _tensor_version(tensor: torch.Tensor) -> Optional[int]: + """Return the mutation version when the tensor tracks one.""" + try: + return tensor._version + except RuntimeError: + return None + + +def _saved_tensor_signature(tensor: torch.Tensor) -> Tuple[Any, ...]: + """Describe the layout needed to reproduce a tensor in a static arena.""" + if tensor.layout != torch.strided: + raise RuntimeError( + "CUDA graph saved-tensor arenas only support strided tensors, " + f"but got layout={tensor.layout}." + ) + if any(stride < 0 for stride in tensor.stride()): + raise RuntimeError("CUDA graph saved-tensor arenas do not support negative strides.") + + if tensor.numel() == 0: + storage_numel = 0 + else: + storage_numel = 1 + sum( + (size - 1) * stride for size, stride in zip(tensor.shape, tensor.stride()) + ) + return ( + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + tensor.device, + tensor.requires_grad, + storage_numel * tensor.element_size(), + ) + + +def _input_staging_key(tensor: torch.Tensor) -> Tuple[Any, ...]: + """Describe user inputs that can use one forward-only staging surface.""" + return (tensor.layout, tensor.storage_offset(), *_saved_tensor_signature(tensor)) + + +def _align_up(value: int, alignment: int = 256) -> int: + """Align byte offsets for typed tensor views into a uint8 arena.""" + return (value + alignment - 1) // alignment * alignment + + +def _io_tensor_plan(tensor: Any, kind: str) -> Optional[Tuple[Any, ...]]: + """Return an arena plan for plain CUDA tensors exposed across graph boundaries.""" + if ( + type(tensor) is not torch.Tensor + or not tensor.is_cuda + or tensor.layout != torch.strided + or any(stride < 0 for stride in tensor.stride()) + ): + return None + return (kind, None, *_saved_tensor_signature(tensor)) + + +def _arena_view(arena: torch.Tensor, offset: int, spec: Tuple[Any, ...]) -> torch.Tensor: + """Materialize a typed tensor view at a byte offset in an arena.""" + target = torch.empty((0,), dtype=spec[4], device=spec[5]) + return target.set_( + arena.untyped_storage(), + (arena.storage_offset() + offset) // target.element_size(), + spec[2], + spec[3], + ) + + def set_capture_start() -> None: """Record beginning of `make_graphed_callables`.""" global _IS_GRAPH_CAPTURING @@ -108,6 +183,7 @@ def _make_graphed_callables( pool: Optional[Tuple[int, ...]] = None, retain_graph_in_backward: bool = False, _reuse_graph_input_output_buffers: bool = False, + _graph_memory_slots: Optional[Sequence[Tuple[int, ...]]] = None, pre_warmup_hook: Optional[Callable] = None, post_warmup_hook: Optional[Callable] = None, ) -> SingleOrTuple[Callable]: @@ -251,10 +327,29 @@ def _make_graphed_callables( f"for {len(sample_args)} sample_args" ) - # Check reuse graph conditions and reorganize sample_args and sample_kwargs. - # Note: When capturing a graph, we hold onto the args and kwargs so we have static buffers - # when the graph is replayed. If two model chunk microbatches have no overlap between their - # forward and backward, then we can reduce memory usage by reusing the same static buffers. + use_slot_memory = _graph_memory_slots is not None + if use_slot_memory: + required_checkpoint_apis = ( + "_cuda_getCheckpointState", + "_cuda_setCheckpointPoolState", + "_cuda_checkPoolLiveAllocations", + "_free_And_Remove_DeleterFn", + ) + missing_checkpoint_apis = [ + name for name in required_checkpoint_apis if not hasattr(torch._C, name) + ] + if missing_checkpoint_apis: + raise RuntimeError( + "CUDA graph slot-branch checkpointing requires PyTorch allocator APIs " + f"{missing_checkpoint_apis}." + ) + if not hasattr(tex, "_graph_checkpoint_detach_storage"): + raise RuntimeError( + "CUDA graph slot-branch checkpointing requires " + "transformer_engine_torch._graph_checkpoint_detach_storage." + ) + + _reuse_graph_input_buffers = _reuse_graph_input_output_buffers and not use_slot_memory if _reuse_graph_input_output_buffers: if _order is None: raise ValueError( @@ -264,6 +359,43 @@ def _make_graphed_callables( raise RuntimeError( "`_reuse_graph_input_output_buffers` is only available in training mode." ) + + saved_tensor_memory_alias_groups = None + saved_tensor_memory_families = None + slot_io_memory_alias_groups = None + slot_io_liveness_groups = None + warmup_plan_alias_groups = None + if use_slot_memory: + if _order is None or not is_training or not _reuse_graph_input_output_buffers: + raise RuntimeError( + "Graph-memory slots require a training graph with `_order` and graph buffer reuse." + ) + if pool is not None: + raise ValueError("Graph-memory slots create and own their CUDA graph memory pool.") + if not hasattr(torch.cuda, "MemPool"): + raise RuntimeError("Graph-memory slots require torch.cuda.MemPool support.") + if len(_graph_memory_slots) != len(sample_args): + raise ValueError( + f"Expected {len(sample_args)} graph-memory slots, got {len(_graph_memory_slots)}." + ) + if any( + not isinstance(slot, tuple) + or len(slot) != 9 + or not all(isinstance(value, int) for value in slot) + for slot in _graph_memory_slots + ): + raise TypeError("Each graph-memory slot must be a tuple of nine integers.") + saved_tensor_memory_alias_groups = [(slot[0], slot[1]) for slot in _graph_memory_slots] + saved_tensor_memory_families = [(slot[2], slot[4], slot[5]) for slot in _graph_memory_slots] + slot_io_memory_alias_groups = [(slot[2], slot[3]) for slot in _graph_memory_slots] + slot_io_liveness_groups = [(slot[4], slot[5]) for slot in _graph_memory_slots] + warmup_plan_alias_groups = [slot[6] for slot in _graph_memory_slots] + + # Check reuse graph conditions and reorganize sample_args and sample_kwargs. + # Note: When capturing a graph, we hold onto the args and kwargs so we have static buffers + # when the graph is replayed. If two model chunk microbatches have no overlap between their + # forward and backward, then we can reduce memory usage by reusing the same static buffers. + if _reuse_graph_input_buffers: if isinstance(sample_args, tuple): sample_args = list(sample_args) if isinstance(sample_kwargs, tuple): @@ -403,14 +535,50 @@ def _make_graphed_callables( graph_callables = [None for _ in range(len(flatten_sample_args))] # For cases with multiple active RNG states, e.g. TP. - if graph_safe_rng_available(): + if graph_safe_rng_available() and not bool( + int(os.getenv("NVTE_DISABLE_GRAPH_SAFE_RNG_REGISTRATION", "0")) + ): for _, state in get_all_rng_states().items(): for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs): fwd_graph.register_generator_state(state) bwd_graph.register_generator_state(state) bwd_dw_graph.register_generator_state(state) - mempool = graph_pool_handle() if pool is None else pool + allocator_settings_to_apply = None + allocator_settings_to_restore = None + if use_slot_memory: + allocator_conf = os.getenv("PYTORCH_CUDA_ALLOC_CONF") or os.getenv("PYTORCH_ALLOC_CONF", "") + allocator_parts = [part.strip() for part in allocator_conf.split(",") if part.strip()] + expandable_enabled = any( + part.split(":", 1)[0].strip() == "expandable_segments" + and part.split(":", 1)[1].strip().lower() == "true" + for part in allocator_parts + if ":" in part + ) + if expandable_enabled: + allocator_settings_setter = getattr(torch._C, "_accelerator_setAllocatorSettings", None) + if allocator_settings_setter is None: + raise RuntimeError( + "Temporarily disabling expandable segments during CUDA graph capture " + "requires torch._C._accelerator_setAllocatorSettings." + ) + disabled_parts = [ + ( + "expandable_segments:False" + if part.split(":", 1)[0].strip() == "expandable_segments" + else part + ) + for part in allocator_parts + ] + allocator_settings_to_apply = ",".join(disabled_parts) + allocator_settings_to_restore = allocator_conf + + if use_slot_memory: + slot_allocator_pool = torch.cuda.MemPool() + mempool = slot_allocator_pool.id + else: + slot_allocator_pool = None + mempool = graph_pool_handle() if pool is None else pool # Warmup # Hopefully prevents cudnn benchmarking and other lazy-initialization cuda work @@ -444,9 +612,206 @@ def _make_graphed_callables( f"Warmup runs {len(warmup_func)} but only {len(set(warmup_func_idx))} are unique." ) + warmup_plan_aliases = {} + if warmup_plan_alias_groups is not None: + templates = {} + unique_warmups = [] + for func_idx, func in zip(warmup_func_idx, warmup_func): + group = warmup_plan_alias_groups[func_idx] + template = templates.get(group) + if template is None: + templates[group] = (func_idx, func) + warmup_plan_aliases[func_idx] = [] + unique_warmups.append((func_idx, func)) + else: + template_idx, template_func = template + if template_func is not func: + raise RuntimeError( + f"Warmup-plan alias group {group} spans different callable objects." + ) + warmup_plan_aliases[template_idx].append(func_idx) + # Alias-group IDs also define the communicator warmup order. Dynamic-CP captures use + # variant 0 for the largest CP group, even though mutually exclusive smaller branches + # must appear first in the formal graph order for memory liveness. Warm the largest + # group first so its P2P ring is fully initialized before switching among subgroups. + ordered_warmups = sorted(unique_warmups, key=lambda item: warmup_plan_alias_groups[item[0]]) + warmup_func_idx = [func_idx for func_idx, _ in ordered_warmups] + warmup_func = [func for _, func in ordered_warmups] + # Filter the TE modules that cudagraph can access. visited_te_modules = {} need_bwd_dw_graph = {} + per_callable_fused_wgrad_params = {} + if use_slot_memory: + num_graph_inputs = len(flatten_sample_args) + per_callable_saved_tensor_plans = [None] * num_graph_inputs + per_callable_saved_tensor_boundary_aliases = [None] * num_graph_inputs + per_callable_output_tensor_plans = [None] * num_graph_inputs + per_callable_user_grad_tensor_plans = [None] * num_graph_inputs + per_callable_param_grad_tensor_targets = None + per_callable_external_storage_ptrs = [ + { + _tensor_storage_ptr(tensor) + for tensor in static_input_surface + if isinstance(tensor, torch.Tensor) + } + for static_input_surface in per_callable_static_input_surfaces + ] + per_callable_snapshot_input_storage_ptrs = [] + for func_idx, args in enumerate(sample_args): + if not args or type(args[0]) is not torch.Tensor or not args[0].is_cuda: + raise RuntimeError( + "Slot user-input snapshots require the first positional argument for " + f"graph input {func_idx} to be a plain CUDA tensor." + ) + if flatten_sample_args[func_idx][0] is not args[0]: + raise RuntimeError( + "Slot user-input snapshots require the first positional tensor to be the " + "first flattened graph input." + ) + per_callable_snapshot_input_storage_ptrs.append(_tensor_storage_ptr(args[0])) + else: + per_callable_saved_tensor_plans = None + per_callable_saved_tensor_boundary_aliases = None + per_callable_output_tensor_plans = None + per_callable_user_grad_tensor_plans = None + per_callable_param_grad_tensor_targets = None + per_callable_external_storage_ptrs = None + per_callable_snapshot_input_storage_ptrs = None + + def clone_warmup_plan(template_idx, target_idx): + """Clone one shape-identical warmup observation onto another static slot.""" + source_args = flatten_sample_args[template_idx] + target_args = flatten_sample_args[target_idx] + if len(source_args) != len(target_args): + raise RuntimeError( + f"Warmup-plan aliases {template_idx} and {target_idx} expose different " + "numbers of user tensors." + ) + + external_storage_map = {} + for source, target in zip(source_args, target_args): + if _input_staging_key(source) != _input_staging_key(target): + raise RuntimeError( + f"Warmup-plan aliases {template_idx} and {target_idx} have incompatible " + "user tensor surfaces." + ) + source_ptr = _tensor_storage_ptr(source) + target_ptr = _tensor_storage_ptr(target) + previous_target = external_storage_map.setdefault(source_ptr, target_ptr) + if previous_target != target_ptr: + raise RuntimeError( + f"Warmup-plan alias {target_idx} changes an input storage alias from " + f"{previous_target} to {target_ptr}." + ) + + source_plan = per_callable_saved_tensor_plans[template_idx] + if source_plan is None: + raise RuntimeError(f"Warmup template {template_idx} has no saved-tensor plan.") + per_callable_saved_tensor_plans[target_idx] = [ + ( + (spec[0], external_storage_map.get(spec[1], spec[1]), *spec[2:]) + if spec[0] == "external" + else spec + ) + for spec in source_plan + ] + per_callable_saved_tensor_boundary_aliases[target_idx] = list( + per_callable_saved_tensor_boundary_aliases[template_idx] + ) + for plans in ( + per_callable_output_tensor_plans, + per_callable_user_grad_tensor_plans, + ): + plans[target_idx] = list(plans[template_idx]) + + per_callable_module_params[target_idx] = per_callable_module_params[template_idx] + per_callable_static_input_surfaces[target_idx] = ( + target_args + per_callable_module_params[target_idx] + ) + visited_te_modules[target_idx] = set(visited_te_modules.get(template_idx, set())) + per_callable_fused_wgrad_params[target_idx] = set( + per_callable_fused_wgrad_params.get(template_idx, set()) + ) + need_bwd_dw_graph[target_idx] = need_bwd_dw_graph.get(template_idx, False) + + def update_warmup_plan(plans, func_idx, observed_plan, phase): + """Record a stable slot-memory plan across warmup iterations.""" + expected_plan = plans[func_idx] + if expected_plan is None: + plans[func_idx] = observed_plan + elif expected_plan != observed_plan: + raise RuntimeError( + f"{phase} saved tensors changed across CUDA graph warmup iterations " + f"for graph input {func_idx}." + ) + + def observe_saved_tensor_boundary_aliases( + func_idx, saved_tensors, saved_versions, outputs, saved_plan + ): + """Record native saves that are byte ranges of public graph boundaries.""" + boundaries = [] + for kind, tensors in ( + ( + "input", + per_callable_static_input_surfaces[func_idx][ + : per_callable_len_user_args[func_idx] + ], + ), + ("output", outputs), + ): + for tensor_idx, tensor in enumerate(tensors): + if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda: + continue + span_bytes = _saved_tensor_signature(tensor)[-1] + start = tensor.storage_offset() * tensor.element_size() + boundaries.append( + ( + tensor.untyped_storage()._cdata, + start, + start + span_bytes, + span_bytes, + kind, + tensor_idx, + _tensor_version(tensor), + ) + ) + + aliases = [] + for tensor, saved_version, spec in zip(saved_tensors, saved_versions, saved_plan): + if spec[0] != "native" or spec[7] == 0: + aliases.append(None) + continue + saved_start = tensor.storage_offset() * tensor.element_size() + saved_end = saved_start + spec[7] + storage_id = tensor.untyped_storage()._cdata + candidates = [ + ( + span_bytes, + kind, + tensor_idx, + saved_start - boundary_start, + saved_version is not None and saved_version == boundary_version, + ) + for ( + boundary_storage_id, + boundary_start, + boundary_end, + span_bytes, + kind, + tensor_idx, + boundary_version, + ) in boundaries + if boundary_storage_id == storage_id + and boundary_start <= saved_start + and saved_end <= boundary_end + ] + aliases.append( + min(candidates, key=lambda candidate: (not candidate[4], candidate[:4])) + if candidates + else None + ) + return aliases # Run warmup and do the above filtering. with torch.cuda.stream(torch.cuda.Stream()): @@ -454,6 +819,10 @@ def _make_graphed_callables( args = sample_args[func_idx] kwargs = sample_kwargs[func_idx] static_input_surface = per_callable_static_input_surfaces[func_idx] + if per_callable_external_storage_ptrs is not None and isinstance(func, torch.nn.Module): + per_callable_external_storage_ptrs[func_idx].update( + _tensor_storage_ptr(buffer) for buffer in func.buffers() + ) def hook_fn( module, inputs, outputs, func_idx=func_idx @@ -487,7 +856,79 @@ def hook_fn( for module in func.modules(): hook = module.register_forward_hook(hook_fn) hooks.append(hook) - outputs, _ = _tree_flatten(func(*args, **kwargs)) + + if use_slot_memory: + observed_saved_tensor_plan = [] + observed_saved_tensors = [] + observed_saved_versions = [] + copied_storages = {} + + def record_saved_tensor(tensor): + observed_saved_tensors.append(tensor) + observed_saved_versions.append(_tensor_version(tensor)) + storage_ptr = _tensor_storage_ptr(tensor) + signature = _saved_tensor_signature(tensor) + snapshot_user_input = ( + tensor.is_cuda + and storage_ptr == per_callable_snapshot_input_storage_ptrs[func_idx] + ) + is_external = not tensor.is_cuda or ( + storage_ptr in per_callable_external_storage_ptrs[func_idx] + and not snapshot_user_input + ) + if not is_external and type(tensor) is not torch.Tensor: + raise RuntimeError( + "CUDA graph saved-tensor arenas do not yet support tensor " + f"subclass {type(tensor).__name__}." + ) + storage_group = None + storage_offset_bytes = None + if not is_external: + storage_identity = (storage_ptr, _tensor_version(tensor)) + storage_group = copied_storages.setdefault( + storage_identity, len(copied_storages) + ) + storage_offset_bytes = tensor.storage_offset() * tensor.element_size() + observed_saved_tensor_plan.append( + ( + "external" if is_external else "native", + storage_ptr if is_external else None, + *signature, + storage_group, + storage_offset_bytes, + ) + ) + return tensor + + with torch.autograd.graph.saved_tensors_hooks(record_saved_tensor, lambda x: x): + outputs, _ = _tree_flatten(func(*args, **kwargs)) + observed_boundary_aliases = observe_saved_tensor_boundary_aliases( + func_idx, + observed_saved_tensors, + observed_saved_versions, + outputs, + observed_saved_tensor_plan, + ) + update_warmup_plan( + per_callable_saved_tensor_plans, + func_idx, + observed_saved_tensor_plan, + "Forward", + ) + update_warmup_plan( + per_callable_saved_tensor_boundary_aliases, + func_idx, + observed_boundary_aliases, + "Forward boundary alias", + ) + update_warmup_plan( + per_callable_output_tensor_plans, + func_idx, + [_io_tensor_plan(output, "output") for output in outputs], + "Output", + ) + else: + outputs, _ = _tree_flatten(func(*args, **kwargs)) for hook in hooks: hook.remove() if is_training: @@ -501,6 +942,27 @@ def hook_fn( grad_tensors=tuple(torch.empty_like(o) for o in outputs_requiring_grad), ) grad_inputs = tuple(input.grad for input in inputs) + if use_slot_memory: + observed_user_grad_tensor_plan = [] + grad_idx = 0 + for input_idx, input_tensor in enumerate(static_input_surface): + grad_input = None + if ( + isinstance(input_tensor, torch.Tensor) + and input_tensor.requires_grad + ): + grad_input = grad_inputs[grad_idx] + grad_idx += 1 + if input_idx < per_callable_len_user_args[func_idx]: + observed_user_grad_tensor_plan.append( + _io_tensor_plan(grad_input, "user_grad") + ) + update_warmup_plan( + per_callable_user_grad_tensor_plans, + func_idx, + observed_user_grad_tensor_plan, + "User-gradient output", + ) # Filter module params that get None grad from grad_inputs and remove them # from static_input_surface. This is to ensure that the backward hooks @@ -512,8 +974,27 @@ def hook_fn( for i, arg in enumerate(static_input_surface): if arg.requires_grad: required_grad_input_idx.append(i) + fused_wgrad_params = set() + if use_slot_memory: + for module in visited_te_modules.get(func_idx, set()): + if not ( + isinstance(module, TransformerEngineBaseModule) + and getattr(module, "fuse_wgrad_accumulation", False) + ): + continue + for name in getattr(module, "weight_names", ()): + param = getattr(module, name, None) + if isinstance(param, torch.nn.Parameter) and param.requires_grad: + fused_wgrad_params.add(param) + get_dummy_wgrad( + list(param.shape), + param.dtype, + zero=getattr(param, "zero_out_wgrad", False), + ) + per_callable_fused_wgrad_params[func_idx] = fused_wgrad_params module_params_with_grad = [] for grad_inputs_idx, inputs_idx in enumerate(required_grad_input_idx): + input_tensor = static_input_surface[inputs_idx] if ( grad_inputs[grad_inputs_idx] is None and grad_inputs_idx < num_required_grad_sample_args @@ -523,11 +1004,14 @@ def hook_fn( "The input tensor requires grad, but the grad is None after" " backward pass." ) - elif ( + elif grad_inputs_idx >= num_required_grad_sample_args and ( grad_inputs[grad_inputs_idx] is not None - and grad_inputs_idx >= num_required_grad_sample_args + or input_tensor in fused_wgrad_params ): - module_params_with_grad.append(static_input_surface[inputs_idx]) + # Fused wgrad writes directly into main_grad. Keep its parameter as + # an autograd input even when no ordinary param.grad was materialized, + # so replay can still trigger AccumulateGrad/DDP hooks. + module_params_with_grad.append(input_tensor) if len(module_params_with_grad) != len(per_callable_module_params[func_idx]): if warmup_iter != 0: raise RuntimeError( @@ -551,14 +1035,1126 @@ def hook_fn( else: grad_inputs = None del outputs, grad_inputs + if is_training: + del outputs_requiring_grad + if use_slot_memory: + grad_input = None if post_warmup_hook is not None: post_warmup_hook() + if warmup_plan_alias_groups is not None: + # Dynamic-CP warmup callables can replace the CP process group while TE still + # has asynchronous CP/TP work queued on auxiliary streams. Drain every observed + # callable before changing groups; otherwise one TP peer can enter the next + # variant while the other is still completing the previous CP ring. + torch.cuda.synchronize() + for target_idx in warmup_plan_aliases.get(func_idx, ()): + clone_warmup_plan(func_idx, target_idx) torch.cuda.synchronize() + if use_slot_memory: + per_callable_param_grad_tensor_targets = [ + [None] * len(static_input_surface) + for static_input_surface in per_callable_static_input_surfaces + ] + + if allocator_settings_to_apply is not None: + torch._C._accelerator_setAllocatorSettings(allocator_settings_to_apply) + + if use_slot_memory: + if isinstance(sample_args, tuple): + sample_args = list(sample_args) + + staging_group_by_key = {} + staging_groups = [] + for func_idx, args in enumerate(sample_args): + old_input = args[0] + saved_arena_id, _ = saved_tensor_memory_alias_groups[func_idx] + staging_key = (saved_arena_id, _input_staging_key(old_input)) + group_idx = staging_group_by_key.get(staging_key) + if group_idx is None: + group_idx = len(staging_groups) + staging_group_by_key[staging_key] = group_idx + staging_groups.append({"members": [], "candidates": {}}) + group = staging_groups[group_idx] + group["members"].append(func_idx) + group["candidates"].setdefault(old_input.untyped_storage()._cdata, old_input) + + # MCore's sample-input plan and the union liveness coloring are each safe in + # isolation, but reusing an arbitrary representative can transitively merge two + # conflicting colors. Match colors onto distinct existing storages first, then + # allocate only when the original CP-variant plans do not provide enough choices. + storage_owner = {} + staging_targets = {} + + def match_staging_group(group_idx, seen_storages): + for storage_id, tensor in staging_groups[group_idx]["candidates"].items(): + if storage_id in seen_storages: + continue + seen_storages.add(storage_id) + previous_group = storage_owner.get(storage_id) + if previous_group is None or match_staging_group(previous_group, seen_storages): + storage_owner[storage_id] = group_idx + staging_targets[group_idx] = tensor + return True + return False + + for group_idx in sorted( + range(len(staging_groups)), + key=lambda index: len(staging_groups[index]["candidates"]), + ): + match_staging_group(group_idx, set()) + + for group_idx, group in enumerate(staging_groups): + input_target = staging_targets.get(group_idx) + if input_target is None: + source = next(iter(group["candidates"].values())) + signature = _saved_tensor_signature(source) + storage_numel = signature[-1] // source.element_size() + backing = torch.empty( + (source.storage_offset() + storage_numel,), + dtype=source.dtype, + device=source.device, + ) + input_target = torch.empty((0,), dtype=source.dtype, device=source.device).set_( + backing.untyped_storage(), + source.storage_offset(), + source.shape, + source.stride(), + ) + with torch.no_grad(): + input_target.copy_(source) + input_target.requires_grad_(source.requires_grad) + staging_targets[group_idx] = input_target + + for func_idx in group["members"]: + args = sample_args[func_idx] + old_input = args[0] + if input_target is old_input: + continue + + args = list(args) + args[0] = input_target + sample_args[func_idx] = tuple(args) + flattened_args = list(flatten_sample_args[func_idx]) + flattened_args[0] = input_target + flatten_sample_args[func_idx] = tuple(flattened_args) + static_input_surface = list(per_callable_static_input_surfaces[func_idx]) + static_input_surface[0] = input_target + per_callable_static_input_surfaces[func_idx] = tuple(static_input_surface) + + def prepare_native_io_targets(per_callable_plans, kind): + """Validate same-slot CP branches and create their lazy alias targets.""" + if per_callable_plans is None: + return None + + plans_by_family = {} + for func_idx, plan in enumerate(per_callable_plans): + arena_id, branch_id = slot_io_memory_alias_groups[func_idx] + family = (arena_id, *slot_io_liveness_groups[func_idx]) + branch_plans = plans_by_family.setdefault(family, {}) + if branch_id in branch_plans: + raise RuntimeError( + f"CUDA graph {kind} family {family} has duplicate branch {branch_id}." + ) + branch_plans[branch_id] = plan + + for family, branch_plans in plans_by_family.items(): + plans = list(branch_plans.values()) + if len({len(plan) for plan in plans}) != 1: + raise RuntimeError( + f"CUDA graph {kind} family {family} exposes different tensor counts." + ) + for tensor_idx, specs in enumerate(zip(*plans)): + if all(spec is None for spec in specs): + continue + if any(spec is None for spec in specs): + raise RuntimeError( + f"CUDA graph {kind} family {family} has an incompatible tensor " + f"at position {tensor_idx}." + ) + layout_keys = {(spec[0], spec[4], spec[5], spec[6]) for spec in specs} + if layout_keys != {(kind, specs[0][4], specs[0][5], specs[0][6])}: + raise RuntimeError( + f"CUDA graph {kind} family {family} has incompatible dtype, device, " + f"or autograd state at position {tensor_idx}." + ) + + return [[None] * len(plan) for plan in per_callable_plans] + + per_callable_output_tensor_targets = prepare_native_io_targets( + per_callable_output_tensor_plans, "output" + ) + per_callable_user_grad_tensor_targets = prepare_native_io_targets( + per_callable_user_grad_tensor_plans, "user_grad" + ) + + native_io_family_sizes = {} + if use_slot_memory: + for func_idx in range(len(flatten_sample_args)): + arena_id, _ = slot_io_memory_alias_groups[func_idx] + family = (arena_id, *slot_io_liveness_groups[func_idx]) + native_io_family_sizes[family] = native_io_family_sizes.get(family, 0) + 1 + native_io_anchors = {"output": {}, "user_grad": {}, "param_grad": {}} + native_io_capture_counts = {"output": {}, "user_grad": {}, "param_grad": {}} + release_native_io_targets = use_slot_memory + + def native_io_alias_target(func_idx, tensor_idx, tensor, spec, kind): + """Alias one CP branch onto the first branch's graph-pool I/O storage.""" + arena_id, _ = slot_io_memory_alias_groups[func_idx] + family = (arena_id, *slot_io_liveness_groups[func_idx]) + key = (*family, tensor_idx) + anchors = native_io_anchors[kind] + counts = native_io_capture_counts[kind] + anchor = anchors.get(key) + if anchor is None: + target = tensor + anchors[key] = tensor + else: + available_bytes = anchor.untyped_storage().nbytes() - ( + anchor.storage_offset() * anchor.element_size() + ) + if spec[7] > available_bytes: + raise RuntimeError( + f"CUDA graph native {kind} alias {key} needs {spec[7]} bytes, " + f"but its first CP branch exposes only {available_bytes} bytes." + ) + target = _arena_view(anchor, 0, spec) + + captured = counts.get(key, 0) + 1 + expected = native_io_family_sizes[family] + if captured > expected: + raise RuntimeError( + f"CUDA graph native {kind} alias {key} captured {captured} of " + f"{expected} CP branches." + ) + if captured == expected: + anchors.pop(key) + counts.pop(key, None) + else: + counts[key] = captured + return target + + def clear_native_io_target_rows(func_indices, clear_outputs=False, clear_grads=False): + """Drop capture-only I/O aliases after the corresponding TE value dies.""" + if not release_native_io_targets: + return + for func_idx in func_indices: + if clear_outputs: + per_callable_output_tensor_targets[func_idx] = [None] * len( + per_callable_output_tensor_targets[func_idx] + ) + if clear_grads: + per_callable_user_grad_tensor_targets[func_idx] = [None] * len( + per_callable_user_grad_tensor_targets[func_idx] + ) + + def copy_outputs_to_slot_arena(func_idx, flatten_outputs): + """Copy public forward outputs to the fixed surface for their physical slot.""" + if per_callable_output_tensor_targets is None: + return flatten_outputs + plan = per_callable_output_tensor_plans[func_idx] + targets = per_callable_output_tensor_targets[func_idx] + if len(flatten_outputs) != len(plan): + raise RuntimeError( + f"CUDA graph input {func_idx} changed its output count during capture." + ) + copied_outputs = [] + for tensor_idx, (output, spec, target) in enumerate(zip(flatten_outputs, plan, targets)): + if spec != _io_tensor_plan(output, "output"): + raise RuntimeError( + f"CUDA graph input {func_idx} changed its output tensor surface during capture." + ) + if target is None: + if spec is None: + copied_outputs.append(output) + continue + target = native_io_alias_target(func_idx, tensor_idx, output, spec, "output") + targets[tensor_idx] = target + if target is not output: + target.copy_(output) + copied_outputs.append(target) + return copied_outputs + + def copy_user_grads_to_slot_arena(func_idx, static_input_surface, grad_inputs): + """Copy returned gradients to the fixed surface for their physical slot.""" + if per_callable_user_grad_tensor_targets is None: + return grad_inputs + plan = per_callable_user_grad_tensor_plans[func_idx] + targets = per_callable_user_grad_tensor_targets[func_idx] + copied_grad_inputs = [] + grad_idx = 0 + for input_idx, input_tensor in enumerate(static_input_surface): + if not (isinstance(input_tensor, torch.Tensor) and input_tensor.requires_grad): + continue + grad_input = grad_inputs[grad_idx] + grad_idx += 1 + if input_idx < per_callable_len_user_args[func_idx]: + spec = plan[input_idx] + target = targets[input_idx] + if spec != _io_tensor_plan(grad_input, "user_grad"): + raise RuntimeError( + f"CUDA graph input {func_idx} changed its user-gradient tensor " + "surface during capture." + ) + if target is None and spec is not None: + target = native_io_alias_target( + func_idx, input_idx, grad_input, spec, "user_grad" + ) + targets[input_idx] = target + if target is not None: + if target is not grad_input: + with torch.no_grad(): + target.copy_(grad_input) + grad_input = target + elif per_callable_param_grad_tensor_targets is not None and grad_input is not None: + spec = _io_tensor_plan(grad_input, "param_grad") + if spec is None: + raise RuntimeError( + f"CUDA graph input {func_idx} produced an unsupported parameter-gradient " + f"tensor at input position {input_idx}." + ) + target = per_callable_param_grad_tensor_targets[func_idx][input_idx] + if target is None: + target = native_io_alias_target( + func_idx, input_idx, grad_input, spec, "param_grad" + ) + per_callable_param_grad_tensor_targets[func_idx][input_idx] = target + if target is not grad_input: + with torch.no_grad(): + target.copy_(grad_input) + grad_input = target + copied_grad_inputs.append(grad_input) + return tuple(copied_grad_inputs) + + per_callable_native_saved_storages = [dict() for _ in flatten_sample_args] + per_callable_native_saved_intervals = [[] for _ in flatten_sample_args] + per_callable_native_saved_capture_targets = [None] * len(flatten_sample_args) + + def plan_native_saved_alias_targets( + plan, + canonical_targets, + measure_spill=False, + protected_storage_ranges=None, + preassigned_targets=None, + ): + """Pack one CP branch's native saved tensors into canonical live storages.""" + protected_storage_ranges = protected_storage_ranges or {} + if preassigned_targets is None: + target_views = [None] * len(plan) + else: + if len(preassigned_targets) != len(plan): + raise RuntimeError("Native saved preassignment does not match its plan.") + target_views = list(preassigned_targets) + storage_banks = [] + seen_storages = set() + for target in canonical_targets: + if target is None: + continue + storage = target.untyped_storage() + if storage._cdata in seen_storages: + continue + seen_storages.add(storage._cdata) + cursor = 0 + for protected_start, protected_end in protected_storage_ranges.get(storage._cdata, ()): + if cursor < protected_start: + storage_banks.append( + { + "storage": storage, + "base_offset": cursor, + "capacity": protected_start - cursor, + "cursor": 0, + } + ) + cursor = max(cursor, protected_end) + if cursor < storage.nbytes(): + storage_banks.append( + { + "storage": storage, + "base_offset": cursor, + "capacity": storage.nbytes() - cursor, + "cursor": 0, + } + ) + records_by_storage_group = {} + for saved_idx, spec in enumerate(plan): + if spec[0] != "native": + continue + if target_views[saved_idx] is not None: + continue + if spec[7] == 0: + target = torch.empty_strided(spec[2], spec[3], dtype=spec[4], device=spec[5]) + target.requires_grad_(spec[6]) + target_views[saved_idx] = target + continue + storage_group = spec[8] + storage_offset_bytes = spec[9] + if storage_group is None or storage_offset_bytes is None: + raise RuntimeError(f"Native saved tensor {saved_idx} has no backing-storage plan.") + records_by_storage_group.setdefault(storage_group, []).append( + (storage_offset_bytes, storage_offset_bytes + spec[7], saved_idx) + ) + + components = [] + for records in records_by_storage_group.values(): + records.sort() + group_components = [] + for start, end, saved_idx in records: + if not group_components or start >= group_components[-1][1]: + group_components.append([start, end, [(start, saved_idx)]]) + else: + group_components[-1][1] = max(group_components[-1][1], end) + group_components[-1][2].append((start, saved_idx)) + components.extend(group_components) + + packed_components = [] + for component_start, component_end, component_records in components: + component_alignment = max( + plan[saved_idx][4].itemsize for _, saved_idx in component_records + ) + component_origin = component_start // component_alignment * component_alignment + packed_components.append( + ( + component_end - component_origin, + component_origin, + component_records, + ) + ) + + banks = list(storage_banks) + spill_bank = None + if measure_spill: + spill_bank = { + "storage": None, + "base_offset": 0, + "capacity": sum(_align_up(item[0]) for item in packed_components), + "cursor": 0, + } + banks.append(spill_bank) + for component_size, component_origin, component_records in sorted( + packed_components, key=lambda item: item[0], reverse=True + ): + candidates = [] + for bank_idx, bank in enumerate(banks): + if bank["storage"] is None: + continue + offset = _align_up(bank["cursor"]) + if offset + component_size <= bank["capacity"]: + candidates.append( + (bank["capacity"] - offset - component_size, bank_idx, offset) + ) + if not candidates and spill_bank is not None: + spill_bank_idx = len(banks) - 1 + spill_offset = _align_up(spill_bank["cursor"]) + if spill_offset + component_size <= spill_bank["capacity"]: + candidates.append((0, spill_bank_idx, spill_offset)) + if not candidates: + raise RuntimeError( + "CUDA graph CP branch native saved tensors do not fit in canonical " + "live storages: " + f"component_bytes={component_size}, " + f"component_sizes={sorted((item[0] for item in packed_components), reverse=True)}, " + f"storage_capacities={sorted((bank['capacity'] for bank in banks), reverse=True)}." + ) + _, bank_idx, component_target_offset = min(candidates) + bank = banks[bank_idx] + bank["cursor"] = component_target_offset + component_size + for source_offset, saved_idx in component_records: + if bank["storage"] is None: + target_views[saved_idx] = True + continue + spec = plan[saved_idx] + target_offset = ( + bank["base_offset"] + component_target_offset + source_offset - component_origin + ) + itemsize = spec[4].itemsize + if target_offset % itemsize: + raise RuntimeError( + f"Native saved tensor {saved_idx} has an unaligned canonical offset." + ) + target = torch.empty((0,), dtype=spec[4], device=spec[5]).set_( + bank["storage"], + target_offset // itemsize, + spec[2], + spec[3], + ) + target.requires_grad_(spec[6]) + target_views[saved_idx] = target + + if measure_spill: + return _align_up(spill_bank["cursor"]) + + missing = [ + saved_idx + for saved_idx, spec in enumerate(plan) + if spec[0] == "native" and target_views[saved_idx] is None + ] + if missing: + raise RuntimeError(f"Native saved tensors have no canonical targets: {missing}.") + return tuple(target_views) + + def semantic_boundary_alias_targets(canonical_func_idx, sibling_func_idx): + """Map boundary-backed sibling saves onto the boundary address used at replay.""" + plan = per_callable_saved_tensor_plans[sibling_func_idx] + aliases = per_callable_saved_tensor_boundary_aliases[sibling_func_idx] + targets = [None] * len(plan) + records_by_storage_group = {} + for saved_idx, spec in enumerate(plan): + if spec[0] != "native" or spec[7] == 0: + continue + records_by_storage_group.setdefault(spec[8], []).append( + (spec[9], spec[9] + spec[7], saved_idx) + ) + + components = [] + for records in records_by_storage_group.values(): + records.sort() + group_components = [] + for start, end, saved_idx in records: + if not group_components or start >= group_components[-1][1]: + group_components.append([start, end, [saved_idx]]) + else: + group_components[-1][1] = max(group_components[-1][1], end) + group_components[-1][2].append(saved_idx) + components.extend(group_components) + + for component_start, component_end, component_saved_indices in components: + candidate_aliases = [ + (saved_idx, aliases[saved_idx]) + for saved_idx in component_saved_indices + if aliases[saved_idx] is not None + ] + if not candidate_aliases: + continue + + # An alias only proves that one saved view is a graph-boundary view. Reusing + # the boundary for its whole overlapping storage component is safe only when + # every byte in that component is part of the same logical boundary tensor. + component_aliases = [] + for saved_idx, alias in candidate_aliases: + boundary_span_bytes, _, _, relative_offset, version_matches = alias + if not version_matches: + continue + source_boundary_start = plan[saved_idx][9] - relative_offset + source_boundary_end = source_boundary_start + boundary_span_bytes + if ( + source_boundary_start <= component_start + and component_end <= source_boundary_end + ): + component_aliases.append((saved_idx, alias)) + if not component_aliases: + continue + + anchor_storage = None + anchor_shift = None + for saved_idx, alias in component_aliases: + _, kind, boundary_idx, relative_offset, _ = alias + if kind == "input": + boundary = per_callable_static_input_surfaces[sibling_func_idx][boundary_idx] + else: + boundary = per_callable_static_outputs[canonical_func_idx][boundary_idx] + if not isinstance(boundary, torch.Tensor) or not boundary.is_cuda: + raise RuntimeError( + f"CUDA graph {kind} boundary {boundary_idx} is not a CUDA tensor." + ) + + spec = plan[saved_idx] + storage = boundary.untyped_storage() + boundary_start = boundary.storage_offset() * boundary.element_size() + shift = boundary_start + relative_offset - spec[9] + if anchor_storage is None: + anchor_storage = storage + anchor_shift = shift + elif anchor_storage._cdata != storage._cdata or anchor_shift != shift: + raise RuntimeError( + "CUDA graph overlapping saved tensors have inconsistent boundary " + f"aliases: func={sibling_func_idx}, saved={component_saved_indices}." + ) + + for saved_idx in component_saved_indices: + spec = plan[saved_idx] + target_offset = spec[9] + anchor_shift + if target_offset < 0 or target_offset + spec[7] > anchor_storage.nbytes(): + raise RuntimeError( + "CUDA graph boundary-backed saved component does not fit its replay " + f"storage: func={sibling_func_idx}, saved={saved_idx}, " + f"offset={target_offset}, bytes={spec[7]}, " + f"storage_bytes={anchor_storage.nbytes()}." + ) + itemsize = spec[4].itemsize + if target_offset % itemsize: + raise RuntimeError( + f"CUDA graph boundary-backed saved tensor {saved_idx} is unaligned." + ) + target = torch.empty((0,), dtype=spec[4], device=spec[5]).set_( + anchor_storage, + target_offset // itemsize, + spec[2], + spec[3], + ) + target.requires_grad_(spec[6]) + targets[saved_idx] = target + return tuple(targets) + + def materialize_native_saved_spill_targets(canonical_func_indices): + """Complete canonical saved-tensor storage for every same-slot CP branch.""" + completed_targets = [] + protected_ranges = {} + preassigned_targets = {} + + for func_idx in canonical_func_indices: + targets = per_callable_native_saved_capture_targets[func_idx] + if targets is None: + raise RuntimeError("CUDA graph canonical CP branch did not retain saved targets.") + + storage_ranges = {} + boundary_tensors = ( + *per_callable_static_input_surfaces[func_idx][ + : per_callable_len_user_args[func_idx] + ], + *per_callable_static_outputs[func_idx], + ) + for tensor in boundary_tensors: + if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda: + continue + storage = tensor.untyped_storage() + storage_ranges.setdefault(storage._cdata, []).append((0, storage.nbytes())) + for storage_id, ranges in storage_ranges.items(): + merged = [] + for start, end in sorted(ranges): + if not merged or start > merged[-1][1]: + merged.append([start, end]) + else: + merged[-1][1] = max(merged[-1][1], end) + storage_ranges[storage_id] = tuple(map(tuple, merged)) + + family = saved_tensor_memory_families[func_idx] + sibling_indices = [ + sibling_idx + for sibling_idx, sibling_family in enumerate(saved_tensor_memory_families) + if sibling_family == family + ] + for sibling_idx in sibling_indices: + protected_ranges[sibling_idx] = storage_ranges + preassigned_targets[sibling_idx] = ( + semantic_boundary_alias_targets(func_idx, sibling_idx) + if sibling_idx != func_idx + else None + ) + + spill_bytes = max( + ( + plan_native_saved_alias_targets( + per_callable_saved_tensor_plans[sibling_idx], + targets, + measure_spill=True, + protected_storage_ranges=storage_ranges, + preassigned_targets=preassigned_targets[sibling_idx], + ) + for sibling_idx in sibling_indices + if sibling_idx != func_idx + ), + default=0, + ) + if spill_bytes: + with torch.cuda.use_mem_pool(slot_allocator_pool): + spill = torch.empty( + (spill_bytes,), + dtype=torch.uint8, + device=torch.cuda.current_device(), + ) + storage = spill.untyped_storage() + per_callable_native_saved_storages[func_idx].setdefault( + storage._cdata, (storage, storage.data_ptr()) + ) + targets = (*targets, spill) + completed_targets.append(tuple(targets)) + return tuple(completed_targets), protected_ranges, preassigned_targets + + @contextlib.contextmanager + def capture_saved_tensors(func_idx, alias_targets=None): + """Capture forward tensors that cross the graph's F/B boundary.""" + if per_callable_saved_tensor_plans is None: + yield + return + + plan = per_callable_saved_tensor_plans[func_idx] + saved_idx = 0 + captured_targets = [None] * len(plan) + per_callable_native_saved_intervals[func_idx].clear() + if alias_targets is not None and len(alias_targets) != len(plan): + raise RuntimeError( + f"CUDA graph input {func_idx} changed its canonical saved-target count." + ) + + def pack_saved_tensor(tensor): + nonlocal saved_idx + if saved_idx >= len(plan): + raise RuntimeError( + f"CUDA graph input {func_idx} saved more forward tensors during capture " + "than warmup." + ) + current_saved_idx = saved_idx + spec = plan[current_saved_idx] + saved_idx += 1 + if spec[2:8] != _saved_tensor_signature(tensor): + raise RuntimeError( + f"CUDA graph input {func_idx} changed forward saved-tensor layout " + "during capture." + ) + if spec[0] == "external": + if spec[1] != _tensor_storage_ptr(tensor): + raise RuntimeError( + f"CUDA graph input {func_idx} changed an external saved tensor." + ) + return tensor + if spec[0] != "native": + raise RuntimeError( + f"CUDA graph input {func_idx} has unsupported saved-tensor mode {spec[0]}." + ) + + if alias_targets is None: + target = torch.empty((0,), dtype=tensor.dtype, device=tensor.device).set_( + tensor.untyped_storage(), + tensor.storage_offset(), + tensor.shape, + tensor.stride(), + ) + target.requires_grad_(tensor.requires_grad) + else: + target = alias_targets[current_saved_idx] + if target is None: + raise RuntimeError( + f"CUDA graph input {func_idx} has no canonical target for native " + f"saved tensor {current_saved_idx}." + ) + same_view = ( + target.data_ptr() == tensor.data_ptr() + and target.shape == tensor.shape + and target.stride() == tensor.stride() + and target.dtype == tensor.dtype + ) + if not same_view: + with torch.no_grad(): + target.copy_(tensor) + tensor = target + + captured_targets[current_saved_idx] = target + storage = tensor.untyped_storage() + per_callable_native_saved_storages[func_idx].setdefault( + storage._cdata, (storage, storage.data_ptr()) + ) + if spec[7]: + start = _tensor_storage_ptr(tensor) + ( + tensor.storage_offset() * tensor.element_size() + ) + per_callable_native_saved_intervals[func_idx].append( + (current_saved_idx, start, start + spec[7]) + ) + return tensor + + with torch.autograd.graph.saved_tensors_hooks(pack_saved_tensor, lambda x: x): + yield + if saved_idx != len(plan): + raise RuntimeError( + f"CUDA graph input {func_idx} saved {saved_idx} forward tensors during " + f"capture, but saved {len(plan)} during warmup." + ) + if alias_targets is None: + per_callable_native_saved_capture_targets[func_idx] = tuple(captured_targets) + + def validate_captured_module_grads(func_idx, static_grad_inputs): + """Require capture to preserve every parameter gradient observed during warmup.""" + if per_callable_saved_tensor_plans is None: + return + module_params = per_callable_module_params[func_idx] + module_grad_inputs = static_grad_inputs[per_callable_len_user_args[func_idx] :] + if len(module_grad_inputs) != len(module_params): + raise RuntimeError( + f"CUDA graph input {func_idx} captured {len(module_grad_inputs)} parameter " + f"gradient slots for {len(module_params)} parameters." + ) + missing_params = [ + param for param, grad in zip(module_params, module_grad_inputs) if grad is None + ] + if not missing_params: + return + + func = graph_callables[func_idx] + param_names = {} + if isinstance(func, torch.nn.Module): + param_names = {id(param): name for name, param in func.named_parameters()} + missing_names = [ + param_names.get(id(param), f"") + for param in missing_params + ] + raise RuntimeError( + f"CUDA graph input {func_idx} lost parameter gradients during capture: {missing_names}." + ) + # All captures here share a mempool. To avoid replays corrupting each other's memory, # the safest approach is to capture all passes in the same order they'll run: # fwd 1, fwd 2, ... fwd N, then bwd N, bwd N-1, ... bwd 1. + branch_capture_groups = None + branch_checkpoint_state = None + branch_checkpoint_live_blocks = None + branch_checkpoint_pool_layout = None + branch_checkpoint_storage_owners = None + branch_pre_checkpoint_state = None + branch_pre_checkpoint_live_blocks = None + branch_pre_checkpoint_storage_owners = None + current_storage_owners = None + branch_canonical_native_saved_intervals = None + branch_canonical_native_saved_targets = None + branch_canonical_native_saved_excluded_storages = None + branch_canonical_native_saved_preassigned_targets = None + native_saved_alias_targets = None + if use_slot_memory: + branch_capture_groups = [None] * len(_order) + group_records = [] + group_fwd_idx = [0] * num_model_chunks + group_bwd_idx = [0] * num_model_chunks + for order_idx, c_id in enumerate(_order): + if ceil(c_id) != c_id: + group_records.append(None) + continue + m_chunk = abs(int(c_id)) - 1 + logical_idx = group_fwd_idx[m_chunk] if c_id > 0 else group_bwd_idx[m_chunk] + first_func_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( + logical_idx * _num_layers_per_chunk[m_chunk] + ) + slot = _graph_memory_slots[first_func_idx] + event_key = ( + c_id > 0, + slot[4], + logical_idx, + slot[2], + _num_layers_per_chunk[m_chunk], + ) + group_records.append((event_key, slot[3])) + if c_id > 0: + group_fwd_idx[m_chunk] += 1 + else: + group_bwd_idx[m_chunk] += 1 + + group_start = 0 + while group_start < len(group_records): + record = group_records[group_start] + group_stop = group_start + 1 + while ( + record is not None + and group_stop < len(group_records) + and group_records[group_stop] is not None + and group_records[group_stop][0] == record[0] + ): + group_stop += 1 + group_size = group_stop - group_start + if group_size > 1: + branch_ids = [group_records[idx][1] for idx in range(group_start, group_stop)] + if len(set(branch_ids)) != group_size: + raise RuntimeError( + "CUDA graph slot checkpoint group contains duplicate branch IDs: " + f"{branch_ids}." + ) + for position, order_idx in enumerate(range(group_start, group_stop)): + branch_capture_groups[order_idx] = (position, group_size) + group_start = group_stop + + def slot_pool_active_blocks(): + """Return active allocation addresses and sizes in the graph-private pool.""" + snapshot = slot_allocator_pool.snapshot(include_traces=False) + segments = snapshot["segments"] if isinstance(snapshot, dict) else snapshot + return { + block["address"]: block["size"] + for segment in segments + for block in segment["blocks"] + if block["state"] == "active_allocated" + } + + def slot_pool_layout(): + """Return the allocator block topology for checkpoint diagnostics.""" + snapshot = slot_allocator_pool.snapshot(include_traces=False) + segments = snapshot["segments"] if isinstance(snapshot, dict) else snapshot + return { + segment["address"]: { + "total_size": segment["total_size"], + "blocks": [ + (block["address"], block["size"], block["state"]) for block in segment["blocks"] + ], + } + for segment in segments + } + + def native_saved_pool_intervals(func_indices, layout, full_storage=False): + """Return native saved-tensor intervals that belong to the slot pool.""" + segment_ranges = tuple( + (address, address + segment["total_size"]) for address, segment in layout.items() + ) + output = [] + for func_idx in func_indices: + intervals = [] + if full_storage: + candidates = ( + (storage_idx, storage_ptr, storage_ptr + storage.nbytes()) + for storage_idx, (storage, storage_ptr) in enumerate( + per_callable_native_saved_storages[func_idx].values() + ) + ) + else: + candidates = iter(per_callable_native_saved_intervals[func_idx]) + for saved_idx, start, end in candidates: + containing = [ + (segment_start, segment_end) + for segment_start, segment_end in segment_ranges + if segment_start <= start and end <= segment_end + ] + if containing: + intervals.append((saved_idx, start, end)) + continue + if any( + start < segment_end and segment_start < end + for segment_start, segment_end in segment_ranges + ): + raise RuntimeError( + "CUDA graph native saved tensor crosses a slot-pool segment boundary: " + f"func={func_idx}, saved={saved_idx}, interval=({start}, {end})." + ) + output.append(tuple(intervals)) + return tuple(output) + + def assert_native_saved_interval_coverage(canonical, current, func_indices, phase): + """Require alternate saved tensors to stay inside canonical live allocations.""" + if len(canonical) != len(current): + raise RuntimeError("CUDA graph CP branch changed its captured layer count.") + for position, (canonical_intervals, current_intervals) in enumerate( + zip(canonical, current) + ): + merged = [] + for _, start, end in sorted(canonical_intervals, key=lambda item: item[1:]): + if merged and start <= merged[-1][1]: + merged[-1] = (merged[-1][0], max(merged[-1][1], end)) + else: + merged.append((start, end)) + for saved_idx, start, end in current_intervals: + if any(left <= start and end <= right for left, right in merged): + continue + raise RuntimeError( + "CUDA graph CP branch placed a native saved tensor outside the " + "canonical slot/layer live range: " + f"phase={phase}, func={func_indices[position]}, saved={saved_idx}, " + f"interval=({start}, {end}), canonical={merged}." + ) + + def drain_slot_pool_pending_frees(): + """Poll completed cross-stream frees before restoring allocator state.""" + layout = slot_pool_layout() + if not any( + state == "active_pending_free" + for segment in layout.values() + for _, _, state in segment["blocks"] + ): + return + + # A non-capturing allocator request runs process_events(). The request stays in + # the default pool, so polling cannot split or merge the slot-private pool. + event_poll = torch.empty((1,), dtype=torch.uint8, device=torch.cuda.current_device()) + del event_poll + remaining = [ + (address, size) + for segment in slot_pool_layout().values() + for address, size, state in segment["blocks"] + if state == "active_pending_free" + ] + if remaining: + raise RuntimeError( + f"CUDA graph slot checkpoint could not drain pending allocator frees: {remaining}." + ) + + def assert_slot_pool_liveness(expected_blocks, phase): + expected_ptrs = set(expected_blocks) + if torch._C._cuda_checkPoolLiveAllocations( + torch.cuda.current_device(), mempool, expected_ptrs + ): + return + actual_blocks = slot_pool_active_blocks() + added = set(actual_blocks).difference(expected_ptrs) + removed = expected_ptrs.difference(actual_blocks) + raise RuntimeError( + f"CUDA graph slot checkpoint changed live allocations during {phase}: " + f"added={len(added)} ({sum(actual_blocks[ptr] for ptr in added)} bytes), " + f"added_sizes={sorted(actual_blocks[ptr] for ptr in added)}, " + f"removed={len(removed)} ({sum(expected_blocks[ptr] for ptr in removed)} bytes), " + f"removed_sizes={sorted(expected_blocks[ptr] for ptr in removed)}." + ) + + def record_replaced_io_storages(sources, targets, stale_storages): + """Keep source storages whose graph-boundary tensors now use canonical storage.""" + for source, target in zip(sources, targets): + if not ( + isinstance(source, torch.Tensor) + and isinstance(target, torch.Tensor) + and source.is_cuda + and target.is_cuda + ): + continue + source_storage = source.untyped_storage() + target_storage = target.untyped_storage() + if ( + source_storage.data_ptr() == target_storage.data_ptr() + and torch._C._has_Standard_Deleter(target_storage._cdata) + ): + continue + stale_storages[source_storage._cdata] = ( + source_storage, + source_storage.data_ptr(), + ) + + def checkpoint_live_storage_owners(expected_blocks, phase, extra_values=()): + """Resolve every checkpoint-live allocation to its owning StorageImpl.""" + owners = {} + visited = set() + + def visit(value): + if value is None or id(value) in visited: + return + if isinstance(value, torch.UntypedStorage): + if value.device.type != "cuda": + return + storage_ptr = value.data_ptr() + for block_ptr, block_size in expected_blocks.items(): + if block_ptr <= storage_ptr < block_ptr + block_size: + if torch._C._has_Standard_Deleter(value._cdata): + previous = owners.setdefault(block_ptr, value) + if previous._cdata != value._cdata: + raise RuntimeError( + "CUDA graph slot checkpoint found multiple owning storages " + f"for allocation {block_ptr}." + ) + break + return + visited.add(id(value)) + if isinstance(value, torch.Tensor): + if not value.is_cuda: + return + storage = value.untyped_storage() + storage_ptr = storage.data_ptr() + for block_ptr, block_size in expected_blocks.items(): + if block_ptr <= storage_ptr < block_ptr + block_size: + if torch._C._has_Standard_Deleter(storage._cdata): + previous = owners.setdefault(block_ptr, storage) + if previous._cdata != storage._cdata: + raise RuntimeError( + "CUDA graph slot checkpoint found multiple owning storages " + f"for allocation {block_ptr}." + ) + break + if value.is_leaf: + visit(value.grad) + for child in vars(value).values(): + visit(child) + return + if isinstance(value, dict): + for child in value.values(): + visit(child) + return + if isinstance(value, (list, tuple, set)): + for child in value: + visit(child) + + for value in ( + sample_args, + flatten_sample_args, + per_callable_static_input_surfaces, + per_callable_static_outputs, + per_callable_static_grad_outputs, + per_callable_static_grad_inputs, + per_callable_native_saved_storages, + per_callable_output_tensor_targets, + per_callable_user_grad_tensor_targets, + per_callable_param_grad_tensor_targets, + static_grad_outputs_dict, + native_io_anchors, + *extra_values, + ): + visit(value) + for func in graph_callables: + if isinstance(func, torch.nn.Module): + for module in func.modules(): + visit(vars(module)) + + missing = set(expected_blocks).difference(owners) + if missing: + raise RuntimeError( + f"CUDA graph slot checkpoint could not resolve owning storages during {phase} " + "for " + f"{len(missing)} live allocations with sizes " + f"{sorted(expected_blocks[ptr] for ptr in missing)}." + ) + return owners + + def release_checkpoint_live_storages(expected_blocks, owners, phase): + """Make checkpoint-live blocks free before allocator topology restoration.""" + if set(owners) != set(expected_blocks): + raise RuntimeError("CUDA graph slot checkpoint live-storage owner set changed.") + for block_ptr, storage in owners.items(): + storage_ptr = storage.data_ptr() + if not (block_ptr <= storage_ptr < block_ptr + expected_blocks[block_ptr]): + raise RuntimeError( + "CUDA graph slot checkpoint live storage changed its allocation." + ) + if not torch._C._has_Standard_Deleter(storage._cdata): + raise RuntimeError( + "CUDA graph slot checkpoint live storage lost its allocator deleter." + ) + torch._C._free_And_Remove_DeleterFn(storage._cdata) + tex._graph_checkpoint_detach_storage(storage._cdata) + if storage.data_ptr() != storage_ptr or torch._C._has_Standard_Deleter(storage._cdata): + raise RuntimeError( + "CUDA graph slot checkpoint could not detach live-storage ownership." + ) + + torch.cuda.synchronize() + drain_slot_pool_pending_frees() + remaining = slot_pool_active_blocks() + if remaining: + raise RuntimeError( + f"CUDA graph slot checkpoint retained {len(remaining)} active allocations " + f"during {phase}." + ) + return [storage._cdata for storage in owners.values()] + + def verify_checkpoint_live_storages(expected_blocks, owners): + """Verify allocator ownership was restored onto the original StorageImpls.""" + for block_ptr, storage in owners.items(): + if not torch._C._has_Standard_Deleter(storage._cdata): + raise RuntimeError( + "CUDA graph slot checkpoint did not restore the live-storage deleter." + ) + storage_ptr = storage.data_ptr() + if not (block_ptr <= storage_ptr < block_ptr + expected_blocks[block_ptr]): + raise RuntimeError( + "CUDA graph slot checkpoint restored a live storage at the wrong address." + ) + + def restore_slot_pool_boundary( + current_blocks, + current_owners, + target_state, + target_blocks, + target_owners, + phase, + ): + """Switch between two allocator boundaries while preserving their StorageImpls.""" + release_checkpoint_live_storages(current_blocks, current_owners, phase) + torch._C._cuda_setCheckpointPoolState( + torch.cuda.current_device(), + target_state, + [], + [storage._cdata for storage in target_owners.values()], + ) + verify_checkpoint_live_storages(target_blocks, target_owners) + assert_slot_pool_liveness(target_blocks, phase) + if _order is not None: # pylint: disable=too-many-nested-blocks per_callable_static_outputs = [None] * len(flatten_sample_args) per_callable_output_unflatten_spec = [None] * len(flatten_sample_args) @@ -570,6 +2166,51 @@ def hook_fn( wgrad_validation_list = [None] * len(_order) previous_chunk_last_callable_bwd_idx = None for i, c_id in enumerate(_order): + branch_group = branch_capture_groups[i] if branch_capture_groups is not None else None + deferred_native_output_target_releases = set() + deferred_native_grad_target_releases = set() + captured_branch_func_indices = [] + branch_stale_storages = {} + branch_boundary_values = [] + if ( + branch_group is not None + and branch_group[0] == 0 + and ( + branch_checkpoint_state is not None + or branch_pre_checkpoint_state is not None + or current_storage_owners is not None + or branch_canonical_native_saved_intervals is not None + or branch_canonical_native_saved_targets is not None + or branch_canonical_native_saved_excluded_storages is not None + or branch_canonical_native_saved_preassigned_targets is not None + or native_saved_alias_targets is not None + ) + ): + raise RuntimeError("CUDA graph slot checkpoint groups overlap.") + if branch_group is not None: + if branch_group[0] == 0: + gc.collect() + torch.cuda.synchronize() + drain_slot_pool_pending_frees() + branch_pre_checkpoint_state = torch._C._cuda_getCheckpointState( + torch.cuda.current_device(), mempool + ) + branch_pre_checkpoint_live_blocks = slot_pool_active_blocks() + branch_pre_checkpoint_storage_owners = checkpoint_live_storage_owners( + branch_pre_checkpoint_live_blocks, + f"{'forward' if c_id > 0 else 'backward'} pre-branch " + f"boundary at order index {i}", + ) + elif branch_group[0] > 1 and c_id < 0: + restore_slot_pool_boundary( + branch_checkpoint_live_blocks, + branch_checkpoint_storage_owners, + branch_pre_checkpoint_state, + branch_pre_checkpoint_live_blocks, + branch_pre_checkpoint_storage_owners, + f"{'forward' if c_id > 0 else 'backward'} branch " + f"{branch_group[0] + 1}/{branch_group[1]} restore pre-boundary", + ) if c_id > 0: if not isinstance(c_id, int): raise TypeError( @@ -582,15 +2223,57 @@ def hook_fn( per_callable_fwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no ) + captured_branch_func_indices.append(per_callable_fwd_idx) args = sample_args[per_callable_fwd_idx] kwargs = sample_kwargs[per_callable_fwd_idx] fwd_graph = fwd_graphs[per_callable_fwd_idx] + native_saved_alias_targets = None + if branch_group is not None and branch_group[0] > 0: + if ( + branch_canonical_native_saved_targets is None + or branch_canonical_native_saved_excluded_storages is None + or branch_canonical_native_saved_preassigned_targets is None + ): + raise RuntimeError( + "CUDA graph CP branch has no canonical saved-tensor targets." + ) + native_saved_alias_targets = plan_native_saved_alias_targets( + per_callable_saved_tensor_plans[per_callable_fwd_idx], + branch_canonical_native_saved_targets[l_no], + protected_storage_ranges=branch_canonical_native_saved_excluded_storages[ + per_callable_fwd_idx + ], + preassigned_targets=branch_canonical_native_saved_preassigned_targets[ + per_callable_fwd_idx + ], + ) with _graph_context_wrapper(fwd_graph, pool=mempool): - outputs = func(*args, **kwargs) - flatten_outputs, spec = _tree_flatten(outputs) + with capture_saved_tensors( + per_callable_fwd_idx, native_saved_alias_targets + ): + outputs = func(*args, **kwargs) + flatten_outputs, spec = _tree_flatten(outputs) + original_flatten_outputs = flatten_outputs + flatten_outputs = copy_outputs_to_slot_arena( + per_callable_fwd_idx, flatten_outputs + ) + if branch_group is not None and branch_group[0] > 0: + record_replaced_io_storages( + original_flatten_outputs, + flatten_outputs, + branch_stale_storages, + ) + branch_stale_storages.update( + per_callable_native_saved_storages[per_callable_fwd_idx] + ) + del original_flatten_outputs + native_saved_alias_targets = None per_callable_static_outputs[per_callable_fwd_idx] = tuple(flatten_outputs) per_callable_output_unflatten_spec[per_callable_fwd_idx] = spec graph_callables[per_callable_fwd_idx] = func + if use_slot_memory: + del outputs + del flatten_outputs fwd_idx[m_chunk] += 1 else: # Capture backward graph for model chunk c_id, microbatch bwd_idx[-c_id-1] @@ -600,6 +2283,7 @@ def hook_fn( per_callable_bwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no ) + captured_branch_func_indices.append(per_callable_bwd_idx) if ceil(c_id) == c_id and need_bwd_dw_graph[per_callable_bwd_idx]: # Check if bwd graph has corresponding wgrad graph: # Number of dgrad backward graphs should be equal to number of @@ -673,13 +2357,13 @@ def hook_fn( static_grad_outputs = static_grad_outputs_dict[static_grad_outputs_keys] else: static_grad_outputs = tuple( - torch.empty_like(o) if o is not None and o.requires_grad else None + (torch.empty_like(o) if o is not None and o.requires_grad else None) for o in static_outputs ) static_grad_outputs_dict[static_grad_outputs_keys] = static_grad_outputs else: static_grad_outputs = tuple( - torch.empty_like(o) if o is not None and o.requires_grad else None + (torch.empty_like(o) if o is not None and o.requires_grad else None) for o in static_outputs ) if is_training: @@ -695,33 +2379,61 @@ def hook_fn( retain_graph=retain_graph_in_backward, ) grad_inputs = tuple(input.grad for input in inputs) + original_grad_inputs = grad_inputs + grad_inputs = copy_user_grads_to_slot_arena( + per_callable_bwd_idx, static_input_surface, grad_inputs + ) + if branch_group is not None and branch_group[0] > 0: + record_replaced_io_storages( + original_grad_inputs, + grad_inputs, + branch_stale_storages, + ) + del original_grad_inputs # Constructs a tuple suitable for returning from Graphed.backward: # Pads out the actually-needed grads with Nones in gradient slots for inputs # that don't require grad. I couldn't think of a one-liner for this pattern. static_grad_inputs = [] grad_idx = 0 + fused_wgrad_params = per_callable_fused_wgrad_params.get( + per_callable_bwd_idx, set() + ) for arg in static_input_surface: if is_training and isinstance(arg, torch.Tensor) and arg.requires_grad: - static_grad_inputs.append(grad_inputs[grad_idx]) + grad_input = grad_inputs[grad_idx] grad_idx += 1 + if grad_input is None and arg in fused_wgrad_params: + main_grad = getattr(arg, "main_grad", arg) + grad_input = get_dummy_wgrad( + list(main_grad.shape), + arg.dtype, + zero=getattr(arg, "zero_out_wgrad", False), + ) + static_grad_inputs.append(grad_input) else: static_grad_inputs.append(None) # type: ignore[arg-type] static_grad_inputs = tuple(static_grad_inputs) # type: ignore[assignment] + validate_captured_module_grads(per_callable_bwd_idx, static_grad_inputs) per_callable_static_grad_outputs[per_callable_bwd_idx] = static_grad_outputs per_callable_static_grad_inputs[per_callable_bwd_idx] = static_grad_inputs + if branch_group is not None: + branch_boundary_values.append(static_grad_inputs) - # Weak ref the static outputs and static grad inputs that are no longer needed - # in the following steps. These two type of tensors are both in cudagraph - # mempool, so we just deallocate them and let PyTorch's memory allocator - # reuse them elsewhere. + # Weak-ref static output and gradient objects after their capture lifetime. + # Their backing storage remains alive either in the graph pool or an explicit + # slot arena, while transient graph-pool references can be reclaimed. if _reuse_graph_input_output_buffers: # Weak ref the static outputs of the forward pass of this backward. It's # no longer needed after the corresponding backward graph is built up. per_callable_static_outputs[per_callable_bwd_idx] = make_weak_ref( static_outputs ) + if branch_group is None: + clear_native_io_target_rows((per_callable_bwd_idx,), clear_outputs=True) + else: + deferred_native_output_target_releases.add(per_callable_bwd_idx) # Weak ref the static grad inputs of the previous backward pass within the # same chunk. @@ -730,6 +2442,10 @@ def hook_fn( per_callable_static_grad_inputs[idx] = make_weak_ref( per_callable_static_grad_inputs[idx] ) + if branch_group is None: + clear_native_io_target_rows((idx,), clear_grads=True) + else: + deferred_native_grad_target_releases.add(idx) previous_per_callable_bwd_idx = per_callable_bwd_idx # Weak ref the static grad inputs of the previous chunk's last backward @@ -743,9 +2459,117 @@ def hook_fn( per_callable_static_grad_inputs[idx] = make_weak_ref( per_callable_static_grad_inputs[idx] ) + if branch_group is None: + clear_native_io_target_rows((idx,), clear_grads=True) + else: + deferred_native_grad_target_releases.add(idx) previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx + if use_slot_memory and branch_group is None: + per_callable_native_saved_storages[per_callable_bwd_idx].clear() + del static_outputs if ceil(c_id) == c_id: bwd_idx[m_chunk] += 1 + + if branch_group is not None: + if c_id < 0: + for captured_func_idx in captured_branch_func_indices: + per_callable_static_grad_inputs[captured_func_idx] = make_weak_ref( + per_callable_static_grad_inputs[captured_func_idx] + ) + gc.collect() + torch.cuda.synchronize() + drain_slot_pool_pending_frees() + if branch_group[0] == 0: + if c_id > 0: + ( + branch_canonical_native_saved_targets, + branch_canonical_native_saved_excluded_storages, + branch_canonical_native_saved_preassigned_targets, + ) = materialize_native_saved_spill_targets(captured_branch_func_indices) + for func_idx in captured_branch_func_indices: + per_callable_native_saved_capture_targets[func_idx] = None + branch_checkpoint_state = torch._C._cuda_getCheckpointState( + torch.cuda.current_device(), mempool + ) + branch_checkpoint_live_blocks = slot_pool_active_blocks() + branch_checkpoint_pool_layout = slot_pool_layout() + if c_id > 0: + branch_canonical_native_saved_intervals = native_saved_pool_intervals( + captured_branch_func_indices, + branch_checkpoint_pool_layout, + full_storage=True, + ) + branch_checkpoint_storage_owners = checkpoint_live_storage_owners( + branch_checkpoint_live_blocks, + f"{'forward' if c_id > 0 else 'backward'} branch " + f"1/{branch_group[1]} at order index {i}", + (branch_stale_storages, branch_boundary_values), + ) + if c_id < 0: + restore_slot_pool_boundary( + branch_checkpoint_live_blocks, + branch_checkpoint_storage_owners, + branch_pre_checkpoint_state, + branch_pre_checkpoint_live_blocks, + branch_pre_checkpoint_storage_owners, + f"backward branch 1/{branch_group[1]} restore pre-boundary", + ) + else: + current_live_blocks = slot_pool_active_blocks() + current_layout = slot_pool_layout() + if c_id > 0: + assert_native_saved_interval_coverage( + branch_canonical_native_saved_intervals, + native_saved_pool_intervals( + captured_branch_func_indices, current_layout + ), + captured_branch_func_indices, + f"branch {branch_group[0] + 1}/{branch_group[1]} at order index {i}", + ) + current_storage_owners = checkpoint_live_storage_owners( + current_live_blocks, + f"{'forward' if c_id > 0 else 'backward'} branch " + f"{branch_group[0] + 1}/{branch_group[1]} before post-boundary restore", + ( + branch_stale_storages, + branch_boundary_values, + branch_pre_checkpoint_storage_owners, + branch_checkpoint_storage_owners, + ), + ) + restore_slot_pool_boundary( + current_live_blocks, + current_storage_owners, + branch_checkpoint_state, + branch_checkpoint_live_blocks, + branch_checkpoint_storage_owners, + f"{'forward' if c_id > 0 else 'backward'} branch " + f"{branch_group[0] + 1}/{branch_group[1]} restore post-boundary", + ) + current_storage_owners = None + clear_native_io_target_rows( + deferred_native_output_target_releases, clear_outputs=True + ) + clear_native_io_target_rows(deferred_native_grad_target_releases, clear_grads=True) + if c_id < 0: + for captured_func_idx in captured_branch_func_indices: + per_callable_native_saved_storages[captured_func_idx].clear() + if branch_group[0] == branch_group[1] - 1: + branch_checkpoint_state = None + branch_checkpoint_live_blocks = None + branch_checkpoint_pool_layout = None + branch_checkpoint_storage_owners = None + branch_pre_checkpoint_state = None + branch_pre_checkpoint_live_blocks = None + branch_pre_checkpoint_storage_owners = None + branch_canonical_native_saved_intervals = None + branch_canonical_native_saved_targets = None + branch_canonical_native_saved_excluded_storages = None + branch_canonical_native_saved_preassigned_targets = None + gc.collect() + torch.cuda.synchronize() + drain_slot_pool_pending_frees() + else: # Capture forward graphs per_callable_static_outputs = [] @@ -764,7 +2588,13 @@ def hook_fn( # Capture backward graphs in reverse order per_callable_static_grad_outputs = [] per_callable_static_grad_inputs = [] - for static_input_surface, static_outputs, bwd_graph, bwd_dw_graph, bwd_idx in zip( + for ( + static_input_surface, + static_outputs, + bwd_graph, + bwd_dw_graph, + bwd_idx, + ) in zip( reversed(per_callable_static_input_surfaces), reversed(per_callable_static_outputs), reversed(bwd_graphs), @@ -812,6 +2642,21 @@ def hook_fn( # Reverses the most recent two lists per_callable_static_grad_outputs = list(reversed(per_callable_static_grad_outputs)) per_callable_static_grad_inputs = list(reversed(per_callable_static_grad_inputs)) + + if branch_checkpoint_state is not None or branch_pre_checkpoint_state is not None: + raise RuntimeError("CUDA graph capture ended inside a slot checkpoint group.") + + if allocator_settings_to_restore is not None: + torch._C._accelerator_setAllocatorSettings(allocator_settings_to_restore) + + if use_slot_memory and ( + any(native_io_anchors.values()) or any(native_io_capture_counts.values()) + ): + raise RuntimeError( + "CUDA graph capture ended with incomplete native I/O aliases: " + f"anchors={native_io_anchors}, counts={native_io_capture_counts}." + ) + # Now for every per_callable list, per_callable_*[i] holds the stuff for the ith callable. def make_graphed_autograd_function( @@ -830,7 +2675,13 @@ class Graphed(torch.autograd.Function): """Autograd function for graph replay.""" @staticmethod - def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *inputs): + def forward( + ctx, + skip_fp8_weight_update, + cuda_graph_stream, + cuda_graph_event, + *inputs, + ): # pylint: disable=missing-function-docstring # Set flag for whether to update FP8 weight updates @@ -1065,6 +2916,8 @@ def new_fwd(*user_args, **user_kwargs): backward_dw_func, reset_func = make_graphed_attribute_functions(i) setattr(ret[-1], "backward_dw", backward_dw_func) setattr(ret[-1], "reset", reset_func) + if slot_allocator_pool is not None: + setattr(ret[-1], "_te_cuda_graph_allocator_pool", slot_allocator_pool) if just_one_callable: return ret[0] @@ -1143,6 +2996,7 @@ def make_graphed_callables( pool: Optional[Tuple[int, ...]] = None, retain_graph_in_backward: bool = False, _reuse_graph_input_output_buffers: bool = False, + _graph_memory_slots: Optional[Sequence[Tuple[int, ...]]] = None, pre_warmup_hook: Optional[Callable] = None, post_warmup_hook: Optional[Callable] = None, ) -> Union[Callable, Tuple[Callable, ...]]: @@ -1183,6 +3037,13 @@ def make_graphed_callables( graphs. Only supported with Mcore interleaved pipeline parallelism, i.e. when `_order` is provided. All callables in `modules` are assumed to have inputs and outputs with the same dtype and shape. + _graph_memory_slots: sequence of 7- or 9-int tuples, default = None + Private liveness plan for mutually exclusive graph variants. Each tuple describes + saved-tensor, graph-I/O, and warmup alias groups for one graph input. Nine-field plans + additionally provide a frame ID and conflict mask for cross-slot validation. Requires the + first positional sample argument of every graph input to be a plain CUDA tensor; it + is snapshotted into the slot arenas whenever forward saves it for backward, so + shape-identical graph inputs can share one input staging surface. pre_warmup_hook: callable, default = None A hook function that will be called before the warmup iterations. post_warmup_hook: callable, default = None @@ -1378,6 +3239,7 @@ def call_func(self, *args, **kwargs): pool=pool, retain_graph_in_backward=retain_graph_in_backward, _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers, + _graph_memory_slots=_graph_memory_slots, pre_warmup_hook=pre_warmup_hook, post_warmup_hook=post_warmup_hook, ) From 4a7b8fd3a74ffbd5df9b7056c8633427899c2f6c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 12 Aug 2026 09:09:44 +0000 Subject: [PATCH 02/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../attention/dot_product_attention/context_parallel.py | 1 + transformer_engine/pytorch/csrc/extensions/pybind.cpp | 5 ++--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 177f3b38a0..b3b3145106 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -75,6 +75,7 @@ def _get_cp_p2p_transport_group(cp_group): _cp_p2p_transport_groups.pop(cp_group, None) return cp_group, False + # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index cac5fec488..1a86caef57 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -6,15 +6,14 @@ #include "pybind.h" +#include +#include #include #include #include #include #include -#include -#include - #include #include #include From c095214af883862527b5e6d35d978acd9b272844 Mon Sep 17 00:00:00 2001 From: xiaoyao0115 <1804647152@qq.com> Date: Wed, 12 Aug 2026 02:19:27 -0700 Subject: [PATCH 03/16] Fix CUDA graph slot memory lint Signed-off-by: xiaoyao0115 <1804647152@qq.com> --- transformer_engine/pytorch/graph.py | 127 ++++++++++++++++------------ 1 file changed, 71 insertions(+), 56 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 4d99d12215..614426ab8c 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -13,10 +13,10 @@ from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union import torch -import transformer_engine_torch as tex from torch.utils._pytree import tree_flatten as _tree_flatten from torch.utils._pytree import tree_unflatten as _tree_unflatten from torch._C import _graph_pool_handle +import transformer_engine_torch as tex from transformer_engine.common.recipe import DelayedScaling, Recipe from transformer_engine.pytorch.constants import dist_group_type @@ -93,7 +93,7 @@ def _align_up(value: int, alignment: int = 256) -> int: def _io_tensor_plan(tensor: Any, kind: str) -> Optional[Tuple[Any, ...]]: """Return an arena plan for plain CUDA tensors exposed across graph boundaries.""" if ( - type(tensor) is not torch.Tensor + tensor.__class__ is not torch.Tensor or not tensor.is_cuda or tensor.layout != torch.strided or any(stride < 0 for stride in tensor.stride()) @@ -659,7 +659,7 @@ def _make_graphed_callables( ] per_callable_snapshot_input_storage_ptrs = [] for func_idx, args in enumerate(sample_args): - if not args or type(args[0]) is not torch.Tensor or not args[0].is_cuda: + if not args or args[0].__class__ is not torch.Tensor or not args[0].is_cuda: raise RuntimeError( "Slot user-input snapshots require the first positional argument for " f"graph input {func_idx} to be a plain CUDA tensor." @@ -813,6 +813,51 @@ def observe_saved_tensor_boundary_aliases( ) return aliases + def make_saved_tensor_recorder( + func_idx, + observed_saved_tensors, + observed_saved_versions, + copied_storages, + observed_saved_tensor_plan, + ): + """Bind one warmup iteration's saved-tensor observation state.""" + + def record_saved_tensor(tensor): + observed_saved_tensors.append(tensor) + observed_saved_versions.append(_tensor_version(tensor)) + storage_ptr = _tensor_storage_ptr(tensor) + signature = _saved_tensor_signature(tensor) + snapshot_user_input = ( + tensor.is_cuda and storage_ptr == per_callable_snapshot_input_storage_ptrs[func_idx] + ) + is_external = not tensor.is_cuda or ( + storage_ptr in per_callable_external_storage_ptrs[func_idx] + and not snapshot_user_input + ) + if not is_external and tensor.__class__ is not torch.Tensor: + raise RuntimeError( + "CUDA graph saved-tensor arenas do not yet support tensor " + f"subclass {type(tensor).__name__}." + ) + storage_group = None + storage_offset_bytes = None + if not is_external: + storage_identity = (storage_ptr, _tensor_version(tensor)) + storage_group = copied_storages.setdefault(storage_identity, len(copied_storages)) + storage_offset_bytes = tensor.storage_offset() * tensor.element_size() + observed_saved_tensor_plan.append( + ( + "external" if is_external else "native", + storage_ptr if is_external else None, + *signature, + storage_group, + storage_offset_bytes, + ) + ) + return tensor + + return record_saved_tensor + # Run warmup and do the above filtering. with torch.cuda.stream(torch.cuda.Stream()): for func_idx, func in zip(warmup_func_idx, warmup_func): @@ -862,43 +907,13 @@ def hook_fn( observed_saved_tensors = [] observed_saved_versions = [] copied_storages = {} - - def record_saved_tensor(tensor): - observed_saved_tensors.append(tensor) - observed_saved_versions.append(_tensor_version(tensor)) - storage_ptr = _tensor_storage_ptr(tensor) - signature = _saved_tensor_signature(tensor) - snapshot_user_input = ( - tensor.is_cuda - and storage_ptr == per_callable_snapshot_input_storage_ptrs[func_idx] - ) - is_external = not tensor.is_cuda or ( - storage_ptr in per_callable_external_storage_ptrs[func_idx] - and not snapshot_user_input - ) - if not is_external and type(tensor) is not torch.Tensor: - raise RuntimeError( - "CUDA graph saved-tensor arenas do not yet support tensor " - f"subclass {type(tensor).__name__}." - ) - storage_group = None - storage_offset_bytes = None - if not is_external: - storage_identity = (storage_ptr, _tensor_version(tensor)) - storage_group = copied_storages.setdefault( - storage_identity, len(copied_storages) - ) - storage_offset_bytes = tensor.storage_offset() * tensor.element_size() - observed_saved_tensor_plan.append( - ( - "external" if is_external else "native", - storage_ptr if is_external else None, - *signature, - storage_group, - storage_offset_bytes, - ) - ) - return tensor + record_saved_tensor = make_saved_tensor_recorder( + func_idx, + observed_saved_tensors, + observed_saved_versions, + copied_storages, + observed_saved_tensor_plan, + ) with torch.autograd.graph.saved_tensors_hooks(record_saved_tensor, lambda x: x): outputs, _ = _tree_flatten(func(*args, **kwargs)) @@ -1326,7 +1341,7 @@ def copy_user_grads_to_slot_arena(func_idx, static_input_surface, grad_inputs): copied_grad_inputs.append(grad_input) return tuple(copied_grad_inputs) - per_callable_native_saved_storages = [dict() for _ in flatten_sample_args] + per_callable_native_saved_storages = [{} for _ in flatten_sample_args] per_callable_native_saved_intervals = [[] for _ in flatten_sample_args] per_callable_native_saved_capture_targets = [None] * len(flatten_sample_args) @@ -2098,7 +2113,7 @@ def release_checkpoint_live_storages(expected_blocks, owners, phase): raise RuntimeError("CUDA graph slot checkpoint live-storage owner set changed.") for block_ptr, storage in owners.items(): storage_ptr = storage.data_ptr() - if not (block_ptr <= storage_ptr < block_ptr + expected_blocks[block_ptr]): + if not block_ptr <= storage_ptr < block_ptr + expected_blocks[block_ptr]: raise RuntimeError( "CUDA graph slot checkpoint live storage changed its allocation." ) @@ -2131,7 +2146,7 @@ def verify_checkpoint_live_storages(expected_blocks, owners): "CUDA graph slot checkpoint did not restore the live-storage deleter." ) storage_ptr = storage.data_ptr() - if not (block_ptr <= storage_ptr < block_ptr + expected_blocks[block_ptr]): + if not block_ptr <= storage_ptr < block_ptr + expected_blocks[block_ptr]: raise RuntimeError( "CUDA graph slot checkpoint restored a live storage at the wrong address." ) @@ -2172,20 +2187,20 @@ def restore_slot_pool_boundary( captured_branch_func_indices = [] branch_stale_storages = {} branch_boundary_values = [] - if ( - branch_group is not None - and branch_group[0] == 0 - and ( - branch_checkpoint_state is not None - or branch_pre_checkpoint_state is not None - or current_storage_owners is not None - or branch_canonical_native_saved_intervals is not None - or branch_canonical_native_saved_targets is not None - or branch_canonical_native_saved_excluded_storages is not None - or branch_canonical_native_saved_preassigned_targets is not None - or native_saved_alias_targets is not None + branch_checkpoint_active = any( + value is not None + for value in ( + branch_checkpoint_state, + branch_pre_checkpoint_state, + current_storage_owners, + branch_canonical_native_saved_intervals, + branch_canonical_native_saved_targets, + branch_canonical_native_saved_excluded_storages, + branch_canonical_native_saved_preassigned_targets, + native_saved_alias_targets, ) - ): + ) + if branch_group is not None and branch_group[0] == 0 and branch_checkpoint_active: raise RuntimeError("CUDA graph slot checkpoint groups overlap.") if branch_group is not None: if branch_group[0] == 0: From 6047839f07907ceb2fd8b78c90c7fe7ece4061ed Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Mon, 17 Aug 2026 12:37:59 +0900 Subject: [PATCH 04/16] Fix dynamic-CP CUDA graph slot wrap Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 69 ++++++ transformer_engine/pytorch/graph.py | 330 ++++++++++------------------ 2 files changed, 183 insertions(+), 216 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 7ad712f5a5..4548b5b279 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1327,6 +1327,75 @@ def forward(self, inp): reset_graphs(graphed) +def test_slot_memory_saved_arenas_cover_alternate_schedule() -> None: + """Saved tensors must follow union liveness, not only the capture schedule.""" + + class Module(torch.nn.Module): + def forward(self, inp): + hidden = inp.sin() + return hidden.square() + + module = Module().cuda() + samples = tuple((torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(4)) + graphed = make_graphed_callables( + (module,), + samples, + num_warmup_iters=2, + _order=[1, 1, -1, 1, -1, 1, -1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple(_slot(index, index, index) for index in range(4)), + ) + + try: + inputs = [torch.randn(4096, device="cuda", requires_grad=True) for _ in range(4)] + outputs = [graphed[index](inputs[index]) for index in range(3)] + outputs[0].sum().backward() + torch.testing.assert_close( + inputs[0].grad, 2.0 * inputs[0].detach().sin() * inputs[0].detach().cos() + ) + outputs.append(graphed[3](inputs[3])) + for index in (1, 2, 3): + outputs[index].sum().backward() + torch.testing.assert_close( + inputs[index].grad, + 2.0 * inputs[index].detach().sin() * inputs[index].detach().cos(), + ) + finally: + reset_graphs(graphed) + + +def test_slot_memory_reuses_user_grad_surface_across_chunks() -> None: + """Adjacent backward chunks may reuse one physical-slot gradient surface.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() + + module = Module().cuda() + samples = tuple((torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(2)) + graphed = make_graphed_callables( + (module, module), + samples, + num_warmup_iters=2, + _order=[1, 2, -2, -1], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=( + _slot(0, 0, 0, overlap=0, warmup=0), + _slot(1, 1, 0, overlap=1, warmup=1), + ), + ) + + try: + inp = torch.randn(4096, device="cuda", requires_grad=True) + graphed[1](graphed[0](inp)).sum().backward() + torch.testing.assert_close(inp.grad, 4.0 * inp.detach().pow(3)) + assert tuple(graphed[-1]._te_cuda_graph_user_grad_arenas) == (0,) + finally: + reset_graphs(graphed) + + def test_slot_memory_coalesces_overlapping_saved_views() -> None: """Saved views of one storage should occupy only their byte union in each live slot.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 614426ab8c..7350dc03ac 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -361,7 +361,6 @@ def _make_graphed_callables( ) saved_tensor_memory_alias_groups = None - saved_tensor_memory_families = None slot_io_memory_alias_groups = None slot_io_liveness_groups = None warmup_plan_alias_groups = None @@ -386,7 +385,6 @@ def _make_graphed_callables( ): raise TypeError("Each graph-memory slot must be a tuple of nine integers.") saved_tensor_memory_alias_groups = [(slot[0], slot[1]) for slot in _graph_memory_slots] - saved_tensor_memory_families = [(slot[2], slot[4], slot[5]) for slot in _graph_memory_slots] slot_io_memory_alias_groups = [(slot[2], slot[3]) for slot in _graph_memory_slots] slot_io_liveness_groups = [(slot[4], slot[5]) for slot in _graph_memory_slots] warmup_plan_alias_groups = [slot[6] for slot in _graph_memory_slots] @@ -1342,8 +1340,6 @@ def copy_user_grads_to_slot_arena(func_idx, static_input_surface, grad_inputs): return tuple(copied_grad_inputs) per_callable_native_saved_storages = [{} for _ in flatten_sample_args] - per_callable_native_saved_intervals = [[] for _ in flatten_sample_args] - per_callable_native_saved_capture_targets = [None] * len(flatten_sample_args) def plan_native_saved_alias_targets( plan, @@ -1507,10 +1503,10 @@ def plan_native_saved_alias_targets( raise RuntimeError(f"Native saved tensors have no canonical targets: {missing}.") return tuple(target_views) - def semantic_boundary_alias_targets(canonical_func_idx, sibling_func_idx): - """Map boundary-backed sibling saves onto the boundary address used at replay.""" - plan = per_callable_saved_tensor_plans[sibling_func_idx] - aliases = per_callable_saved_tensor_boundary_aliases[sibling_func_idx] + def semantic_boundary_alias_targets(func_idx, outputs): + """Map boundary-backed saves onto the boundary address used at replay.""" + plan = per_callable_saved_tensor_plans[func_idx] + aliases = per_callable_saved_tensor_boundary_aliases[func_idx] targets = [None] * len(plan) records_by_storage_group = {} for saved_idx, spec in enumerate(plan): @@ -1564,9 +1560,9 @@ def semantic_boundary_alias_targets(canonical_func_idx, sibling_func_idx): for saved_idx, alias in component_aliases: _, kind, boundary_idx, relative_offset, _ = alias if kind == "input": - boundary = per_callable_static_input_surfaces[sibling_func_idx][boundary_idx] + boundary = per_callable_static_input_surfaces[func_idx][boundary_idx] else: - boundary = per_callable_static_outputs[canonical_func_idx][boundary_idx] + boundary = outputs[boundary_idx] if not isinstance(boundary, torch.Tensor) or not boundary.is_cuda: raise RuntimeError( f"CUDA graph {kind} boundary {boundary_idx} is not a CUDA tensor." @@ -1582,7 +1578,7 @@ def semantic_boundary_alias_targets(canonical_func_idx, sibling_func_idx): elif anchor_storage._cdata != storage._cdata or anchor_shift != shift: raise RuntimeError( "CUDA graph overlapping saved tensors have inconsistent boundary " - f"aliases: func={sibling_func_idx}, saved={component_saved_indices}." + f"aliases: func={func_idx}, saved={component_saved_indices}." ) for saved_idx in component_saved_indices: @@ -1591,7 +1587,7 @@ def semantic_boundary_alias_targets(canonical_func_idx, sibling_func_idx): if target_offset < 0 or target_offset + spec[7] > anchor_storage.nbytes(): raise RuntimeError( "CUDA graph boundary-backed saved component does not fit its replay " - f"storage: func={sibling_func_idx}, saved={saved_idx}, " + f"storage: func={func_idx}, saved={saved_idx}, " f"offset={target_offset}, bytes={spec[7]}, " f"storage_bytes={anchor_storage.nbytes()}." ) @@ -1610,80 +1606,99 @@ def semantic_boundary_alias_targets(canonical_func_idx, sibling_func_idx): targets[saved_idx] = target return tuple(targets) - def materialize_native_saved_spill_targets(canonical_func_indices): - """Complete canonical saved-tensor storage for every same-slot CP branch.""" - completed_targets = [] - protected_ranges = {} - preassigned_targets = {} - - for func_idx in canonical_func_indices: - targets = per_callable_native_saved_capture_targets[func_idx] - if targets is None: - raise RuntimeError("CUDA graph canonical CP branch did not retain saved targets.") - - storage_ranges = {} - boundary_tensors = ( - *per_callable_static_input_surfaces[func_idx][ - : per_callable_len_user_args[func_idx] - ], - *per_callable_static_outputs[func_idx], + def slot_tensor_targets(plan, arena=None): + """Lay out graph-boundary tensors contiguously in an arena.""" + targets = [] + offset = 0 + for spec in plan: + if spec is None: + targets.append(None) + continue + offset = _align_up(offset) + target = None + if arena is not None: + target = _arena_view(arena, offset, spec) + targets.append(target) + offset += spec[7] + return tuple(targets), _align_up(offset) + + slot_saved_arenas = {} + per_callable_slot_saved_targets = None + if use_slot_memory: + arena_sizes = {} + for func_idx, plan in enumerate(per_callable_saved_tensor_plans): + output_plan = per_callable_output_tensor_plans[func_idx] + _, output_bytes = slot_tensor_targets(output_plan) + temporary_arena = None + if output_bytes: + with torch.cuda.use_mem_pool(slot_allocator_pool): + temporary_arena = torch.empty( + (output_bytes,), dtype=torch.uint8, device=torch.cuda.current_device() + ) + temporary_outputs, _ = slot_tensor_targets(output_plan, temporary_arena) + preassigned_targets = semantic_boundary_alias_targets(func_idx, temporary_outputs) + spill_bytes = plan_native_saved_alias_targets( + plan, + (), + measure_spill=True, + preassigned_targets=preassigned_targets, ) - for tensor in boundary_tensors: - if not isinstance(tensor, torch.Tensor) or not tensor.is_cuda: - continue - storage = tensor.untyped_storage() - storage_ranges.setdefault(storage._cdata, []).append((0, storage.nbytes())) - for storage_id, ranges in storage_ranges.items(): - merged = [] - for start, end in sorted(ranges): - if not merged or start > merged[-1][1]: - merged.append([start, end]) - else: - merged[-1][1] = max(merged[-1][1], end) - storage_ranges[storage_id] = tuple(map(tuple, merged)) - - family = saved_tensor_memory_families[func_idx] - sibling_indices = [ - sibling_idx - for sibling_idx, sibling_family in enumerate(saved_tensor_memory_families) - if sibling_family == family - ] - for sibling_idx in sibling_indices: - protected_ranges[sibling_idx] = storage_ranges - preassigned_targets[sibling_idx] = ( - semantic_boundary_alias_targets(func_idx, sibling_idx) - if sibling_idx != func_idx - else None + arena_id, _ = saved_tensor_memory_alias_groups[func_idx] + arena_sizes[arena_id] = max(arena_sizes.get(arena_id, 0), output_bytes + spill_bytes) + del temporary_outputs, temporary_arena, preassigned_targets + + with torch.cuda.use_mem_pool(slot_allocator_pool): + slot_saved_arenas = { + arena_id: torch.empty( + (required_bytes,), dtype=torch.uint8, device=torch.cuda.current_device() ) + for arena_id, required_bytes in arena_sizes.items() + if required_bytes > 0 + } - spill_bytes = max( - ( - plan_native_saved_alias_targets( - per_callable_saved_tensor_plans[sibling_idx], - targets, - measure_spill=True, - protected_storage_ranges=storage_ranges, - preassigned_targets=preassigned_targets[sibling_idx], - ) - for sibling_idx in sibling_indices - if sibling_idx != func_idx - ), - default=0, + per_callable_slot_saved_targets = [] + for func_idx, plan in enumerate(per_callable_saved_tensor_plans): + arena_id, _ = saved_tensor_memory_alias_groups[func_idx] + arena = slot_saved_arenas.get(arena_id) + output_targets, output_bytes = slot_tensor_targets( + per_callable_output_tensor_plans[func_idx], arena ) - if spill_bytes: - with torch.cuda.use_mem_pool(slot_allocator_pool): - spill = torch.empty( - (spill_bytes,), - dtype=torch.uint8, - device=torch.cuda.current_device(), - ) - storage = spill.untyped_storage() - per_callable_native_saved_storages[func_idx].setdefault( - storage._cdata, (storage, storage.data_ptr()) + per_callable_output_tensor_targets[func_idx] = list(output_targets) + preassigned_targets = semantic_boundary_alias_targets(func_idx, output_targets) + protected_ranges = {} + if arena is not None and output_bytes: + protected_ranges[arena.untyped_storage()._cdata] = ((0, output_bytes),) + canonical_targets = () if arena is None else (arena,) + per_callable_slot_saved_targets.append( + plan_native_saved_alias_targets( + plan, + canonical_targets, + protected_storage_ranges=protected_ranges, + preassigned_targets=preassigned_targets, + ) + ) + + slot_user_grad_arenas = {} + if use_slot_memory: + slot_sizes = {} + for func_idx, plan in enumerate(per_callable_user_grad_tensor_plans): + _, required_bytes = slot_tensor_targets(plan) + physical_slot, _ = slot_io_memory_alias_groups[func_idx] + slot_sizes[physical_slot] = max(slot_sizes.get(physical_slot, 0), required_bytes) + + with torch.cuda.use_mem_pool(slot_allocator_pool): + slot_user_grad_arenas = { + physical_slot: torch.empty( + (required_bytes,), dtype=torch.uint8, device=torch.cuda.current_device() ) - targets = (*targets, spill) - completed_targets.append(tuple(targets)) - return tuple(completed_targets), protected_ranges, preassigned_targets + for physical_slot, required_bytes in slot_sizes.items() + if required_bytes > 0 + } + + for func_idx, plan in enumerate(per_callable_user_grad_tensor_plans): + physical_slot, _ = slot_io_memory_alias_groups[func_idx] + targets, _ = slot_tensor_targets(plan, slot_user_grad_arenas.get(physical_slot)) + per_callable_user_grad_tensor_targets[func_idx] = list(targets) @contextlib.contextmanager def capture_saved_tensors(func_idx, alias_targets=None): @@ -1694,8 +1709,6 @@ def capture_saved_tensors(func_idx, alias_targets=None): plan = per_callable_saved_tensor_plans[func_idx] saved_idx = 0 - captured_targets = [None] * len(plan) - per_callable_native_saved_intervals[func_idx].clear() if alias_targets is not None and len(alias_targets) != len(plan): raise RuntimeError( f"CUDA graph input {func_idx} changed its canonical saved-target count." @@ -1727,6 +1740,10 @@ def pack_saved_tensor(tensor): f"CUDA graph input {func_idx} has unsupported saved-tensor mode {spec[0]}." ) + source_storage = tensor.untyped_storage() + per_callable_native_saved_storages[func_idx].setdefault( + source_storage._cdata, (source_storage, source_storage.data_ptr()) + ) if alias_targets is None: target = torch.empty((0,), dtype=tensor.dtype, device=tensor.device).set_( tensor.untyped_storage(), @@ -1753,18 +1770,10 @@ def pack_saved_tensor(tensor): target.copy_(tensor) tensor = target - captured_targets[current_saved_idx] = target storage = tensor.untyped_storage() per_callable_native_saved_storages[func_idx].setdefault( storage._cdata, (storage, storage.data_ptr()) ) - if spec[7]: - start = _tensor_storage_ptr(tensor) + ( - tensor.storage_offset() * tensor.element_size() - ) - per_callable_native_saved_intervals[func_idx].append( - (current_saved_idx, start, start + spec[7]) - ) return tensor with torch.autograd.graph.saved_tensors_hooks(pack_saved_tensor, lambda x: x): @@ -1774,8 +1783,6 @@ def pack_saved_tensor(tensor): f"CUDA graph input {func_idx} saved {saved_idx} forward tensors during " f"capture, but saved {len(plan)} during warmup." ) - if alias_targets is None: - per_callable_native_saved_capture_targets[func_idx] = tuple(captured_targets) def validate_captured_module_grads(func_idx, static_grad_inputs): """Require capture to preserve every parameter gradient observed during warmup.""" @@ -1813,16 +1820,11 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): branch_capture_groups = None branch_checkpoint_state = None branch_checkpoint_live_blocks = None - branch_checkpoint_pool_layout = None branch_checkpoint_storage_owners = None branch_pre_checkpoint_state = None branch_pre_checkpoint_live_blocks = None branch_pre_checkpoint_storage_owners = None current_storage_owners = None - branch_canonical_native_saved_intervals = None - branch_canonical_native_saved_targets = None - branch_canonical_native_saved_excluded_storages = None - branch_canonical_native_saved_preassigned_targets = None native_saved_alias_targets = None if use_slot_memory: branch_capture_groups = [None] * len(_order) @@ -1900,66 +1902,6 @@ def slot_pool_layout(): for segment in segments } - def native_saved_pool_intervals(func_indices, layout, full_storage=False): - """Return native saved-tensor intervals that belong to the slot pool.""" - segment_ranges = tuple( - (address, address + segment["total_size"]) for address, segment in layout.items() - ) - output = [] - for func_idx in func_indices: - intervals = [] - if full_storage: - candidates = ( - (storage_idx, storage_ptr, storage_ptr + storage.nbytes()) - for storage_idx, (storage, storage_ptr) in enumerate( - per_callable_native_saved_storages[func_idx].values() - ) - ) - else: - candidates = iter(per_callable_native_saved_intervals[func_idx]) - for saved_idx, start, end in candidates: - containing = [ - (segment_start, segment_end) - for segment_start, segment_end in segment_ranges - if segment_start <= start and end <= segment_end - ] - if containing: - intervals.append((saved_idx, start, end)) - continue - if any( - start < segment_end and segment_start < end - for segment_start, segment_end in segment_ranges - ): - raise RuntimeError( - "CUDA graph native saved tensor crosses a slot-pool segment boundary: " - f"func={func_idx}, saved={saved_idx}, interval=({start}, {end})." - ) - output.append(tuple(intervals)) - return tuple(output) - - def assert_native_saved_interval_coverage(canonical, current, func_indices, phase): - """Require alternate saved tensors to stay inside canonical live allocations.""" - if len(canonical) != len(current): - raise RuntimeError("CUDA graph CP branch changed its captured layer count.") - for position, (canonical_intervals, current_intervals) in enumerate( - zip(canonical, current) - ): - merged = [] - for _, start, end in sorted(canonical_intervals, key=lambda item: item[1:]): - if merged and start <= merged[-1][1]: - merged[-1] = (merged[-1][0], max(merged[-1][1], end)) - else: - merged.append((start, end)) - for saved_idx, start, end in current_intervals: - if any(left <= start and end <= right for left, right in merged): - continue - raise RuntimeError( - "CUDA graph CP branch placed a native saved tensor outside the " - "canonical slot/layer live range: " - f"phase={phase}, func={func_indices[position]}, saved={saved_idx}, " - f"interval=({start}, {end}), canonical={merged}." - ) - def drain_slot_pool_pending_frees(): """Poll completed cross-stream frees before restoring allocator state.""" layout = slot_pool_layout() @@ -2089,6 +2031,8 @@ def visit(value): per_callable_param_grad_tensor_targets, static_grad_outputs_dict, native_io_anchors, + slot_saved_arenas, + slot_user_grad_arenas, *extra_values, ): visit(value) @@ -2193,10 +2137,6 @@ def restore_slot_pool_boundary( branch_checkpoint_state, branch_pre_checkpoint_state, current_storage_owners, - branch_canonical_native_saved_intervals, - branch_canonical_native_saved_targets, - branch_canonical_native_saved_excluded_storages, - branch_canonical_native_saved_preassigned_targets, native_saved_alias_targets, ) ) @@ -2242,26 +2182,11 @@ def restore_slot_pool_boundary( args = sample_args[per_callable_fwd_idx] kwargs = sample_kwargs[per_callable_fwd_idx] fwd_graph = fwd_graphs[per_callable_fwd_idx] - native_saved_alias_targets = None - if branch_group is not None and branch_group[0] > 0: - if ( - branch_canonical_native_saved_targets is None - or branch_canonical_native_saved_excluded_storages is None - or branch_canonical_native_saved_preassigned_targets is None - ): - raise RuntimeError( - "CUDA graph CP branch has no canonical saved-tensor targets." - ) - native_saved_alias_targets = plan_native_saved_alias_targets( - per_callable_saved_tensor_plans[per_callable_fwd_idx], - branch_canonical_native_saved_targets[l_no], - protected_storage_ranges=branch_canonical_native_saved_excluded_storages[ - per_callable_fwd_idx - ], - preassigned_targets=branch_canonical_native_saved_preassigned_targets[ - per_callable_fwd_idx - ], - ) + native_saved_alias_targets = ( + per_callable_slot_saved_targets[per_callable_fwd_idx] + if use_slot_memory + else None + ) with _graph_context_wrapper(fwd_graph, pool=mempool): with capture_saved_tensors( per_callable_fwd_idx, native_saved_alias_targets @@ -2272,15 +2197,16 @@ def restore_slot_pool_boundary( flatten_outputs = copy_outputs_to_slot_arena( per_callable_fwd_idx, flatten_outputs ) - if branch_group is not None and branch_group[0] > 0: + if branch_group is not None: record_replaced_io_storages( original_flatten_outputs, flatten_outputs, branch_stale_storages, ) - branch_stale_storages.update( - per_callable_native_saved_storages[per_callable_fwd_idx] - ) + if branch_group[0] > 0: + branch_stale_storages.update( + per_callable_native_saved_storages[per_callable_fwd_idx] + ) del original_flatten_outputs native_saved_alias_targets = None per_callable_static_outputs[per_callable_fwd_idx] = tuple(flatten_outputs) @@ -2495,25 +2421,10 @@ def restore_slot_pool_boundary( torch.cuda.synchronize() drain_slot_pool_pending_frees() if branch_group[0] == 0: - if c_id > 0: - ( - branch_canonical_native_saved_targets, - branch_canonical_native_saved_excluded_storages, - branch_canonical_native_saved_preassigned_targets, - ) = materialize_native_saved_spill_targets(captured_branch_func_indices) - for func_idx in captured_branch_func_indices: - per_callable_native_saved_capture_targets[func_idx] = None branch_checkpoint_state = torch._C._cuda_getCheckpointState( torch.cuda.current_device(), mempool ) branch_checkpoint_live_blocks = slot_pool_active_blocks() - branch_checkpoint_pool_layout = slot_pool_layout() - if c_id > 0: - branch_canonical_native_saved_intervals = native_saved_pool_intervals( - captured_branch_func_indices, - branch_checkpoint_pool_layout, - full_storage=True, - ) branch_checkpoint_storage_owners = checkpoint_live_storage_owners( branch_checkpoint_live_blocks, f"{'forward' if c_id > 0 else 'backward'} branch " @@ -2531,16 +2442,6 @@ def restore_slot_pool_boundary( ) else: current_live_blocks = slot_pool_active_blocks() - current_layout = slot_pool_layout() - if c_id > 0: - assert_native_saved_interval_coverage( - branch_canonical_native_saved_intervals, - native_saved_pool_intervals( - captured_branch_func_indices, current_layout - ), - captured_branch_func_indices, - f"branch {branch_group[0] + 1}/{branch_group[1]} at order index {i}", - ) current_storage_owners = checkpoint_live_storage_owners( current_live_blocks, f"{'forward' if c_id > 0 else 'backward'} branch " @@ -2572,15 +2473,10 @@ def restore_slot_pool_boundary( if branch_group[0] == branch_group[1] - 1: branch_checkpoint_state = None branch_checkpoint_live_blocks = None - branch_checkpoint_pool_layout = None branch_checkpoint_storage_owners = None branch_pre_checkpoint_state = None branch_pre_checkpoint_live_blocks = None branch_pre_checkpoint_storage_owners = None - branch_canonical_native_saved_intervals = None - branch_canonical_native_saved_targets = None - branch_canonical_native_saved_excluded_storages = None - branch_canonical_native_saved_preassigned_targets = None gc.collect() torch.cuda.synchronize() drain_slot_pool_pending_frees() @@ -2933,6 +2829,8 @@ def new_fwd(*user_args, **user_kwargs): setattr(ret[-1], "reset", reset_func) if slot_allocator_pool is not None: setattr(ret[-1], "_te_cuda_graph_allocator_pool", slot_allocator_pool) + setattr(ret[-1], "_te_cuda_graph_saved_arenas", slot_saved_arenas) + setattr(ret[-1], "_te_cuda_graph_user_grad_arenas", slot_user_grad_arenas) if just_one_callable: return ret[0] From 85e4dcdd77fb58563cec7bae562a218574ad86ad Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Wed, 19 Aug 2026 10:49:33 +0900 Subject: [PATCH 05/16] Speed up CUDA graph slot owner lookup Signed-off-by: Tailai Ma --- transformer_engine/pytorch/graph.py | 48 ++++++++++++++--------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 7350dc03ac..134c46e790 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -4,6 +4,7 @@ """Functions for CUDA Graphs support in FP8""" +from bisect import bisect_right from collections.abc import Iterable import contextlib import gc @@ -1970,41 +1971,40 @@ def checkpoint_live_storage_owners(expected_blocks, phase, extra_values=()): """Resolve every checkpoint-live allocation to its owning StorageImpl.""" owners = {} visited = set() + visited_storage_impls = set() + block_starts = sorted(expected_blocks) + + def record_storage(storage): + if storage.device.type != "cuda" or storage._cdata in visited_storage_impls: + return + visited_storage_impls.add(storage._cdata) + storage_ptr = storage.data_ptr() + block_idx = bisect_right(block_starts, storage_ptr) - 1 + if block_idx < 0: + return + block_ptr = block_starts[block_idx] + if storage_ptr >= block_ptr + expected_blocks[block_ptr]: + return + if torch._C._has_Standard_Deleter(storage._cdata): + previous = owners.setdefault(block_ptr, storage) + if previous._cdata != storage._cdata: + raise RuntimeError( + "CUDA graph slot checkpoint found multiple owning storages " + f"for allocation {block_ptr}." + ) def visit(value): if value is None or id(value) in visited: return if isinstance(value, torch.UntypedStorage): - if value.device.type != "cuda": - return - storage_ptr = value.data_ptr() - for block_ptr, block_size in expected_blocks.items(): - if block_ptr <= storage_ptr < block_ptr + block_size: - if torch._C._has_Standard_Deleter(value._cdata): - previous = owners.setdefault(block_ptr, value) - if previous._cdata != value._cdata: - raise RuntimeError( - "CUDA graph slot checkpoint found multiple owning storages " - f"for allocation {block_ptr}." - ) - break + record_storage(value) return visited.add(id(value)) if isinstance(value, torch.Tensor): if not value.is_cuda: return storage = value.untyped_storage() - storage_ptr = storage.data_ptr() - for block_ptr, block_size in expected_blocks.items(): - if block_ptr <= storage_ptr < block_ptr + block_size: - if torch._C._has_Standard_Deleter(storage._cdata): - previous = owners.setdefault(block_ptr, storage) - if previous._cdata != storage._cdata: - raise RuntimeError( - "CUDA graph slot checkpoint found multiple owning storages " - f"for allocation {block_ptr}." - ) - break + record_storage(storage) if value.is_leaf: visit(value.grad) for child in vars(value).values(): From 9b48ea57cc34f836c89b9e47f28cbc3a153cbeef Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Fri, 21 Aug 2026 17:14:14 +0900 Subject: [PATCH 06/16] Harden CUDA graph slot capture lifecycle Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 152 +++++++++++- transformer_engine/pytorch/graph.py | 346 +++++++++++++++++----------- 2 files changed, 358 insertions(+), 140 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 4548b5b279..593c271efc 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. +import contextlib import gc import weakref from typing import Callable, Dict, Iterable, List, Tuple, Union @@ -29,6 +30,7 @@ set_cp_p2p_transport_group, ) import transformer_engine.pytorch.ops as te_ops +import transformer_engine.pytorch.graph as te_graph from transformer_engine.common import recipe from utils import ModelConfig, reset_rng_states @@ -777,7 +779,153 @@ def test_make_graphed_callables_with_interleaved_pipeline_parallelism( def _slot(saved_arena, branch, io_arena, overlap=0, frame=0, warmup=0): """Build one private graph-memory slot used by the focused tests below.""" - return (saved_arena, branch, io_arena, branch, overlap, frame, warmup, 0, 0) + return (saved_arena, io_arena, branch, overlap, frame, warmup) + + +def test_graph_capture_contexts_restore_process_state_on_error(monkeypatch) -> None: + """Capture failures must restore GC and input gradients.""" + gc_was_enabled = gc.isenabled() + gc.enable() + monkeypatch.setattr(torch.cuda, "graph", lambda *args, **kwargs: contextlib.nullcontext()) + try: + with pytest.raises(RuntimeError, match="capture failed"): + with te_graph._graph_context_wrapper(): + raise RuntimeError("capture failed") + assert gc.isenabled() + finally: + if not gc_was_enabled: + gc.disable() + + inp = torch.ones(1, requires_grad=True) + original_grad = torch.full_like(inp, 2.0) + inp.grad = original_grad + with pytest.raises(RuntimeError, match="capture failed"): + with te_graph._none_grad_context_wrapper((inp,)): + assert inp.grad is None + raise RuntimeError("capture failed") + assert inp.grad is original_grad + + +def test_temporary_forward_hooks_are_removed_on_error() -> None: + """Warmup failures must not leave hooks installed on user modules.""" + module = torch.nn.Sequential(torch.nn.Identity()) + + with pytest.raises(RuntimeError, match="warmup failed"): + with te_graph._module_forward_hooks(module.modules(), lambda *args: None): + assert module._forward_hooks + assert module[0]._forward_hooks + raise RuntimeError("warmup failed") + + assert not module._forward_hooks + assert not module[0]._forward_hooks + + +def test_allocator_settings_guard_restores_once() -> None: + """Temporary allocator settings have idempotent failure cleanup.""" + settings = [] + guard = te_graph._AllocatorSettingsGuard() + + guard.apply(settings.append, "expandable_segments:False", "expandable_segments:True") + guard.restore() + guard.restore() + + assert settings == ["expandable_segments:False", "expandable_segments:True"] + + +def test_make_graphed_callables_restores_process_state_on_error(monkeypatch) -> None: + """The public graph API must unwind every process-wide capture mutation.""" + + class TestModule(torch.nn.Module): + def forward(self, inp): + return inp + + module = TestModule() + original_call = TestModule.__call__ + fp8_state = object() + rng_state = object() + restored_fp8 = [] + restored_rng = [] + allocator_settings = [] + + monkeypatch.setattr(te_graph, "save_fp8_tensors", lambda *args, **kwargs: fp8_state) + monkeypatch.setattr( + te_graph, + "restore_fp8_tensors", + lambda modules, state: restored_fp8.append((modules, state)), + ) + monkeypatch.setattr(te_graph, "graph_safe_rng_available", lambda: False) + monkeypatch.setattr(torch.cuda, "get_rng_state", lambda: rng_state) + monkeypatch.setattr(torch.cuda, "set_rng_state", restored_rng.append) + + def fail_capture(*args, **kwargs): + assert te_graph.is_graph_capturing() + kwargs["_allocator_settings_guard"].apply( + allocator_settings.append, + "expandable_segments:False", + "expandable_segments:True", + ) + raise RuntimeError("capture failed") + + monkeypatch.setattr(te_graph, "_make_graphed_callables", fail_capture) + + assert not te_graph.is_graph_capturing() + with pytest.raises(RuntimeError, match="capture failed"): + te_graph.make_graphed_callables(module, (torch.ones(1),)) + + assert not te_graph.is_graph_capturing() + assert TestModule.__call__ is original_call + assert restored_fp8 == [((module,), fp8_state)] + assert restored_rng == [rng_state] + assert allocator_settings == ["expandable_segments:False", "expandable_segments:True"] + + +def test_make_graphed_callables_restores_wrappers_on_preparation_error(monkeypatch) -> None: + """Preparation failures before capture starts must also restore global wrappers.""" + + class TestModule(torch.nn.Module): + def forward(self, inp): + return inp + + module = TestModule() + original_call = TestModule.__call__ + fp8_state = object() + restored_fp8 = [] + + monkeypatch.setattr(te_graph, "save_fp8_tensors", lambda *args, **kwargs: fp8_state) + monkeypatch.setattr( + te_graph, + "restore_fp8_tensors", + lambda modules, state: restored_fp8.append((modules, state)), + ) + + def fail_rng_preparation(): + raise RuntimeError("rng preparation failed") + + monkeypatch.setattr(te_graph, "graph_safe_rng_available", fail_rng_preparation) + + with pytest.raises(RuntimeError, match="rng preparation failed"): + te_graph.make_graphed_callables(module, (torch.ones(1),)) + + assert not te_graph.is_graph_capturing() + assert TestModule.__call__ is original_call + assert restored_fp8 == [((module,), fp8_state)] + + +def test_slot_memory_rejects_non_native_allocator(monkeypatch) -> None: + """Allocator checkpoints are only defined for the native caching allocator.""" + + module = torch.nn.Identity() + monkeypatch.setattr(torch.cuda.memory, "get_allocator_backend", lambda: "cudaMallocAsync") + + with pytest.raises(RuntimeError, match="requires the native CUDA caching allocator"): + te_graph._make_graphed_callables( + module, + (torch.ones(1),), + _order=[1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0),), + ) @pytest.mark.parametrize("elements", (0, 4096), ids=("empty", "nonempty")) @@ -1133,7 +1281,7 @@ def forward(self, inp): def test_slot_memory_native_io_aliases_graph_pool_storage() -> None: - """Nine-field DCP slots default to forked graph-pool I/O aliases.""" + """DCP slot plans default to forked graph-pool I/O aliases.""" class Module(torch.nn.Module): def forward(self, inp): diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 134c46e790..56c9e256f9 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -42,6 +42,31 @@ SingleOrTuple = Union[_T, Tuple[_T, ...]] +class _AllocatorSettingsGuard: + """Restore temporary allocator settings even when graph capture fails.""" + + def __init__(self) -> None: + self._setter = None + self._settings_to_restore = None + + def apply(self, setter: Callable[[str], None], settings: str, restore: str) -> None: + if self._setter is not None: + raise RuntimeError("CUDA allocator settings guard is already active.") + self._setter = setter + self._settings_to_restore = restore + setter(settings) + + def restore(self) -> None: + if self._setter is None: + return + setter = self._setter + settings = self._settings_to_restore + assert settings is not None + setter(settings) + self._setter = None + self._settings_to_restore = None + + def _tensor_storage_ptr(tensor: torch.Tensor) -> int: """Return the base storage pointer used to recognize static graph inputs.""" return tensor.untyped_storage().data_ptr() @@ -145,12 +170,14 @@ def _none_grad_context_wrapper(inputs): in case the backward pass makes grad accumulations. """ original_input_grads = [] - for input_tensor in inputs: - original_input_grads.append(input_tensor.grad) - input_tensor.grad = None - yield - for input_tensor, original_grad in zip(inputs, original_input_grads): - input_tensor.grad = original_grad + try: + for input_tensor in inputs: + original_input_grads.append(input_tensor.grad) + input_tensor.grad = None + yield + finally: + for input_tensor, original_grad in zip(inputs, original_input_grads): + input_tensor.grad = original_grad @contextlib.contextmanager @@ -166,10 +193,25 @@ def _graph_context_wrapper(*args, **kwargs): gc_is_enabled = gc.isenabled() if gc_is_enabled: gc.disable() - with torch.cuda.graph(*args, **kwargs): + try: + with torch.cuda.graph(*args, **kwargs): + yield + finally: + if gc_is_enabled: + gc.enable() + + +@contextlib.contextmanager +def _module_forward_hooks(modules, hook_fn): + """Remove temporary warmup hooks even when a module raises.""" + hooks = [] + try: + for module in modules: + hooks.append(module.register_forward_hook(hook_fn)) yield - if gc_is_enabled: - gc.enable() + finally: + for hook in reversed(hooks): + hook.remove() def _make_graphed_callables( @@ -185,6 +227,7 @@ def _make_graphed_callables( retain_graph_in_backward: bool = False, _reuse_graph_input_output_buffers: bool = False, _graph_memory_slots: Optional[Sequence[Tuple[int, ...]]] = None, + _allocator_settings_guard: Optional[_AllocatorSettingsGuard] = None, pre_warmup_hook: Optional[Callable] = None, post_warmup_hook: Optional[Callable] = None, ) -> SingleOrTuple[Callable]: @@ -330,6 +373,12 @@ def _make_graphed_callables( use_slot_memory = _graph_memory_slots is not None if use_slot_memory: + allocator_backend = torch.cuda.memory.get_allocator_backend() + if allocator_backend != "native": + raise RuntimeError( + "CUDA graph slot-branch checkpointing requires the native CUDA caching " + f"allocator, but the active backend is {allocator_backend!r}." + ) required_checkpoint_apis = ( "_cuda_getCheckpointState", "_cuda_setCheckpointPoolState", @@ -361,7 +410,7 @@ def _make_graphed_callables( "`_reuse_graph_input_output_buffers` is only available in training mode." ) - saved_tensor_memory_alias_groups = None + saved_tensor_arena_ids = None slot_io_memory_alias_groups = None slot_io_liveness_groups = None warmup_plan_alias_groups = None @@ -380,15 +429,15 @@ def _make_graphed_callables( ) if any( not isinstance(slot, tuple) - or len(slot) != 9 + or len(slot) != 6 or not all(isinstance(value, int) for value in slot) for slot in _graph_memory_slots ): - raise TypeError("Each graph-memory slot must be a tuple of nine integers.") - saved_tensor_memory_alias_groups = [(slot[0], slot[1]) for slot in _graph_memory_slots] - slot_io_memory_alias_groups = [(slot[2], slot[3]) for slot in _graph_memory_slots] - slot_io_liveness_groups = [(slot[4], slot[5]) for slot in _graph_memory_slots] - warmup_plan_alias_groups = [slot[6] for slot in _graph_memory_slots] + raise TypeError("Each graph-memory slot must be a tuple of six integers.") + saved_tensor_arena_ids = [slot[0] for slot in _graph_memory_slots] + slot_io_memory_alias_groups = [(slot[1], slot[2]) for slot in _graph_memory_slots] + slot_io_liveness_groups = [(slot[3], slot[4]) for slot in _graph_memory_slots] + warmup_plan_alias_groups = [slot[5] for slot in _graph_memory_slots] # Check reuse graph conditions and reorganize sample_args and sample_kwargs. # Note: When capturing a graph, we hold onto the args and kwargs so we have static buffers @@ -545,6 +594,7 @@ def _make_graphed_callables( allocator_settings_to_apply = None allocator_settings_to_restore = None + allocator_settings_setter = None if use_slot_memory: allocator_conf = os.getenv("PYTORCH_CUDA_ALLOC_CONF") or os.getenv("PYTORCH_ALLOC_CONF", "") allocator_parts = [part.strip() for part in allocator_conf.split(",") if part.strip()] @@ -896,55 +946,51 @@ def hook_fn( if pre_warmup_hook is not None: pre_warmup_hook() for warmup_iter in range(num_warmup_iters): - hooks = [] - for module in func.modules(): - hook = module.register_forward_hook(hook_fn) - hooks.append(hook) - - if use_slot_memory: - observed_saved_tensor_plan = [] - observed_saved_tensors = [] - observed_saved_versions = [] - copied_storages = {} - record_saved_tensor = make_saved_tensor_recorder( - func_idx, - observed_saved_tensors, - observed_saved_versions, - copied_storages, - observed_saved_tensor_plan, - ) + with _module_forward_hooks(func.modules(), hook_fn): + if use_slot_memory: + observed_saved_tensor_plan = [] + observed_saved_tensors = [] + observed_saved_versions = [] + copied_storages = {} + record_saved_tensor = make_saved_tensor_recorder( + func_idx, + observed_saved_tensors, + observed_saved_versions, + copied_storages, + observed_saved_tensor_plan, + ) - with torch.autograd.graph.saved_tensors_hooks(record_saved_tensor, lambda x: x): + with torch.autograd.graph.saved_tensors_hooks( + record_saved_tensor, lambda x: x + ): + outputs, _ = _tree_flatten(func(*args, **kwargs)) + observed_boundary_aliases = observe_saved_tensor_boundary_aliases( + func_idx, + observed_saved_tensors, + observed_saved_versions, + outputs, + observed_saved_tensor_plan, + ) + update_warmup_plan( + per_callable_saved_tensor_plans, + func_idx, + observed_saved_tensor_plan, + "Forward", + ) + update_warmup_plan( + per_callable_saved_tensor_boundary_aliases, + func_idx, + observed_boundary_aliases, + "Forward boundary alias", + ) + update_warmup_plan( + per_callable_output_tensor_plans, + func_idx, + [_io_tensor_plan(output, "output") for output in outputs], + "Output", + ) + else: outputs, _ = _tree_flatten(func(*args, **kwargs)) - observed_boundary_aliases = observe_saved_tensor_boundary_aliases( - func_idx, - observed_saved_tensors, - observed_saved_versions, - outputs, - observed_saved_tensor_plan, - ) - update_warmup_plan( - per_callable_saved_tensor_plans, - func_idx, - observed_saved_tensor_plan, - "Forward", - ) - update_warmup_plan( - per_callable_saved_tensor_boundary_aliases, - func_idx, - observed_boundary_aliases, - "Forward boundary alias", - ) - update_warmup_plan( - per_callable_output_tensor_plans, - func_idx, - [_io_tensor_plan(output, "output") for output in outputs], - "Output", - ) - else: - outputs, _ = _tree_flatten(func(*args, **kwargs)) - for hook in hooks: - hook.remove() if is_training: inputs = tuple(i for i in static_input_surface if i.requires_grad) with _none_grad_context_wrapper(inputs): @@ -1072,7 +1118,13 @@ def hook_fn( ] if allocator_settings_to_apply is not None: - torch._C._accelerator_setAllocatorSettings(allocator_settings_to_apply) + if _allocator_settings_guard is None or allocator_settings_setter is None: + raise RuntimeError("CUDA graph slot capture is missing its allocator settings guard.") + _allocator_settings_guard.apply( + allocator_settings_setter, + allocator_settings_to_apply, + allocator_settings_to_restore, + ) if use_slot_memory: if isinstance(sample_args, tuple): @@ -1082,7 +1134,7 @@ def hook_fn( staging_groups = [] for func_idx, args in enumerate(sample_args): old_input = args[0] - saved_arena_id, _ = saved_tensor_memory_alias_groups[func_idx] + saved_arena_id = saved_tensor_arena_ids[func_idx] staging_key = (saved_arena_id, _input_staging_key(old_input)) group_idx = staging_group_by_key.get(staging_key) if group_idx is None: @@ -1091,7 +1143,7 @@ def hook_fn( staging_groups.append({"members": [], "candidates": {}}) group = staging_groups[group_idx] group["members"].append(func_idx) - group["candidates"].setdefault(old_input.untyped_storage()._cdata, old_input) + group["candidates"].setdefault(_tensor_storage_ptr(old_input), old_input) # MCore's sample-input plan and the union liveness coloring are each safe in # isolation, but reusing an arbitrary representative can transitively merge two @@ -1124,11 +1176,12 @@ def match_staging_group(group_idx, seen_storages): source = next(iter(group["candidates"].values())) signature = _saved_tensor_signature(source) storage_numel = signature[-1] // source.element_size() - backing = torch.empty( - (source.storage_offset() + storage_numel,), - dtype=source.dtype, - device=source.device, - ) + with torch.cuda.use_mem_pool(slot_allocator_pool): + backing = torch.empty( + (source.storage_offset() + storage_numel,), + dtype=source.dtype, + device=source.device, + ) input_target = torch.empty((0,), dtype=source.dtype, device=source.device).set_( backing.untyped_storage(), source.storage_offset(), @@ -1644,7 +1697,7 @@ def slot_tensor_targets(plan, arena=None): measure_spill=True, preassigned_targets=preassigned_targets, ) - arena_id, _ = saved_tensor_memory_alias_groups[func_idx] + arena_id = saved_tensor_arena_ids[func_idx] arena_sizes[arena_id] = max(arena_sizes.get(arena_id, 0), output_bytes + spill_bytes) del temporary_outputs, temporary_arena, preassigned_targets @@ -1659,7 +1712,7 @@ def slot_tensor_targets(plan, arena=None): per_callable_slot_saved_targets = [] for func_idx, plan in enumerate(per_callable_saved_tensor_plans): - arena_id, _ = saved_tensor_memory_alias_groups[func_idx] + arena_id = saved_tensor_arena_ids[func_idx] arena = slot_saved_arenas.get(arena_id) output_targets, output_bytes = slot_tensor_targets( per_callable_output_tensor_plans[func_idx], arena @@ -1844,12 +1897,12 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): slot = _graph_memory_slots[first_func_idx] event_key = ( c_id > 0, - slot[4], + slot[3], logical_idx, - slot[2], + slot[1], _num_layers_per_chunk[m_chunk], ) - group_records.append((event_key, slot[3])) + group_records.append((event_key, slot[2])) if c_id > 0: group_fwd_idx[m_chunk] += 1 else: @@ -2055,6 +2108,8 @@ def release_checkpoint_live_storages(expected_blocks, owners, phase): """Make checkpoint-live blocks free before allocator topology restoration.""" if set(owners) != set(expected_blocks): raise RuntimeError("CUDA graph slot checkpoint live-storage owner set changed.") + # Validate every owner before mutating any StorageImpl. A validation failure must not + # leave a partially detached checkpoint boundary. for block_ptr, storage in owners.items(): storage_ptr = storage.data_ptr() if not block_ptr <= storage_ptr < block_ptr + expected_blocks[block_ptr]: @@ -2065,6 +2120,9 @@ def release_checkpoint_live_storages(expected_blocks, owners, phase): raise RuntimeError( "CUDA graph slot checkpoint live storage lost its allocator deleter." ) + + for storage in owners.values(): + storage_ptr = storage.data_ptr() torch._C._free_And_Remove_DeleterFn(storage._cdata) tex._graph_checkpoint_detach_storage(storage._cdata) if storage.data_ptr() != storage_ptr or torch._C._has_Standard_Deleter(storage._cdata): @@ -2558,7 +2616,7 @@ def restore_slot_pool_boundary( raise RuntimeError("CUDA graph capture ended inside a slot checkpoint group.") if allocator_settings_to_restore is not None: - torch._C._accelerator_setAllocatorSettings(allocator_settings_to_restore) + _allocator_settings_guard.restore() if use_slot_memory and ( any(native_io_anchors.values()) or any(native_io_capture_counts.values()) @@ -2950,13 +3008,13 @@ def make_graphed_callables( graphs. Only supported with Mcore interleaved pipeline parallelism, i.e. when `_order` is provided. All callables in `modules` are assumed to have inputs and outputs with the same dtype and shape. - _graph_memory_slots: sequence of 7- or 9-int tuples, default = None + _graph_memory_slots: sequence of 6-int tuples, default = None Private liveness plan for mutually exclusive graph variants. Each tuple describes - saved-tensor, graph-I/O, and warmup alias groups for one graph input. Nine-field plans - additionally provide a frame ID and conflict mask for cross-slot validation. Requires the - first positional sample argument of every graph input to be a plain CUDA tensor; it - is snapshotted into the slot arenas whenever forward saves it for backward, so - shape-identical graph inputs can share one input staging surface. + the saved-tensor arena, physical I/O slot, I/O branch, model chunk, layer, and warmup + alias group for one graph input. Requires the first positional sample argument of every + graph input to be a plain CUDA tensor; it is snapshotted into the slot arenas whenever + forward saves it for backward, so shape-identical graph inputs can share one input + staging surface. pre_warmup_hook: callable, default = None A hook function that will be called before the warmup iterations. post_warmup_hook: callable, default = None @@ -3066,8 +3124,6 @@ def make_graphed_callables( if cache_quantized_params is None: cache_quantized_params = False - set_capture_start() - # Handle single module. just_one_callable = False if not isinstance(modules, tuple): @@ -3091,8 +3147,9 @@ def make_graphed_callables( recipe = None module_uses_fp8 = dict(zip((id(m) for m in modules), enabled)) - # Store FP8 tensors to reset later. - saved_fp8_tensors = save_fp8_tensors(modules, recipe=recipe) + for module in modules: + if not isinstance(module, torch.nn.Module): + raise TypeError(f"Graphing for {type(module)} is not supported.") # FP8 wrapper. old_call_funcs = {} @@ -3118,58 +3175,71 @@ def call_func(self, *args, **kwargs): block_cls.__call__ = call_func - forward_funcs = [] - for module in modules: - if not isinstance(module, torch.nn.Module): - raise TypeError(f"Graphing for {type(module)} is not supported.") - wrap_autocast(module) - forward_funcs.append(module) - - if just_one_callable: - forward_funcs = forward_funcs[0] - else: - forward_funcs = tuple(forward_funcs) - - # Save RNG state. - if graph_safe_rng_available(): - generators = [ - torch.cuda.default_generators[torch.cuda.current_device()], - *get_all_rng_states().values(), - ] - original_rng_states = [state.get_state() for state in generators] - else: - original_rng_states = torch.cuda.get_rng_state() - - graphed_callables = _make_graphed_callables( - forward_funcs, - sample_args, - num_warmup_iters=num_warmup_iters, - allow_unused_input=allow_unused_input, - cache_quantized_params=cache_quantized_params, - sample_kwargs=sample_kwargs, - _order=_order, - _num_layers_per_chunk=_num_layers_per_chunk, - pool=pool, - retain_graph_in_backward=retain_graph_in_backward, - _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers, - _graph_memory_slots=_graph_memory_slots, - pre_warmup_hook=pre_warmup_hook, - post_warmup_hook=post_warmup_hook, - ) - - # Ensures warmup does not affect numerics for ops such as dropout. - if graph_safe_rng_available(): - for gen, state in zip(generators, original_rng_states): - gen.set_state(state) - else: - torch.cuda.set_rng_state(original_rng_states) + allocator_settings_guard = _AllocatorSettingsGuard() + saved_fp8_tensors = None + fp8_state_saved = False + rng_restore_callbacks = [] + capture_started = False + try: + # Store all process-wide state before capture and register enough information to restore + # anything that was already changed if a later preparation step raises. + saved_fp8_tensors = save_fp8_tensors(modules, recipe=recipe) + fp8_state_saved = True + + forward_funcs = [] + for module in modules: + wrap_autocast(module) + forward_funcs.append(module) + + if just_one_callable: + forward_funcs = forward_funcs[0] + else: + forward_funcs = tuple(forward_funcs) - # Remove FP8 wrapper. - for module_cls, old_call in old_call_funcs.items(): - module_cls.__call__ = old_call + if graph_safe_rng_available(): + generators = [ + torch.cuda.default_generators[torch.cuda.current_device()], + *get_all_rng_states().values(), + ] + original_rng_states = [state.get_state() for state in generators] + rng_restore_callbacks = [ + (generator.set_state, state) + for generator, state in zip(generators, original_rng_states) + ] + else: + original_rng_state = torch.cuda.get_rng_state() + rng_restore_callbacks = [(torch.cuda.set_rng_state, original_rng_state)] - # Restore FP8 state. - restore_fp8_tensors(modules, saved_fp8_tensors) + set_capture_start() + capture_started = True + graphed_callables = _make_graphed_callables( + forward_funcs, + sample_args, + num_warmup_iters=num_warmup_iters, + allow_unused_input=allow_unused_input, + cache_quantized_params=cache_quantized_params, + sample_kwargs=sample_kwargs, + _order=_order, + _num_layers_per_chunk=_num_layers_per_chunk, + pool=pool, + retain_graph_in_backward=retain_graph_in_backward, + _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers, + _graph_memory_slots=_graph_memory_slots, + _allocator_settings_guard=allocator_settings_guard, + pre_warmup_hook=pre_warmup_hook, + post_warmup_hook=post_warmup_hook, + ) + finally: + # ExitStack runs every callback even if an earlier restoration fails. + with contextlib.ExitStack() as capture_cleanup: + if capture_started: + capture_cleanup.callback(set_capture_end) + if fp8_state_saved: + capture_cleanup.callback(restore_fp8_tensors, modules, saved_fp8_tensors) + for module_cls, old_call in old_call_funcs.items(): + capture_cleanup.callback(setattr, module_cls, "__call__", old_call) + for restore_rng_state, state in rng_restore_callbacks: + capture_cleanup.callback(restore_rng_state, state) + capture_cleanup.callback(allocator_settings_guard.restore) - set_capture_end() return graphed_callables From 8fa96bb365bc9aaf33c9053de712e520c8262c92 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Fri, 21 Aug 2026 17:23:07 +0900 Subject: [PATCH 07/16] Fix allocator guard lint Signed-off-by: Tailai Ma --- transformer_engine/pytorch/graph.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 56c9e256f9..7eef61b238 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -50,6 +50,7 @@ def __init__(self) -> None: self._settings_to_restore = None def apply(self, setter: Callable[[str], None], settings: str, restore: str) -> None: + """Apply temporary allocator settings and remember how to restore them.""" if self._setter is not None: raise RuntimeError("CUDA allocator settings guard is already active.") self._setter = setter @@ -57,6 +58,7 @@ def apply(self, setter: Callable[[str], None], settings: str, restore: str) -> N setter(settings) def restore(self) -> None: + """Restore the allocator settings saved by :meth:`apply`.""" if self._setter is None: return setter = self._setter From 209e757d3085616230008db12ecb522acb703588 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Fri, 21 Aug 2026 20:24:30 +0900 Subject: [PATCH 08/16] Honor user gradient liveness in graph slots Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 16 +++++++++------- transformer_engine/pytorch/graph.py | 28 +++++++++++++++------------- 2 files changed, 24 insertions(+), 20 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 593c271efc..429b86b466 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -777,9 +777,11 @@ def test_make_graphed_callables_with_interleaved_pipeline_parallelism( assert_all_equal(outputs, graph_outputs) -def _slot(saved_arena, branch, io_arena, overlap=0, frame=0, warmup=0): +def _slot(saved_arena, branch, io_arena, overlap=0, frame=0, warmup=0, user_grad=None): """Build one private graph-memory slot used by the focused tests below.""" - return (saved_arena, io_arena, branch, overlap, frame, warmup) + if user_grad is None: + user_grad = io_arena + return (saved_arena, io_arena, branch, overlap, frame, warmup, user_grad) def test_graph_capture_contexts_restore_process_state_on_error(monkeypatch) -> None: @@ -1513,8 +1515,8 @@ def forward(self, inp): reset_graphs(graphed) -def test_slot_memory_reuses_user_grad_surface_across_chunks() -> None: - """Adjacent backward chunks may reuse one physical-slot gradient surface.""" +def test_slot_memory_honors_user_grad_liveness_groups() -> None: + """The private plan may keep adjacent asynchronous gradient consumers disjoint.""" class Module(torch.nn.Module): def forward(self, inp): @@ -1530,8 +1532,8 @@ def forward(self, inp): _num_layers_per_chunk=[1, 1], _reuse_graph_input_output_buffers=True, _graph_memory_slots=( - _slot(0, 0, 0, overlap=0, warmup=0), - _slot(1, 1, 0, overlap=1, warmup=1), + _slot(0, 0, 0, overlap=0, warmup=0, user_grad=0), + _slot(1, 1, 0, overlap=1, warmup=1, user_grad=1), ), ) @@ -1539,7 +1541,7 @@ def forward(self, inp): inp = torch.randn(4096, device="cuda", requires_grad=True) graphed[1](graphed[0](inp)).sum().backward() torch.testing.assert_close(inp.grad, 4.0 * inp.detach().pow(3)) - assert tuple(graphed[-1]._te_cuda_graph_user_grad_arenas) == (0,) + assert tuple(graphed[-1]._te_cuda_graph_user_grad_arenas) == (0, 1) finally: reset_graphs(graphed) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 7eef61b238..bc0f7c42b2 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -416,6 +416,7 @@ def _make_graphed_callables( slot_io_memory_alias_groups = None slot_io_liveness_groups = None warmup_plan_alias_groups = None + user_grad_arena_ids = None if use_slot_memory: if _order is None or not is_training or not _reuse_graph_input_output_buffers: raise RuntimeError( @@ -431,15 +432,16 @@ def _make_graphed_callables( ) if any( not isinstance(slot, tuple) - or len(slot) != 6 + or len(slot) != 7 or not all(isinstance(value, int) for value in slot) for slot in _graph_memory_slots ): - raise TypeError("Each graph-memory slot must be a tuple of six integers.") + raise TypeError("Each graph-memory slot must be a tuple of seven integers.") saved_tensor_arena_ids = [slot[0] for slot in _graph_memory_slots] slot_io_memory_alias_groups = [(slot[1], slot[2]) for slot in _graph_memory_slots] slot_io_liveness_groups = [(slot[3], slot[4]) for slot in _graph_memory_slots] warmup_plan_alias_groups = [slot[5] for slot in _graph_memory_slots] + user_grad_arena_ids = [slot[6] for slot in _graph_memory_slots] # Check reuse graph conditions and reorganize sample_args and sample_kwargs. # Note: When capturing a graph, we hold onto the args and kwargs so we have static buffers @@ -1739,21 +1741,21 @@ def slot_tensor_targets(plan, arena=None): slot_sizes = {} for func_idx, plan in enumerate(per_callable_user_grad_tensor_plans): _, required_bytes = slot_tensor_targets(plan) - physical_slot, _ = slot_io_memory_alias_groups[func_idx] - slot_sizes[physical_slot] = max(slot_sizes.get(physical_slot, 0), required_bytes) + arena_id = user_grad_arena_ids[func_idx] + slot_sizes[arena_id] = max(slot_sizes.get(arena_id, 0), required_bytes) with torch.cuda.use_mem_pool(slot_allocator_pool): slot_user_grad_arenas = { - physical_slot: torch.empty( + arena_id: torch.empty( (required_bytes,), dtype=torch.uint8, device=torch.cuda.current_device() ) - for physical_slot, required_bytes in slot_sizes.items() + for arena_id, required_bytes in slot_sizes.items() if required_bytes > 0 } for func_idx, plan in enumerate(per_callable_user_grad_tensor_plans): - physical_slot, _ = slot_io_memory_alias_groups[func_idx] - targets, _ = slot_tensor_targets(plan, slot_user_grad_arenas.get(physical_slot)) + arena_id = user_grad_arena_ids[func_idx] + targets, _ = slot_tensor_targets(plan, slot_user_grad_arenas.get(arena_id)) per_callable_user_grad_tensor_targets[func_idx] = list(targets) @contextlib.contextmanager @@ -3010,13 +3012,13 @@ def make_graphed_callables( graphs. Only supported with Mcore interleaved pipeline parallelism, i.e. when `_order` is provided. All callables in `modules` are assumed to have inputs and outputs with the same dtype and shape. - _graph_memory_slots: sequence of 6-int tuples, default = None + _graph_memory_slots: sequence of 7-int tuples, default = None Private liveness plan for mutually exclusive graph variants. Each tuple describes the saved-tensor arena, physical I/O slot, I/O branch, model chunk, layer, and warmup - alias group for one graph input. Requires the first positional sample argument of every - graph input to be a plain CUDA tensor; it is snapshotted into the slot arenas whenever - forward saves it for backward, so shape-identical graph inputs can share one input - staging surface. + alias group, followed by the returned user-gradient arena for one graph input. Requires + the first positional sample argument of every graph input to be a plain CUDA tensor; it is + snapshotted into the slot arenas whenever forward saves it for backward, so shape-identical + graph inputs can share one input staging surface. pre_warmup_hook: callable, default = None A hook function that will be called before the warmup iterations. post_warmup_hook: callable, default = None From 6110d0a46ac914c6442c6dfae3e13dec03ef14fd Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Fri, 21 Aug 2026 23:41:59 +0900 Subject: [PATCH 09/16] Protect saved CUDA graph kwargs across slot schedules Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 36 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 22 ++++++++++++++---- 2 files changed, 53 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 429b86b466..49e311dd03 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1012,6 +1012,42 @@ def forward(self, inp): reset_graphs(graphed) +def test_slot_memory_snapshots_shared_kwarg_across_alternate_liveness() -> None: + """A later forward must not overwrite a shared kwarg needed by backward.""" + + class Module(torch.nn.Module): + def forward(self, inp, scale): + return inp * scale + + module = Module().cuda() + samples = tuple((torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(2)) + shared_scale = torch.ones(4096, device="cuda") + sample_kwargs = ({"scale": shared_scale}, {"scale": shared_scale}) + graphed = make_graphed_callables( + (module,), + samples, + sample_kwargs=sample_kwargs, + num_warmup_iters=2, + allow_unused_input=True, + _order=[1, -1, 1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0), _slot(1, 1, 1)), + ) + + try: + inp0 = torch.ones(4096, device="cuda", requires_grad=True) + inp1 = torch.ones(4096, device="cuda", requires_grad=True) + out0 = graphed[0](inp0, scale=torch.full_like(inp0, 2.0)) + out1 = graphed[1](inp1, scale=torch.full_like(inp1, 3.0)) + out0.sum().backward() + out1.sum().backward() + torch.testing.assert_close(inp0.grad, torch.full_like(inp0, 2.0)) + torch.testing.assert_close(inp1.grad, torch.full_like(inp1, 3.0)) + finally: + reset_graphs(graphed) + + @pytest.mark.parametrize("reverse_replay", (False, True), ids=("forward", "reverse")) def test_slot_memory_checkpoint_reuses_lockstep_branches(reverse_replay) -> None: """Lockstep CP branches must restore one live slot boundary between captures.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index bc0f7c42b2..2f1acfaf40 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -722,7 +722,13 @@ def _make_graphed_callables( "Slot user-input snapshots require the first positional tensor to be the " "first flattened graph input." ) - per_callable_snapshot_input_storage_ptrs.append(_tensor_storage_ptr(args[0])) + per_callable_snapshot_input_storage_ptrs.append( + { + _tensor_storage_ptr(tensor) + for tensor in flatten_sample_args[func_idx] + if tensor.__class__ is torch.Tensor and tensor.is_cuda + } + ) else: per_callable_saved_tensor_plans = None per_callable_saved_tensor_boundary_aliases = None @@ -881,7 +887,7 @@ def record_saved_tensor(tensor): storage_ptr = _tensor_storage_ptr(tensor) signature = _saved_tensor_signature(tensor) snapshot_user_input = ( - tensor.is_cuda and storage_ptr == per_callable_snapshot_input_storage_ptrs[func_idx] + tensor.is_cuda and storage_ptr in per_callable_snapshot_input_storage_ptrs[func_idx] ) is_external = not tensor.is_cuda or ( storage_ptr in per_callable_external_storage_ptrs[func_idx] @@ -1591,6 +1597,12 @@ def semantic_boundary_alias_targets(func_idx, outputs): (saved_idx, aliases[saved_idx]) for saved_idx in component_saved_indices if aliases[saved_idx] is not None + # Only the leading input is rebound to a union-liveness staging surface. + # Other user inputs may share MCore capture-order buffers that overlap in a + # different runtime schedule, so they must use the saved arena instead. + and not ( + aliases[saved_idx][1] == "input" and aliases[saved_idx][2] != 0 + ) ] if not candidate_aliases: continue @@ -3016,9 +3028,9 @@ def make_graphed_callables( Private liveness plan for mutually exclusive graph variants. Each tuple describes the saved-tensor arena, physical I/O slot, I/O branch, model chunk, layer, and warmup alias group, followed by the returned user-gradient arena for one graph input. Requires - the first positional sample argument of every graph input to be a plain CUDA tensor; it is - snapshotted into the slot arenas whenever forward saves it for backward, so shape-identical - graph inputs can share one input staging surface. + the first positional sample argument of every graph input to be a plain CUDA tensor. Plain + CUDA user inputs are snapshotted into the slot arenas whenever forward saves them for + backward, so shape-identical graph inputs can safely share staging surfaces. pre_warmup_hook: callable, default = None A hook function that will be called before the warmup iterations. post_warmup_hook: callable, default = None From ffef962eefa53e9ba3d0e40fb8e46c3937f99578 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 00:28:22 +0900 Subject: [PATCH 10/16] Roll back failed CUDA graph allocator restores Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 79 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 49 ++++++++++++++---- 2 files changed, 119 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 49e311dd03..cc11d30d70 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1113,6 +1113,85 @@ def forward(self, inp): reset_graphs(graphed) +@pytest.mark.parametrize("failure_timing", ("before", "after")) +def test_slot_memory_checkpoint_rolls_back_restore_failure(monkeypatch, failure_timing) -> None: + """A failed allocator restore must put the original owners back before raising.""" + + class Module(torch.nn.Module): + def forward(self, inp): + transient = torch.cat((inp.square(), inp.sin(), inp.cos()), dim=0) + return transient[: inp.numel()] * 3.0 + + module = Module().cuda() + variants = 2 + samples = tuple( + (torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(variants) + ) + slots = tuple(_slot(0, variant, 1, warmup=variant) for variant in range(variants)) + real_set_state = torch._C._cuda_setCheckpointPoolState + real_detach = te_graph.tex._graph_checkpoint_detach_storage + detached_storage_impls = [] + original_owner_impls = [] + set_state_calls = 0 + + def record_detach(storage_impl_ptr): + detached_storage_impls.append(storage_impl_ptr) + real_detach(storage_impl_ptr) + + def fail_first_set_state(*args, **kwargs): + nonlocal set_state_calls + set_state_calls += 1 + if set_state_calls == 1: + original_owner_impls.extend(detached_storage_impls) + if failure_timing == "before": + raise RuntimeError("injected checkpoint restore failure") + result = real_set_state(*args, **kwargs) + if set_state_calls == 1: + raise RuntimeError("injected checkpoint restore failure") + return result + + monkeypatch.setattr(te_graph.tex, "_graph_checkpoint_detach_storage", record_detach) + monkeypatch.setattr(torch._C, "_cuda_setCheckpointPoolState", fail_first_set_state) + + with pytest.raises(RuntimeError, match="injected checkpoint restore failure") as exc_info: + make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[1, 2, -1, -2], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=slots, + ) + assert set_state_calls == 2 + assert original_owner_impls + assert all( + torch._C._has_Standard_Deleter(storage_impl_ptr) + for storage_impl_ptr in original_owner_impls + ) + + del exc_info + gc.collect() + graphed = make_graphed_callables( + (module,) * variants, + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[1, 2, -1, -2], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=slots, + ) + try: + for graph in graphed: + inp = torch.randn(4096, device="cuda", requires_grad=True) + graph(inp).sum().backward() + torch.testing.assert_close(inp.grad, 6.0 * inp.detach()) + finally: + reset_graphs(graphed) + + def test_slot_memory_checkpoint_reclaims_retained_branch_outputs() -> None: """A module-held source output must not pin a mutually exclusive branch allocation.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 2f1acfaf40..82b5bf1845 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -2178,15 +2178,46 @@ def restore_slot_pool_boundary( phase, ): """Switch between two allocator boundaries while preserving their StorageImpls.""" - release_checkpoint_live_storages(current_blocks, current_owners, phase) - torch._C._cuda_setCheckpointPoolState( - torch.cuda.current_device(), - target_state, - [], - [storage._cdata for storage in target_owners.values()], - ) - verify_checkpoint_live_storages(target_blocks, target_owners) - assert_slot_pool_liveness(target_blocks, phase) + device = torch.cuda.current_device() + rollback_state = torch._C._cuda_getCheckpointState(device, mempool) + current_released = False + try: + release_checkpoint_live_storages(current_blocks, current_owners, phase) + current_released = True + torch._C._cuda_setCheckpointPoolState( + device, + target_state, + [], + [storage._cdata for storage in target_owners.values()], + ) + verify_checkpoint_live_storages(target_blocks, target_owners) + assert_slot_pool_liveness(target_blocks, phase) + except BaseException as restore_error: + if not current_released: + raise + try: + rollback_blocks = slot_pool_active_blocks() + rollback_owners = checkpoint_live_storage_owners( + rollback_blocks, + f"{phase} rollback", + (current_owners, target_owners), + ) + release_checkpoint_live_storages( + rollback_blocks, rollback_owners, f"{phase} rollback" + ) + torch._C._cuda_setCheckpointPoolState( + device, + rollback_state, + [], + [storage._cdata for storage in current_owners.values()], + ) + verify_checkpoint_live_storages(current_blocks, current_owners) + assert_slot_pool_liveness(current_blocks, f"{phase} rollback") + except BaseException as rollback_error: + raise RuntimeError( + f"CUDA graph slot checkpoint rollback failed after {restore_error!r}." + ) from rollback_error + raise if _order is not None: # pylint: disable=too-many-nested-blocks per_callable_static_outputs = [None] * len(flatten_sample_args) From db1f42b244dac840385f48e2ade06875573eb008 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 00:47:11 +0900 Subject: [PATCH 11/16] Roll back partial CUDA graph owner release Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 13 ++++++++++--- transformer_engine/pytorch/graph.py | 18 +++++++++++------- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index cc11d30d70..4294f88ce4 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1113,7 +1113,7 @@ def forward(self, inp): reset_graphs(graphed) -@pytest.mark.parametrize("failure_timing", ("before", "after")) +@pytest.mark.parametrize("failure_timing", ("detach", "before", "after")) def test_slot_memory_checkpoint_rolls_back_restore_failure(monkeypatch, failure_timing) -> None: """A failed allocator restore must put the original owners back before raising.""" @@ -1137,6 +1137,9 @@ def forward(self, inp): def record_detach(storage_impl_ptr): detached_storage_impls.append(storage_impl_ptr) real_detach(storage_impl_ptr) + if failure_timing == "detach" and len(detached_storage_impls) == 1: + original_owner_impls.extend(detached_storage_impls) + raise RuntimeError("injected checkpoint detach failure") def fail_first_set_state(*args, **kwargs): nonlocal set_state_calls @@ -1153,7 +1156,9 @@ def fail_first_set_state(*args, **kwargs): monkeypatch.setattr(te_graph.tex, "_graph_checkpoint_detach_storage", record_detach) monkeypatch.setattr(torch._C, "_cuda_setCheckpointPoolState", fail_first_set_state) - with pytest.raises(RuntimeError, match="injected checkpoint restore failure") as exc_info: + with pytest.raises( + RuntimeError, match="injected checkpoint (detach|restore) failure" + ) as exc_info: make_graphed_callables( (module,) * variants, samples, @@ -1164,7 +1169,9 @@ def fail_first_set_state(*args, **kwargs): _reuse_graph_input_output_buffers=True, _graph_memory_slots=slots, ) - assert set_state_calls == 2 + assert set_state_calls == (1 if failure_timing == "detach" else 2) + if failure_timing == "detach": + original_owner_impls = list(dict.fromkeys(detached_storage_impls)) assert original_owner_impls assert all( torch._C._has_Standard_Deleter(storage_impl_ptr) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 82b5bf1845..0a88792b49 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -2120,12 +2120,10 @@ def visit(value): ) return owners - def release_checkpoint_live_storages(expected_blocks, owners, phase): - """Make checkpoint-live blocks free before allocator topology restoration.""" + def validate_checkpoint_live_storages(expected_blocks, owners): + """Validate checkpoint owners before changing any StorageImpl.""" if set(owners) != set(expected_blocks): raise RuntimeError("CUDA graph slot checkpoint live-storage owner set changed.") - # Validate every owner before mutating any StorageImpl. A validation failure must not - # leave a partially detached checkpoint boundary. for block_ptr, storage in owners.items(): storage_ptr = storage.data_ptr() if not block_ptr <= storage_ptr < block_ptr + expected_blocks[block_ptr]: @@ -2137,6 +2135,8 @@ def release_checkpoint_live_storages(expected_blocks, owners, phase): "CUDA graph slot checkpoint live storage lost its allocator deleter." ) + def release_checkpoint_live_storages(expected_blocks, owners, phase): + """Make prevalidated checkpoint-live blocks free for topology restoration.""" for storage in owners.values(): storage_ptr = storage.data_ptr() torch._C._free_And_Remove_DeleterFn(storage._cdata) @@ -2180,10 +2180,13 @@ def restore_slot_pool_boundary( """Switch between two allocator boundaries while preserving their StorageImpls.""" device = torch.cuda.current_device() rollback_state = torch._C._cuda_getCheckpointState(device, mempool) - current_released = False + current_release_started = False try: + # Complete validation before the transaction starts. Once the first owner can be + # mutated, every failure must restore the original allocator boundary. + validate_checkpoint_live_storages(current_blocks, current_owners) + current_release_started = True release_checkpoint_live_storages(current_blocks, current_owners, phase) - current_released = True torch._C._cuda_setCheckpointPoolState( device, target_state, @@ -2193,7 +2196,7 @@ def restore_slot_pool_boundary( verify_checkpoint_live_storages(target_blocks, target_owners) assert_slot_pool_liveness(target_blocks, phase) except BaseException as restore_error: - if not current_released: + if not current_release_started: raise try: rollback_blocks = slot_pool_active_blocks() @@ -2202,6 +2205,7 @@ def restore_slot_pool_boundary( f"{phase} rollback", (current_owners, target_owners), ) + validate_checkpoint_live_storages(rollback_blocks, rollback_owners) release_checkpoint_live_storages( rollback_blocks, rollback_owners, f"{phase} rollback" ) From 461d00fc27e849da2d347e10fbc9e9f90dab6b26 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 21 Aug 2026 16:03:48 +0000 Subject: [PATCH 12/16] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_cuda_graphs.py | 4 +--- transformer_engine/pytorch/graph.py | 4 +--- 2 files changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 4294f88ce4..beb982e1e4 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1124,9 +1124,7 @@ def forward(self, inp): module = Module().cuda() variants = 2 - samples = tuple( - (torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(variants) - ) + samples = tuple((torch.ones(4096, device="cuda", requires_grad=True),) for _ in range(variants)) slots = tuple(_slot(0, variant, 1, warmup=variant) for variant in range(variants)) real_set_state = torch._C._cuda_setCheckpointPoolState real_detach = te_graph.tex._graph_checkpoint_detach_storage diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 0a88792b49..7184fce1f9 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -1600,9 +1600,7 @@ def semantic_boundary_alias_targets(func_idx, outputs): # Only the leading input is rebound to a union-liveness staging surface. # Other user inputs may share MCore capture-order buffers that overlap in a # different runtime schedule, so they must use the saved arena instead. - and not ( - aliases[saved_idx][1] == "input" and aliases[saved_idx][2] != 0 - ) + and not (aliases[saved_idx][1] == "input" and aliases[saved_idx][2] != 0) ] if not candidate_aliases: continue From 032b62f338edcafcb957aac2c48ec3cd8c089149 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 01:26:52 +0900 Subject: [PATCH 13/16] Restore CUDA graph warmup hooks on failure Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 11 ++++++++++- transformer_engine/pytorch/graph.py | 24 ++++++++++++++++++++++-- 2 files changed, 32 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index beb982e1e4..5ab2ba726f 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -848,6 +848,7 @@ def forward(self, inp): restored_fp8 = [] restored_rng = [] allocator_settings = [] + warmup_hooks = [] monkeypatch.setattr(te_graph, "save_fp8_tensors", lambda *args, **kwargs: fp8_state) monkeypatch.setattr( @@ -861,6 +862,8 @@ def forward(self, inp): def fail_capture(*args, **kwargs): assert te_graph.is_graph_capturing() + kwargs["pre_warmup_hook"]() + assert warmup_hooks == ["pre"] kwargs["_allocator_settings_guard"].apply( allocator_settings.append, "expandable_segments:False", @@ -872,13 +875,19 @@ def fail_capture(*args, **kwargs): assert not te_graph.is_graph_capturing() with pytest.raises(RuntimeError, match="capture failed"): - te_graph.make_graphed_callables(module, (torch.ones(1),)) + te_graph.make_graphed_callables( + module, + (torch.ones(1),), + pre_warmup_hook=lambda: warmup_hooks.append("pre"), + post_warmup_hook=lambda: warmup_hooks.append("post"), + ) assert not te_graph.is_graph_capturing() assert TestModule.__call__ is original_call assert restored_fp8 == [((module,), fp8_state)] assert restored_rng == [rng_state] assert allocator_settings == ["expandable_segments:False", "expandable_segments:True"] + assert warmup_hooks == ["pre", "post"] def test_make_graphed_callables_restores_wrappers_on_preparation_error(monkeypatch) -> None: diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 7184fce1f9..99a64b2c02 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -3224,6 +3224,24 @@ def call_func(self, *args, **kwargs): block_cls.__call__ = call_func + warmup_cleanup_pending = False + guarded_pre_warmup_hook = pre_warmup_hook + guarded_post_warmup_hook = post_warmup_hook + if post_warmup_hook is not None: + + def guarded_pre_warmup_hook(): + nonlocal warmup_cleanup_pending + if pre_warmup_hook is not None: + pre_warmup_hook() + warmup_cleanup_pending = True + + def guarded_post_warmup_hook(): + nonlocal warmup_cleanup_pending + if not warmup_cleanup_pending: + return + warmup_cleanup_pending = False + post_warmup_hook() + allocator_settings_guard = _AllocatorSettingsGuard() saved_fp8_tensors = None fp8_state_saved = False @@ -3275,8 +3293,8 @@ def call_func(self, *args, **kwargs): _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers, _graph_memory_slots=_graph_memory_slots, _allocator_settings_guard=allocator_settings_guard, - pre_warmup_hook=pre_warmup_hook, - post_warmup_hook=post_warmup_hook, + pre_warmup_hook=guarded_pre_warmup_hook, + post_warmup_hook=guarded_post_warmup_hook, ) finally: # ExitStack runs every callback even if an earlier restoration fails. @@ -3290,5 +3308,7 @@ def call_func(self, *args, **kwargs): for restore_rng_state, state in rng_restore_callbacks: capture_cleanup.callback(restore_rng_state, state) capture_cleanup.callback(allocator_settings_guard.restore) + if guarded_post_warmup_hook is not None: + capture_cleanup.callback(guarded_post_warmup_hook) return graphed_callables From de59f9642e44a26924d0142fb04e9fcbfc090e4b Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 03:43:05 +0900 Subject: [PATCH 14/16] Preserve aliased outputs in CUDA graph slots Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 109 ++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 114 ++++++++++++++++++++++++---- 2 files changed, 208 insertions(+), 15 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 5ab2ba726f..16b66f8e5a 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -986,6 +986,115 @@ def forward(self, inp): reset_graphs(graphed) +def test_slot_memory_preserves_aliased_public_outputs() -> None: + """Slot output arenas retain overlapping public views across CP branches.""" + + class Module(torch.nn.Module): + def forward(self, inp): + output = inp.square() + return output, output[4:] + + module = Module().cuda() + samples = tuple((torch.ones(16, device="cuda", requires_grad=True),) for _ in range(2)) + graphed = make_graphed_callables( + (module, module), + samples, + num_warmup_iters=2, + _order=[1, 2, -1, -2], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0), _slot(0, 1, 0, warmup=1)), + ) + try: + for graph in graphed: + inp = torch.randn(16, device="cuda", requires_grad=True) + output, output_view = graph(inp) + assert output_view.data_ptr() - output.data_ptr() == 4 * output.element_size() + assert output_view.data_ptr() < ( + output.data_ptr() + output.numel() * output.element_size() + ) + (output.sum() + output_view.sum()).backward() + expected_grad = 2.0 * inp.detach() + expected_grad[4:] *= 2.0 + torch.testing.assert_close(inp.grad, expected_grad) + previous = output[4:].clone() + with torch.no_grad(): + output_view.add_(1.0) + torch.testing.assert_close(output[4:], previous + 1.0) + finally: + reset_graphs(graphed) + + +@pytest.mark.parametrize("state_kind", ("parameter", "buffer")) +def test_slot_memory_preserves_public_outputs_aliased_to_module_state(state_kind) -> None: + """Persistent module-state views remain outside the slot allocator pool.""" + + class Module(torch.nn.Module): + def __init__(self): + super().__init__() + state = torch.randn(16, device="cuda") + if state_kind == "parameter": + self.state = torch.nn.Parameter(state) + else: + self.register_buffer("state", state) + + def forward(self, inp): + return inp.square(), self.state.view_as(self.state) + + module = Module() + samples = tuple((torch.ones(16, device="cuda", requires_grad=True),) for _ in range(2)) + graphed = make_graphed_callables( + (module, module), + samples, + num_warmup_iters=2, + allow_unused_input=True, + _order=[1, 2, -1, -2], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0), _slot(0, 1, 0, warmup=1)), + ) + try: + for graph in graphed: + inp = torch.randn(16, device="cuda", requires_grad=True) + output, state_output = graph(inp) + assert ( + state_output.untyped_storage().data_ptr() + == module.state.untyped_storage().data_ptr() + ) + (output.sum() + state_output.sum()).backward() + torch.testing.assert_close(inp.grad, 2.0 * inp.detach()) + if state_kind == "parameter": + torch.testing.assert_close(module.state.grad, torch.ones_like(module.state)) + module.state.grad = None + finally: + reset_graphs(graphed) + + +def test_slot_memory_rejects_public_cuda_tensor_subclass_outputs() -> None: + """Unsupported CUDA outputs must fail before allocator checkpoint mutation.""" + + class OutputTensor(torch.Tensor): + pass + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square().as_subclass(OutputTensor) + + module = Module().cuda() + sample = (torch.ones(16, device="cuda", requires_grad=True),) + with pytest.raises(RuntimeError, match="tensor subclasses"): + make_graphed_callables( + module, + sample, + num_warmup_iters=2, + _order=[1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0),), + ) + assert not te_graph.is_graph_capturing() + + def test_slot_memory_input_staging_respects_overlapping_liveness() -> None: """Live microbatches must not overwrite inputs still needed by backward.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 99a64b2c02..b46e0fc0a7 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -130,6 +130,41 @@ def _io_tensor_plan(tensor: Any, kind: str) -> Optional[Tuple[Any, ...]]: return (kind, None, *_saved_tensor_signature(tensor)) +def _slot_output_plan(outputs: Sequence[Any], func: Callable) -> List[Optional[Tuple[Any, ...]]]: + """Describe public outputs while retaining storage aliases and module-state views.""" + module_state_storages = set() + if isinstance(func, torch.nn.Module): + for tensors in (func.parameters(), func.buffers()): + for tensor in tensors: + if tensor.is_cuda and tensor.untyped_storage().nbytes() > 0: + module_state_storages.add(_tensor_storage_ptr(tensor)) + + storage_groups = {} + plan = [] + for output_idx, output in enumerate(outputs): + spec = _io_tensor_plan(output, "output") + if spec is None and isinstance(output, torch.Tensor) and output.is_cuda: + raise RuntimeError( + "CUDA graph slot memory does not support public CUDA outputs that are " + "tensor subclasses or use a non-strided layout: " + f"output tensor {output_idx} has type {type(output).__name__} " + f"and layout {output.layout}." + ) + if spec is None: + plan.append(None) + continue + + storage_id = _tensor_storage_ptr(output) + storage_offset_bytes = output.storage_offset() * output.element_size() + if storage_id in module_state_storages: + plan.append(("external_output", storage_id, *spec[2:], None, storage_offset_bytes)) + continue + + storage_group = storage_groups.setdefault(storage_id, len(storage_groups)) + plan.append((*spec, storage_group, storage_offset_bytes)) + return plan + + def _arena_view(arena: torch.Tensor, offset: int, spec: Tuple[Any, ...]) -> torch.Tensor: """Materialize a typed tensor view at a byte offset in an arena.""" target = torch.empty((0,), dtype=spec[4], device=spec[5]) @@ -974,6 +1009,7 @@ def hook_fn( record_saved_tensor, lambda x: x ): outputs, _ = _tree_flatten(func(*args, **kwargs)) + observed_output_plan = _slot_output_plan(outputs, func) observed_boundary_aliases = observe_saved_tensor_boundary_aliases( func_idx, observed_saved_tensors, @@ -996,7 +1032,7 @@ def hook_fn( update_warmup_plan( per_callable_output_tensor_plans, func_idx, - [_io_tensor_plan(output, "output") for output in outputs], + observed_output_plan, "Output", ) else: @@ -1249,6 +1285,14 @@ def prepare_native_io_targets(per_callable_plans, kind): f"CUDA graph {kind} family {family} has an incompatible tensor " f"at position {tensor_idx}." ) + modes = {spec[0] for spec in specs} + if modes == {"external_output"}: + continue + if modes != {kind}: + raise RuntimeError( + f"CUDA graph {kind} family {family} has incompatible storage modes " + f"at position {tensor_idx}: {sorted(modes)}." + ) layout_keys = {(spec[0], spec[4], spec[5], spec[6]) for spec in specs} if layout_keys != {(kind, specs[0][4], specs[0][5], specs[0][6])}: raise RuntimeError( @@ -1325,7 +1369,7 @@ def clear_native_io_target_rows(func_indices, clear_outputs=False, clear_grads=F per_callable_user_grad_tensor_targets[func_idx] ) - def copy_outputs_to_slot_arena(func_idx, flatten_outputs): + def copy_outputs_to_slot_arena(func_idx, flatten_outputs, func): """Copy public forward outputs to the fixed surface for their physical slot.""" if per_callable_output_tensor_targets is None: return flatten_outputs @@ -1335,12 +1379,16 @@ def copy_outputs_to_slot_arena(func_idx, flatten_outputs): raise RuntimeError( f"CUDA graph input {func_idx} changed its output count during capture." ) + if _slot_output_plan(flatten_outputs, func) != plan: + raise RuntimeError( + f"CUDA graph input {func_idx} changed its output tensor storage plan " + "during capture." + ) copied_outputs = [] for tensor_idx, (output, spec, target) in enumerate(zip(flatten_outputs, plan, targets)): - if spec != _io_tensor_plan(output, "output"): - raise RuntimeError( - f"CUDA graph input {func_idx} changed its output tensor surface during capture." - ) + if spec is not None and spec[0] == "external_output": + copied_outputs.append(output) + continue if target is None: if spec is None: copied_outputs.append(output) @@ -1676,6 +1724,38 @@ def semantic_boundary_alias_targets(func_idx, outputs): def slot_tensor_targets(plan, arena=None): """Lay out graph-boundary tensors contiguously in an arena.""" + if any(spec is not None and len(spec) > 8 for spec in plan): + records_by_storage_group = {} + for tensor_idx, spec in enumerate(plan): + if spec is None or spec[0] == "external_output": + continue + storage_group = spec[8] + storage_offset_bytes = spec[9] + records_by_storage_group.setdefault(storage_group, []).append( + (storage_offset_bytes, storage_offset_bytes + spec[7], tensor_idx) + ) + + placements = {} + offset = 0 + for storage_group, records in records_by_storage_group.items(): + alignment = max(plan[tensor_idx][4].itemsize for _, _, tensor_idx in records) + component_start = min(start for start, _, _ in records) + component_end = max(end for _, end, _ in records) + component_origin = component_start // alignment * alignment + offset = _align_up(offset) + placements[storage_group] = (offset, component_origin) + offset += component_end - component_origin + + targets = [] + for spec in plan: + if spec is None or spec[0] == "external_output" or arena is None: + targets.append(None) + continue + group_offset, component_origin = placements[spec[8]] + target_offset = group_offset + spec[9] - component_origin + targets.append(_arena_view(arena, target_offset, spec)) + return tuple(targets), _align_up(offset) + targets = [] offset = 0 for spec in plan: @@ -2133,7 +2213,7 @@ def validate_checkpoint_live_storages(expected_blocks, owners): "CUDA graph slot checkpoint live storage lost its allocator deleter." ) - def release_checkpoint_live_storages(expected_blocks, owners, phase): + def release_checkpoint_live_storages(owners, phase): """Make prevalidated checkpoint-live blocks free for topology restoration.""" for storage in owners.values(): storage_ptr = storage.data_ptr() @@ -2184,7 +2264,7 @@ def restore_slot_pool_boundary( # mutated, every failure must restore the original allocator boundary. validate_checkpoint_live_storages(current_blocks, current_owners) current_release_started = True - release_checkpoint_live_storages(current_blocks, current_owners, phase) + release_checkpoint_live_storages(current_owners, phase) torch._C._cuda_setCheckpointPoolState( device, target_state, @@ -2204,9 +2284,7 @@ def restore_slot_pool_boundary( (current_owners, target_owners), ) validate_checkpoint_live_storages(rollback_blocks, rollback_owners) - release_checkpoint_live_storages( - rollback_blocks, rollback_owners, f"{phase} rollback" - ) + release_checkpoint_live_storages(rollback_owners, f"{phase} rollback") torch._C._cuda_setCheckpointPoolState( device, rollback_state, @@ -2302,7 +2380,7 @@ def restore_slot_pool_boundary( flatten_outputs, spec = _tree_flatten(outputs) original_flatten_outputs = flatten_outputs flatten_outputs = copy_outputs_to_slot_arena( - per_callable_fwd_idx, flatten_outputs + per_callable_fwd_idx, flatten_outputs, func ) if branch_group is not None: record_replaced_io_storages( @@ -3063,7 +3141,10 @@ def make_graphed_callables( alias group, followed by the returned user-gradient arena for one graph input. Requires the first positional sample argument of every graph input to be a plain CUDA tensor. Plain CUDA user inputs are snapshotted into the slot arenas whenever forward saves them for - backward, so shape-identical graph inputs can safely share staging surfaces. + backward, so shape-identical graph inputs can safely share staging surfaces. Public CUDA + outputs must be plain strided tensors. Output views that share storage retain their relative + byte offsets in the slot arena, while views of module parameters or buffers remain external + to the graph pool. pre_warmup_hook: callable, default = None A hook function that will be called before the warmup iterations. post_warmup_hook: callable, default = None @@ -3229,19 +3310,22 @@ def call_func(self, *args, **kwargs): guarded_post_warmup_hook = post_warmup_hook if post_warmup_hook is not None: - def guarded_pre_warmup_hook(): + def run_pre_warmup_hook(): nonlocal warmup_cleanup_pending if pre_warmup_hook is not None: pre_warmup_hook() warmup_cleanup_pending = True - def guarded_post_warmup_hook(): + def run_post_warmup_hook(): nonlocal warmup_cleanup_pending if not warmup_cleanup_pending: return warmup_cleanup_pending = False post_warmup_hook() + guarded_pre_warmup_hook = run_pre_warmup_hook + guarded_post_warmup_hook = run_post_warmup_hook + allocator_settings_guard = _AllocatorSettingsGuard() saved_fp8_tensors = None fp8_state_saved = False From b07d7df77085964df35f5e3a2dd5678f7bc75d68 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 07:54:23 +0900 Subject: [PATCH 15/16] Skip empty pipeline chunks in graph slot grouping Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 40 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 28 +++++++++++--------- 2 files changed, 56 insertions(+), 12 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 16b66f8e5a..6444741e0d 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -986,6 +986,46 @@ def forward(self, inp): reset_graphs(graphed) +@pytest.mark.parametrize( + "num_layers_per_chunk", + ([1, 0], [1, 0, 1], [1, 0, 0, 1]), + ids=("trailing", "middle", "consecutive-middle"), +) +def test_slot_memory_skips_zero_layer_pipeline_chunks(num_layers_per_chunk) -> None: + """Pipeline schedule entries without graphable layers must not consume slot metadata.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() + + module = Module().cuda() + num_chunks = len(num_layers_per_chunk) + num_graphable_layers = sum(num_layers_per_chunk) + samples = tuple( + (torch.ones(16, device="cuda", requires_grad=True),) for _ in range(num_graphable_layers) + ) + graphed = make_graphed_callables( + (module,) * num_graphable_layers, + samples, + num_warmup_iters=2, + _order=[*range(1, num_chunks + 1), *range(-num_chunks, 0)], + _num_layers_per_chunk=num_layers_per_chunk, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _slot(layer, layer, layer, overlap=layer, warmup=layer) + for layer in range(num_graphable_layers) + ), + ) + + try: + for graph in graphed: + inp = torch.randn(16, device="cuda", requires_grad=True) + graph(inp).sum().backward() + torch.testing.assert_close(inp.grad, 2.0 * inp.detach()) + finally: + reset_graphs(graphed) + + def test_slot_memory_preserves_aliased_public_outputs() -> None: """Slot output arenas retain overlapping public views across CP branches.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index b46e0fc0a7..6297f5a736 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -1985,18 +1985,22 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): continue m_chunk = abs(int(c_id)) - 1 logical_idx = group_fwd_idx[m_chunk] if c_id > 0 else group_bwd_idx[m_chunk] - first_func_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( - logical_idx * _num_layers_per_chunk[m_chunk] - ) - slot = _graph_memory_slots[first_func_idx] - event_key = ( - c_id > 0, - slot[3], - logical_idx, - slot[1], - _num_layers_per_chunk[m_chunk], - ) - group_records.append((event_key, slot[2])) + num_layers = _num_layers_per_chunk[m_chunk] + if num_layers == 0: + group_records.append(None) + else: + first_func_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( + logical_idx * num_layers + ) + slot = _graph_memory_slots[first_func_idx] + event_key = ( + c_id > 0, + slot[3], + logical_idx, + slot[1], + num_layers, + ) + group_records.append((event_key, slot[2])) if c_id > 0: group_fwd_idx[m_chunk] += 1 else: From d12ff87e505635e9ab978202661bdbb34b24eba5 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 11:26:13 +0900 Subject: [PATCH 16/16] Harden CUDA graph slot prerequisites Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 18 ++++++++++++++++++ transformer_engine/pytorch/graph.py | 5 ++--- 2 files changed, 20 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 6444741e0d..9511fbfc8e 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -939,6 +939,24 @@ def test_slot_memory_rejects_non_native_allocator(monkeypatch) -> None: ) +def test_slot_memory_rejects_missing_allocator_api(monkeypatch) -> None: + """Every private allocator API must be checked before slot preparation.""" + + module = torch.nn.Identity() + monkeypatch.setattr(torch.cuda.memory, "get_allocator_backend", lambda: "native") + monkeypatch.delattr(torch._C, "_has_Standard_Deleter") + + with pytest.raises(RuntimeError, match="_has_Standard_Deleter"): + te_graph._make_graphed_callables( + module, + (torch.ones(1),), + _order=[1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0),), + ) + + @pytest.mark.parametrize("elements", (0, 4096), ids=("empty", "nonempty")) def test_slot_memory_variants_share_one_backing(elements: int) -> None: """Mutually exclusive variants must use identical slot storage and output addresses.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 6297f5a736..573abd70b9 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -421,6 +421,7 @@ def _make_graphed_callables( "_cuda_setCheckpointPoolState", "_cuda_checkPoolLiveAllocations", "_free_And_Remove_DeleterFn", + "_has_Standard_Deleter", ) missing_checkpoint_apis = [ name for name in required_checkpoint_apis if not hasattr(torch._C, name) @@ -622,9 +623,7 @@ def _make_graphed_callables( graph_callables = [None for _ in range(len(flatten_sample_args))] # For cases with multiple active RNG states, e.g. TP. - if graph_safe_rng_available() and not bool( - int(os.getenv("NVTE_DISABLE_GRAPH_SAFE_RNG_REGISTRATION", "0")) - ): + if graph_safe_rng_available(): for _, state in get_all_rng_states().items(): for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs): fwd_graph.register_generator_state(state)