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/30] 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/30] [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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] 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/30] [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/30] 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/30] 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/30] 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/30] 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) From 28980c2a35d9c5c4491beaa94601bb0bad3060b0 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 12:17:46 +0900 Subject: [PATCH 17/30] Fix typed CUDA graph arena offsets Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 10 ++++++++++ transformer_engine/pytorch/graph.py | 5 ++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 9511fbfc8e..968e7b95af 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -74,6 +74,16 @@ class Group: assert self_transport_group_ref() is None +def test_slot_memory_arena_view_uses_typed_storage_offset(): + backing = torch.empty(16, dtype=torch.float32, device="cuda") + arena = backing[4:] + spec = ("output", None, (1,), (1,), torch.float32, backing.device, False, 4) + + view = te_graph._arena_view(arena, 4, spec) + + assert view.data_ptr() == backing.data_ptr() + 5 * backing.element_size() + + def nvfp4_vanilla(): nvfp4_recipe = recipe.NVFP4BlockScaling() nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 573abd70b9..e509188c5b 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -168,9 +168,12 @@ def _slot_output_plan(outputs: Sequence[Any], func: Callable) -> List[Optional[T 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]) + storage_offset_bytes = arena.storage_offset() * arena.element_size() + offset + if storage_offset_bytes % target.element_size(): + raise RuntimeError("CUDA graph arena view has an unaligned byte offset.") return target.set_( arena.untyped_storage(), - (arena.storage_offset() + offset) // target.element_size(), + storage_offset_bytes // target.element_size(), spec[2], spec[3], ) From 9f4ec8ab66844bc9205d12249184a43c1194dbf5 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 13:09:14 +0900 Subject: [PATCH 18/30] Test CUDA graph staging input rebinding Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 39 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 3 ++- 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 968e7b95af..7dc4631cfc 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1198,6 +1198,45 @@ def forward(self, inp): reset_graphs(graphed) +def test_slot_memory_rebinds_mutable_sample_args_to_staging_surfaces() -> None: + """Mutable sample banks expose the exact static inputs selected by slot staging.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() + + module = Module().cuda() + shared = torch.ones(4096, device="cuda", requires_grad=True) + samples = [(shared,), (shared,)] + 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: + assert any(tensor is shared for tensor in (samples[0][0], samples[1][0])) + assert samples[0][0].data_ptr() != samples[1][0].data_ptr() + + 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) + + def test_slot_memory_snapshots_shared_kwarg_across_alternate_liveness() -> None: """A later forward must not overwrite a shared kwarg needed by backward.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index e509188c5b..3973dffb4e 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -3150,7 +3150,8 @@ def make_graphed_callables( 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. + to the graph pool. When ``sample_args`` is a mutable list, entries whose leading input is + rebound to a staging surface are updated in place. pre_warmup_hook: callable, default = None A hook function that will be called before the warmup iterations. post_warmup_hook: callable, default = None From 7741feaef72d9fcb94f190d6b5dbb12b4d208d62 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sat, 22 Aug 2026 23:58:01 +0900 Subject: [PATCH 19/30] Preserve CUDA graph staging input aliases Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 69 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 32 ++++++++++++- 2 files changed, 99 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 7dc4631cfc..97b8778dd8 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1237,6 +1237,75 @@ def forward(self, inp): reset_graphs(graphed) +def test_slot_memory_rejects_warmup_input_alias_merge() -> None: + """Warmup-plan aliases must preserve the full input-storage alias topology.""" + + class Module(torch.nn.Module): + def forward(self, left, right): + return left.square() + right.square() + + module = Module().cuda() + left = torch.ones(4096, device="cuda", requires_grad=True) + right = torch.ones(4096, device="cuda", requires_grad=True) + merged = torch.ones(4096, device="cuda", requires_grad=True) + + with pytest.raises(RuntimeError, match="merge distinct input storages"): + make_graphed_callables( + (module,), + ((left, right), (merged, merged)), + 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), + ), + ) + assert not te_graph.is_graph_capturing() + + +def test_slot_memory_staging_preserves_nonleading_input_aliases() -> None: + """Staging must not choose a leading-input address used by another argument.""" + + class Module(torch.nn.Module): + def forward(self, left, right): + return left.square() + right.square() + + module = Module().cuda() + branch0_left = torch.ones(4096, device="cuda", requires_grad=True) + branch0_right = torch.ones(4096, device="cuda", requires_grad=True) + branch1_left = torch.ones(4096, device="cuda", requires_grad=True) + samples = [ + (branch0_left, branch0_right), + (branch1_left, branch0_left), + ] + 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, warmup=0), + _slot(0, 1, 0, warmup=1), + ), + ) + + try: + assert samples[1][0].data_ptr() != samples[1][1].data_ptr() + left = torch.full((4096,), 2.0, device="cuda", requires_grad=True) + right = torch.full((4096,), 3.0, device="cuda", requires_grad=True) + output = graphed[1](left, right) + torch.testing.assert_close(output, left.detach().square() + right.detach().square()) + output.sum().backward() + torch.testing.assert_close(left.grad, 2.0 * left.detach()) + torch.testing.assert_close(right.grad, 2.0 * right.detach()) + finally: + 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.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 3973dffb4e..83e91c9ccb 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -786,6 +786,7 @@ def clone_warmup_plan(template_idx, target_idx): ) external_storage_map = {} + source_storage_by_target = {} for source, target in zip(source_args, target_args): if _input_staging_key(source) != _input_staging_key(target): raise RuntimeError( @@ -800,6 +801,12 @@ def clone_warmup_plan(template_idx, target_idx): f"Warmup-plan alias {target_idx} changes an input storage alias from " f"{previous_target} to {target_ptr}." ) + previous_source = source_storage_by_target.setdefault(target_ptr, source_ptr) + if previous_source != source_ptr: + raise RuntimeError( + f"Warmup-plan alias {target_idx} would merge distinct input storages " + f"{previous_source} and {source_ptr} into {target_ptr}." + ) source_plan = per_callable_saved_tensor_plans[template_idx] if source_plan is None: @@ -1200,9 +1207,22 @@ def hook_fn( storage_owner = {} staging_targets = {} + def staging_target_preserves_aliases(group_idx, storage_id): + """Check that rebinding the leading input keeps every input-storage alias.""" + for func_idx in staging_groups[group_idx]["members"]: + static_input_surface = per_callable_static_input_surfaces[func_idx] + old_storage_id = _tensor_storage_ptr(static_input_surface[0]) + for other_input in static_input_surface[1:]: + other_storage_id = _tensor_storage_ptr(other_input) + if (old_storage_id == other_storage_id) != (storage_id == other_storage_id): + return False + return True + 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: + if storage_id in seen_storages or not staging_target_preserves_aliases( + group_idx, storage_id + ): continue seen_storages.add(storage_id) previous_group = storage_owner.get(storage_id) @@ -1214,13 +1234,21 @@ def match_staging_group(group_idx, seen_storages): for group_idx in sorted( range(len(staging_groups)), - key=lambda index: len(staging_groups[index]["candidates"]), + key=lambda index: ( + staging_target_preserves_aliases(index, None), + 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: + if not staging_target_preserves_aliases(group_idx, None): + raise RuntimeError( + "CUDA graph slot staging cannot preserve a leading-input storage alias " + f"for group {group_idx}." + ) source = next(iter(group["candidates"].values())) signature = _saved_tensor_signature(source) storage_numel = signature[-1] // source.element_size() From 6ca7a88309186a8741fc85aeb8e2d60d3c60ee1f Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 00:26:05 +0900 Subject: [PATCH 20/30] Reject incompatible CUDA graph output aliases Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 40 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 24 +++++++++++++++++ 2 files changed, 64 insertions(+) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 97b8778dd8..53c4d0d787 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1093,6 +1093,46 @@ def forward(self, inp): reset_graphs(graphed) +@pytest.mark.parametrize( + "first_mode,second_mode", + (("independent", "offset8"), ("offset8", "independent"), ("offset4", "offset8")), +) +def test_slot_memory_rejects_branch_output_alias_mismatch(first_mode, second_mode) -> None: + """CP branches sharing an output slot must expose the same storage aliases.""" + + class Module(torch.nn.Module): + def __init__(self, mode): + super().__init__() + self.mode = mode + + def forward(self, inp): + output = inp.square() + first = output[:8] + if self.mode == "independent": + return first, output[8:].clone() + offset = 4 if self.mode == "offset4" else 8 + return first, output[offset : offset + 8] + + modules = (Module(first_mode).cuda(), Module(second_mode).cuda()) + samples = tuple((torch.ones(16, device="cuda", requires_grad=True),) for _ in modules) + graphed = None + try: + with pytest.raises(RuntimeError, match="incompatible output storage aliases"): + graphed = make_graphed_callables( + modules, + 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)), + ) + assert not te_graph.is_graph_capturing() + finally: + if graphed is not None: + 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.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 83e91c9ccb..77caae2b15 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -1290,6 +1290,25 @@ def prepare_native_io_targets(per_callable_plans, kind): if per_callable_plans is None: return None + def storage_alias_topology(plan): + indexed_specs = [ + (tensor_idx, spec) + for tensor_idx, spec in enumerate(plan) + if spec is not None and spec[0] == kind + ] + if any(len(spec) < 10 for _, spec in indexed_specs): + return None + return tuple( + ( + left_idx, + right_idx, + left[8] == right[8], + right[9] - left[9] if left[8] == right[8] else None, + ) + for position, (left_idx, left) in enumerate(indexed_specs) + for right_idx, right in indexed_specs[position + 1 :] + ) + plans_by_family = {} for func_idx, plan in enumerate(per_callable_plans): arena_id, branch_id = slot_io_memory_alias_groups[func_idx] @@ -1307,6 +1326,11 @@ def prepare_native_io_targets(per_callable_plans, kind): raise RuntimeError( f"CUDA graph {kind} family {family} exposes different tensor counts." ) + alias_topologies = {storage_alias_topology(plan) for plan in plans} + if len(alias_topologies) != 1: + raise RuntimeError( + f"CUDA graph {kind} family {family} has incompatible {kind} storage aliases." + ) for tensor_idx, specs in enumerate(zip(*plans)): if all(spec is None for spec in specs): continue From f79956296b5585f88db3b5e60df1c4603693db30 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 01:12:53 +0900 Subject: [PATCH 21/30] Reject CUDA graph inputs aliased to module state Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 31 +++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 35 ++++++++++++++++++++++++----- 2 files changed, 60 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 53c4d0d787..f2d700d876 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1346,6 +1346,37 @@ def forward(self, left, right): reset_graphs(graphed) +@pytest.mark.parametrize("state_kind", ("parameter", "buffer")) +def test_slot_memory_rejects_user_inputs_aliased_to_module_state(state_kind) -> None: + """Runtime input staging must not overwrite parameters or buffers.""" + + class Module(torch.nn.Module): + def __init__(self): + super().__init__() + state = torch.ones(16, device="cuda") + if state_kind == "parameter": + self.state = torch.nn.Parameter(state) + else: + self.register_buffer("state", state) + + def forward(self, inp): + return 2.0 * inp + self.state + + module = Module() + sample = module.state.detach().view_as(module.state).requires_grad_(True) + with pytest.raises(RuntimeError, match="sharing storage with a module parameter or buffer"): + 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_snapshots_shared_kwarg_across_alternate_liveness() -> None: """A later forward must not overwrite a shared kwarg needed by backward.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 77caae2b15..95fe8f694e 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -590,6 +590,7 @@ def _make_graphed_callables( # torch.cuda.make_graphed_callables. per_callable_len_user_args = [len(args) for args in flatten_sample_args] if _order is None: + per_callable_funcs = callables per_callable_module_params = [ tuple(c.parameters()) if isinstance(c, torch.nn.Module) else () for c in callables ] @@ -597,17 +598,15 @@ def _make_graphed_callables( flatten_sample_args[i] + per_callable_module_params[i] for i in range(len(callables)) ] else: + per_callable_funcs = [] per_callable_module_params = [] for m_chunk in range(num_model_chunks): for _ in range(num_microbatches): for l_no in range(_num_layers_per_chunk[m_chunk]): + func = callables[_prefix_num_layers[m_chunk] + l_no] + per_callable_funcs.append(func) per_callable_module_params.append( - tuple(callables[_prefix_num_layers[m_chunk] + l_no].parameters()) - if isinstance( - callables[_prefix_num_layers[m_chunk] + l_no], - torch.nn.Module, - ) - else () + tuple(func.parameters()) if isinstance(func, torch.nn.Module) else () ) if len(per_callable_module_params) != len(flatten_sample_args): raise ValueError( @@ -620,6 +619,30 @@ def _make_graphed_callables( for i in range(len(flatten_sample_args)) ] + if use_slot_memory: + for func_idx, (user_args, func) in enumerate(zip(flatten_sample_args, per_callable_funcs)): + module_state = ( + (*func.parameters(), *func.buffers()) + if isinstance(func, torch.nn.Module) + else () + ) + module_state_storages = { + _tensor_storage_ptr(tensor) + for tensor in module_state + if tensor.is_cuda and tensor.untyped_storage().nbytes() > 0 + } + for arg_idx, arg in enumerate(user_args): + if ( + arg.is_cuda + and arg.untyped_storage().nbytes() > 0 + and _tensor_storage_ptr(arg) in module_state_storages + ): + raise RuntimeError( + "CUDA graph slot memory does not support user input tensor " + f"{arg_idx} sharing storage with a module parameter or buffer for " + f"graph input {func_idx}." + ) + fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] bwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] bwd_dw_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] From 0c71af03018c6700a87de4cc9532325018d784ea Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 02:03:56 +0900 Subject: [PATCH 22/30] Reject cross-callable CUDA graph state aliases Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 37 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 26 ++++++++++---------- 2 files changed, 50 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index f2d700d876..f2f8e950d9 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1377,6 +1377,43 @@ def forward(self, inp): assert not te_graph.is_graph_capturing() +@pytest.mark.parametrize("state_kind", ("parameter", "buffer")) +def test_slot_memory_rejects_user_inputs_aliased_to_other_callable_state(state_kind) -> None: + """Staging candidates must not alias state owned by another graph callable.""" + + class StatelessModule(torch.nn.Module): + def forward(self, inp): + return inp.square() + + class StatefulModule(torch.nn.Module): + def __init__(self): + super().__init__() + state = torch.ones(16, device="cuda") + if state_kind == "parameter": + self.state = torch.nn.Parameter(state) + else: + self.register_buffer("state", state) + + def forward(self, inp): + return 2.0 * inp + self.state + + stateless = StatelessModule() + stateful = StatefulModule() + cross_callable_alias = stateful.state.detach().view_as(stateful.state).requires_grad_(True) + independent = torch.ones_like(stateful.state, requires_grad=True) + with pytest.raises(RuntimeError, match="sharing storage with a module parameter or buffer"): + make_graphed_callables( + (stateless, stateful), + ((cross_callable_alias,), (independent,)), + 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)), + ) + assert not te_graph.is_graph_capturing() + + def test_slot_memory_snapshots_shared_kwarg_across_alternate_liveness() -> None: """A later forward must not overwrite a shared kwarg needed by backward.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 95fe8f694e..6960865516 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -620,17 +620,17 @@ def _make_graphed_callables( ] if use_slot_memory: - for func_idx, (user_args, func) in enumerate(zip(flatten_sample_args, per_callable_funcs)): - module_state = ( - (*func.parameters(), *func.buffers()) - if isinstance(func, torch.nn.Module) - else () - ) - module_state_storages = { - _tensor_storage_ptr(tensor) - for tensor in module_state - if tensor.is_cuda and tensor.untyped_storage().nbytes() > 0 - } + module_state_storages = set() + inspected_modules = set() + for func in per_callable_funcs: + if not isinstance(func, torch.nn.Module) or id(func) in inspected_modules: + continue + inspected_modules.add(id(func)) + for tensor in (*func.parameters(), *func.buffers()): + if tensor.is_cuda and tensor.untyped_storage().nbytes() > 0: + module_state_storages.add(_tensor_storage_ptr(tensor)) + + for func_idx, user_args in enumerate(flatten_sample_args): for arg_idx, arg in enumerate(user_args): if ( arg.is_cuda @@ -639,8 +639,8 @@ def _make_graphed_callables( ): raise RuntimeError( "CUDA graph slot memory does not support user input tensor " - f"{arg_idx} sharing storage with a module parameter or buffer for " - f"graph input {func_idx}." + f"{arg_idx} sharing storage with a module parameter or buffer in the " + f"capture bank for graph input {func_idx}." ) fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] From 785c03cde82a5d8a8b230270aabd9868e8381b3e Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 02:08:51 +0900 Subject: [PATCH 23/30] Recheck CUDA graph state aliases after warmup Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 24 ++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 9 +++++++-- 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index f2f8e950d9..b428bd3a47 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1414,6 +1414,30 @@ def forward(self, inp): assert not te_graph.is_graph_capturing() +def test_slot_memory_rejects_input_promoted_to_module_state_during_warmup() -> None: + """Lazy module state must be rechecked before input staging is selected.""" + + class LazyStateModule(torch.nn.Module): + def forward(self, inp): + if "state" not in self._buffers: + self.register_buffer("state", inp.detach()) + return 2.0 * inp + self.state + + module = LazyStateModule() + sample = torch.ones(16, device="cuda", requires_grad=True) + with pytest.raises(RuntimeError, match="module parameter or buffer.*after warmup"): + 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_snapshots_shared_kwarg_across_alternate_liveness() -> None: """A later forward must not overwrite a shared kwarg needed by backward.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 6960865516..43221859d8 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -619,7 +619,8 @@ def _make_graphed_callables( for i in range(len(flatten_sample_args)) ] - if use_slot_memory: + def validate_slot_user_input_state_aliases(phase): + """Reject user inputs backed by any persistent module-state storage.""" module_state_storages = set() inspected_modules = set() for func in per_callable_funcs: @@ -640,9 +641,12 @@ def _make_graphed_callables( raise RuntimeError( "CUDA graph slot memory does not support user input tensor " f"{arg_idx} sharing storage with a module parameter or buffer in the " - f"capture bank for graph input {func_idx}." + f"capture bank for graph input {func_idx} {phase}." ) + if use_slot_memory: + validate_slot_user_input_state_aliases("before warmup") + fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] bwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] bwd_dw_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] @@ -1190,6 +1194,7 @@ def hook_fn( torch.cuda.synchronize() if use_slot_memory: + validate_slot_user_input_state_aliases("after warmup") per_callable_param_grad_tensor_targets = [ [None] * len(static_input_surface) for static_input_surface in per_callable_static_input_surfaces From be515dfd436defa6c30315a8340230692cf14dd6 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 02:55:08 +0900 Subject: [PATCH 24/30] Reject non-Tensor CUDA graph slot kwargs Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 26 ++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 2 ++ 2 files changed, 28 insertions(+) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index b428bd3a47..6b33e36a0b 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1438,6 +1438,32 @@ def forward(self, inp): assert not te_graph.is_graph_capturing() +def test_slot_memory_rejects_non_tensor_sample_kwargs() -> None: + """Slot-memory static surfaces must reject non-Tensor kwargs explicitly.""" + + class Module(torch.nn.Module): + def forward(self, inp, use_bias=False): + output = inp.square() + if use_bias: + output = output + 1.0 + return output.sum().unsqueeze(0) + + module = Module().cuda() + sample = torch.ones(8, device="cuda", requires_grad=True) + with pytest.raises(TypeError, match="slot memory sample_kwargs must contain only Tensors"): + make_graphed_callables( + (module,), + ((sample,),), + sample_kwargs=({"use_bias": True},), + num_warmup_iters=1, + _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_snapshots_shared_kwarg_across_alternate_liveness() -> None: """A later forward must not overwrite a shared kwarg needed by backward.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 43221859d8..5fb4168257 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -581,6 +581,8 @@ def _make_graphed_callables( "In the beta API, sample_args " + "for each callable must contain only Tensors. Other types are not allowed." ) + if use_slot_memory and not all(isinstance(arg, torch.Tensor) for arg in flatten_kwarg): + raise TypeError("CUDA graph slot memory sample_kwargs must contain only Tensors.") # If a callable is an nn.Module, its graph's full input surface is the args the user explicitly # passes to forward (ie, its sample_args) AND the module's parameter attributes. From f5a36a9141e651aa946a350e028e717f5a20aaf7 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 18:15:42 +0900 Subject: [PATCH 25/30] Simplify CUDA graph slot checkpoint state Signed-off-by: Tailai Ma --- transformer_engine/pytorch/graph.py | 277 ++++++++++++---------------- 1 file changed, 114 insertions(+), 163 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 5fb4168257..104d934751 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -7,6 +7,7 @@ from bisect import bisect_right from collections.abc import Iterable import contextlib +from dataclasses import dataclass import gc import os import warnings @@ -69,6 +70,15 @@ def restore(self) -> None: self._settings_to_restore = None +@dataclass(frozen=True) +class _GraphPoolCheckpoint: + """Allocator state and Python storage owners at one graph-pool boundary.""" + + state: Any + blocks: Dict[int, int] + owners: Dict[int, Any] + + 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() @@ -165,20 +175,25 @@ def _slot_output_plan(outputs: Sequence[Any], func: Callable) -> List[Optional[T 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.""" +def _storage_view(storage, offset: int, spec: Tuple[Any, ...]) -> torch.Tensor: + """Materialize a typed tensor view at a byte offset in a CUDA storage.""" target = torch.empty((0,), dtype=spec[4], device=spec[5]) - storage_offset_bytes = arena.storage_offset() * arena.element_size() + offset - if storage_offset_bytes % target.element_size(): + if offset % target.element_size(): raise RuntimeError("CUDA graph arena view has an unaligned byte offset.") return target.set_( - arena.untyped_storage(), - storage_offset_bytes // target.element_size(), + storage, + offset // target.element_size(), spec[2], spec[3], ) +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.""" + offset += arena.storage_offset() * arena.element_size() + return _storage_view(arena.untyped_storage(), offset, spec) + + def set_capture_start() -> None: """Record beginning of `make_graphed_callables`.""" global _IS_GRAPH_CAPTURING @@ -1315,6 +1330,11 @@ def match_staging_group(group_idx, seen_storages): static_input_surface[0] = input_target per_callable_static_input_surfaces[func_idx] = tuple(static_input_surface) + def native_io_family(func_idx): + """Return the slot/liveness family shared by mutually exclusive CP branches.""" + arena_id, _ = slot_io_memory_alias_groups[func_idx] + return (arena_id, *slot_io_liveness_groups[func_idx]) + 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: @@ -1341,8 +1361,8 @@ def storage_alias_topology(plan): 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_id = slot_io_memory_alias_groups[func_idx] + family = native_io_family(func_idx) branch_plans = plans_by_family.setdefault(family, {}) if branch_id in branch_plans: raise RuntimeError( @@ -1396,8 +1416,7 @@ def storage_alias_topology(plan): 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]) + family = native_io_family(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": {}} @@ -1405,8 +1424,7 @@ def storage_alias_topology(plan): 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]) + family = native_io_family(func_idx) key = (*family, tensor_idx) anchors = native_io_anchors[kind] counts = native_io_capture_counts[kind] @@ -1537,6 +1555,29 @@ def copy_user_grads_to_slot_arena(func_idx, static_input_surface, grad_inputs): per_callable_native_saved_storages = [{} for _ in flatten_sample_args] + def native_saved_storage_components(plan): + """Return overlapping byte-range components for native saved storages.""" + 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, [(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) + return components + def plan_native_saved_alias_targets( plan, canonical_targets, @@ -1582,7 +1623,6 @@ def plan_native_saved_alias_targets( "cursor": 0, } ) - records_by_storage_group = {} for saved_idx, spec in enumerate(plan): if spec[0] != "native": continue @@ -1597,24 +1637,11 @@ def plan_native_saved_alias_targets( 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: + for component_start, component_end, component_records in native_saved_storage_components( + plan + ): component_alignment = max( plan[saved_idx][4].itemsize for _, saved_idx in component_records ) @@ -1678,12 +1705,7 @@ def plan_native_saved_alias_targets( 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 = _storage_view(bank["storage"], target_offset, spec) target.requires_grad_(spec[6]) target_views[saved_idx] = target @@ -1704,27 +1726,10 @@ def semantic_boundary_alias_targets(func_idx, outputs): 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): - 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: + for component_start, component_end, component_records in native_saved_storage_components( + plan + ): + component_saved_indices = [saved_idx for _, saved_idx in component_records] candidate_aliases = [ (saved_idx, aliases[saved_idx]) for saved_idx in component_saved_indices @@ -1796,12 +1801,7 @@ def semantic_boundary_alias_targets(func_idx, outputs): 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 = _storage_view(anchor_storage, target_offset, spec) target.requires_grad_(spec[6]) targets[saved_idx] = target return tuple(targets) @@ -2050,14 +2050,8 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): # 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_storage_owners = None - branch_pre_checkpoint_state = None - branch_pre_checkpoint_live_blocks = None - branch_pre_checkpoint_storage_owners = None - current_storage_owners = None - native_saved_alias_targets = None + branch_checkpoint = None + branch_pre_checkpoint = None if use_slot_memory: branch_capture_groups = [None] * len(_order) group_records = [] @@ -2124,40 +2118,27 @@ def slot_pool_active_blocks(): if block["state"] == "active_allocated" } - def slot_pool_layout(): - """Return the allocator block topology for checkpoint diagnostics.""" + def slot_pool_pending_frees(): + """Return graph-pool blocks waiting for a cross-stream free event.""" 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"] - ], - } + return [ + (block["address"], block["size"]) for segment in segments - } + for block in segment["blocks"] + if block["state"] == "active_pending_free" + ] 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"] - ): + if not slot_pool_pending_frees(): 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" - ] + remaining = slot_pool_pending_frees() if remaining: raise RuntimeError( f"CUDA graph slot checkpoint could not drain pending allocator frees: {remaining}." @@ -2286,6 +2267,17 @@ def visit(value): ) return owners + def checkpoint_slot_pool_boundary(phase, extra_values=(), save_state=True): + """Snapshot one allocator boundary and resolve all of its live storage owners.""" + state = ( + torch._C._cuda_getCheckpointState(torch.cuda.current_device(), mempool) + if save_state + else None + ) + blocks = slot_pool_active_blocks() + owners = checkpoint_live_storage_owners(blocks, phase, extra_values) + return _GraphPoolCheckpoint(state, blocks, owners) + def validate_checkpoint_live_storages(expected_blocks, owners): """Validate checkpoint owners before changing any StorageImpl.""" if set(owners) != set(expected_blocks): @@ -2335,14 +2327,7 @@ def verify_checkpoint_live_storages(expected_blocks, owners): "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, - ): + def restore_slot_pool_boundary(current, target, phase): """Switch between two allocator boundaries while preserving their StorageImpls.""" device = torch.cuda.current_device() rollback_state = torch._C._cuda_getCheckpointState(device, mempool) @@ -2350,17 +2335,17 @@ def restore_slot_pool_boundary( 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) + validate_checkpoint_live_storages(current.blocks, current.owners) current_release_started = True - release_checkpoint_live_storages(current_owners, phase) + release_checkpoint_live_storages(current.owners, phase) torch._C._cuda_setCheckpointPoolState( device, - target_state, + target.state, [], - [storage._cdata for storage in target_owners.values()], + [storage._cdata for storage in target.owners.values()], ) - verify_checkpoint_live_storages(target_blocks, target_owners) - assert_slot_pool_liveness(target_blocks, phase) + verify_checkpoint_live_storages(target.blocks, target.owners) + assert_slot_pool_liveness(target.blocks, phase) except BaseException as restore_error: if not current_release_started: raise @@ -2369,7 +2354,7 @@ def restore_slot_pool_boundary( rollback_owners = checkpoint_live_storage_owners( rollback_blocks, f"{phase} rollback", - (current_owners, target_owners), + (current.owners, target.owners), ) validate_checkpoint_live_storages(rollback_blocks, rollback_owners) release_checkpoint_live_storages(rollback_owners, f"{phase} rollback") @@ -2377,10 +2362,10 @@ def restore_slot_pool_boundary( device, rollback_state, [], - [storage._cdata for storage in current_owners.values()], + [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") + 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}." @@ -2404,14 +2389,8 @@ def restore_slot_pool_boundary( captured_branch_func_indices = [] branch_stale_storages = {} branch_boundary_values = [] - branch_checkpoint_active = any( - value is not None - for value in ( - branch_checkpoint_state, - branch_pre_checkpoint_state, - current_storage_owners, - native_saved_alias_targets, - ) + branch_checkpoint_active = ( + branch_checkpoint is not None or branch_pre_checkpoint is not None ) if branch_group is not None and branch_group[0] == 0 and branch_checkpoint_active: raise RuntimeError("CUDA graph slot checkpoint groups overlap.") @@ -2420,22 +2399,14 @@ def restore_slot_pool_boundary( 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, + branch_pre_checkpoint = checkpoint_slot_pool_boundary( 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, + branch_checkpoint, + branch_pre_checkpoint, f"{'forward' if c_id > 0 else 'backward'} branch " f"{branch_group[0] + 1}/{branch_group[1]} restore pre-boundary", ) @@ -2455,15 +2426,13 @@ 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 = ( + 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 - ): + with capture_saved_tensors(per_callable_fwd_idx, saved_alias_targets): outputs = func(*args, **kwargs) flatten_outputs, spec = _tree_flatten(outputs) original_flatten_outputs = flatten_outputs @@ -2481,7 +2450,6 @@ def restore_slot_pool_boundary( 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 @@ -2694,48 +2662,35 @@ def restore_slot_pool_boundary( torch.cuda.synchronize() drain_slot_pool_pending_frees() if branch_group[0] == 0: - branch_checkpoint_state = torch._C._cuda_getCheckpointState( - torch.cuda.current_device(), mempool - ) - branch_checkpoint_live_blocks = slot_pool_active_blocks() - branch_checkpoint_storage_owners = checkpoint_live_storage_owners( - branch_checkpoint_live_blocks, + branch_checkpoint = checkpoint_slot_pool_boundary( 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, + branch_checkpoint, + branch_pre_checkpoint, f"backward branch 1/{branch_group[1]} restore pre-boundary", ) else: - current_live_blocks = slot_pool_active_blocks() - current_storage_owners = checkpoint_live_storage_owners( - current_live_blocks, + current_checkpoint = checkpoint_slot_pool_boundary( 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, + branch_pre_checkpoint.owners, + branch_checkpoint.owners, ), + save_state=False, ) restore_slot_pool_boundary( - current_live_blocks, - current_storage_owners, - branch_checkpoint_state, - branch_checkpoint_live_blocks, - branch_checkpoint_storage_owners, + current_checkpoint, + branch_checkpoint, 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 ) @@ -2744,12 +2699,8 @@ def restore_slot_pool_boundary( 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_storage_owners = None - branch_pre_checkpoint_state = None - branch_pre_checkpoint_live_blocks = None - branch_pre_checkpoint_storage_owners = None + branch_checkpoint = None + branch_pre_checkpoint = None gc.collect() torch.cuda.synchronize() drain_slot_pool_pending_frees() @@ -2827,7 +2778,7 @@ def restore_slot_pool_boundary( 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: + if branch_checkpoint is not None or branch_pre_checkpoint is not None: raise RuntimeError("CUDA graph capture ended inside a slot checkpoint group.") if allocator_settings_to_restore is not None: From ff6697aa9ac06575e90248735911b51b04440536 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 18:45:13 +0900 Subject: [PATCH 26/30] Remove CUDA graph allocator checkpointing Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 180 +------ .../pytorch/csrc/extensions/pybind.cpp | 9 - transformer_engine/pytorch/graph.py | 502 +----------------- 3 files changed, 30 insertions(+), 661 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 6b33e36a0b..d2c94cc8e3 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -794,6 +794,15 @@ def _slot(saved_arena, branch, io_arena, overlap=0, frame=0, warmup=0, user_grad return (saved_arena, io_arena, branch, overlap, frame, warmup, user_grad) +def _variant_major_order(base_order, variants, model_chunks=1): + """Repeat one complete PP/VPP schedule for each mutually exclusive variant.""" + return [ + (1 if chunk > 0 else -1) * (abs(chunk) + variant * model_chunks) + for variant in range(variants) + for chunk in base_order + ] + + 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() @@ -932,41 +941,6 @@ def fail_rng_preparation(): 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),), - ) - - -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.""" @@ -1179,7 +1153,7 @@ def forward(self, inp): def test_slot_memory_rejects_public_cuda_tensor_subclass_outputs() -> None: - """Unsupported CUDA outputs must fail before allocator checkpoint mutation.""" + """Unsupported CUDA outputs must fail before graph capture starts.""" class OutputTensor(torch.Tensor): pass @@ -1501,8 +1475,8 @@ def forward(self, inp, scale): @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.""" +def test_slot_memory_reuses_variant_major_branches(reverse_replay) -> None: + """Complete CP schedules must reuse one slot backing across variants.""" class Module(torch.nn.Module): def forward(self, inp): @@ -1520,10 +1494,7 @@ def forward(self, inp): samples, num_warmup_iters=2, allow_unused_input=True, - _order=[ - *(variant + 1 for variant in range(variants)), - *(-variant - 1 for variant in range(variants)), - ], + _order=_variant_major_order([1, -1], variants), _num_layers_per_chunk=[1] * variants, _reuse_graph_input_output_buffers=True, _graph_memory_slots=tuple( @@ -1543,7 +1514,7 @@ def forward(self, inp): expected = 6.0 * inp.detach() if not torch.allclose(inp.grad, expected): print( - "CHECKPOINT_GRAD_MISMATCH", + "VARIANT_MAJOR_GRAD_MISMATCH", { "reverse_replay": reverse_replay, "variant": variant, @@ -1565,91 +1536,7 @@ def forward(self, inp): reset_graphs(graphed) -@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.""" - - 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) - 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 - 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 (detach|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 == (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) - 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: +def test_slot_memory_reclaims_retained_variant_outputs() -> None: """A module-held source output must not pin a mutually exclusive branch allocation.""" class Module(torch.nn.Module): @@ -1672,10 +1559,7 @@ def forward(self, inp): samples, num_warmup_iters=2, allow_unused_input=True, - _order=[ - *(variant + 1 for variant in range(variants)), - *(-variant - 1 for variant in range(variants)), - ], + _order=_variant_major_order([1, -1], variants), _num_layers_per_chunk=[1] * variants, _reuse_graph_input_output_buffers=True, _graph_memory_slots=tuple( @@ -1697,8 +1581,8 @@ def forward(self, inp): reset_graphs(graphed) -def test_slot_memory_checkpoint_tracks_native_saved_storage() -> None: - """Checkpoint branches must retain allocator owners hidden in autograd saved tensors.""" +def test_slot_memory_tracks_native_saved_storage() -> None: + """Variant-major capture must preserve native autograd saved tensors.""" class Module(torch.nn.Module): def forward(self, inp): @@ -1712,10 +1596,7 @@ def forward(self, inp): (module,) * variants, samples, num_warmup_iters=2, - _order=[ - *(variant + 1 for variant in range(variants)), - *(-variant - 1 for variant in range(variants)), - ], + _order=_variant_major_order([1, -1], variants), _num_layers_per_chunk=[1] * variants, _reuse_graph_input_output_buffers=True, _graph_memory_slots=tuple( @@ -1753,17 +1634,11 @@ def capture(variants): (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), - ], + _order=_variant_major_order([1, 1, -1, -1], variants), _num_layers_per_chunk=[1] * variants, _reuse_graph_input_output_buffers=True, _graph_memory_slots=tuple( @@ -1812,7 +1687,7 @@ def capture(variants): reset_graphs(graphed) -def test_slot_memory_checkpoint_aliases_parameter_gradients() -> None: +def test_slot_memory_aliases_parameter_gradients() -> None: """Mutually exclusive branches must return parameter grads from one slot address.""" class Module(torch.nn.Module): @@ -1831,10 +1706,7 @@ def forward(self, inp): samples, num_warmup_iters=2, allow_unused_input=True, - _order=[ - *(variant + 1 for variant in range(variants)), - *(-variant - 1 for variant in range(variants)), - ], + _order=_variant_major_order([1, -1], variants), _num_layers_per_chunk=[1] * variants, _reuse_graph_input_output_buffers=True, _graph_memory_slots=tuple( @@ -1983,13 +1855,9 @@ def forward(self, inp): ) ) - # PP=2, VPP=2, rank 0, four-microbatch schedule with lockstep CP branches. + # PP=2, VPP=2, rank 0, four-microbatch schedule repeated per CP variant. 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) + order = _variant_major_order(base_order, variants, model_chunks) graphed = make_graphed_callables( (module,) * (variants * model_chunks), diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 1a86caef57..18da5d0e9f 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -6,8 +6,6 @@ #include "pybind.h" -#include -#include #include #include #include @@ -137,13 +135,6 @@ 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 104d934751..c8dbded65a 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -4,10 +4,8 @@ """Functions for CUDA Graphs support in FP8""" -from bisect import bisect_right from collections.abc import Iterable import contextlib -from dataclasses import dataclass import gc import os import warnings @@ -70,15 +68,6 @@ def restore(self) -> None: self._settings_to_restore = None -@dataclass(frozen=True) -class _GraphPoolCheckpoint: - """Allocator state and Python storage owners at one graph-pool boundary.""" - - state: Any - blocks: Dict[int, int] - owners: Dict[int, Any] - - 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() @@ -427,34 +416,6 @@ 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", - "_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) - ] - 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: @@ -1553,8 +1514,6 @@ 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 = [{} for _ in flatten_sample_args] - def native_saved_storage_components(plan): """Return overlapping byte-range components for native saved storages.""" records_by_storage_group = {} @@ -1972,10 +1931,6 @@ 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(), @@ -2001,11 +1956,6 @@ def pack_saved_tensor(tensor): with torch.no_grad(): target.copy_(tensor) tensor = target - - storage = tensor.untyped_storage() - per_callable_native_saved_storages[func_idx].setdefault( - storage._cdata, (storage, storage.data_ptr()) - ) return tensor with torch.autograd.graph.saved_tensors_hooks(pack_saved_tensor, lambda x: x): @@ -2049,329 +1999,6 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): # 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 = None - branch_pre_checkpoint = 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] - 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: - 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_pending_frees(): - """Return graph-pool blocks waiting for a cross-stream free event.""" - 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_pending_free" - ] - - def drain_slot_pool_pending_frees(): - """Poll completed cross-stream frees before restoring allocator state.""" - if not slot_pool_pending_frees(): - 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 = slot_pool_pending_frees() - 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() - 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): - record_storage(value) - return - visited.add(id(value)) - if isinstance(value, torch.Tensor): - if not value.is_cuda: - return - storage = value.untyped_storage() - record_storage(storage) - 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, - slot_saved_arenas, - slot_user_grad_arenas, - *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 checkpoint_slot_pool_boundary(phase, extra_values=(), save_state=True): - """Snapshot one allocator boundary and resolve all of its live storage owners.""" - state = ( - torch._C._cuda_getCheckpointState(torch.cuda.current_device(), mempool) - if save_state - else None - ) - blocks = slot_pool_active_blocks() - owners = checkpoint_live_storage_owners(blocks, phase, extra_values) - return _GraphPoolCheckpoint(state, blocks, owners) - - 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.") - 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." - ) - - 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() - 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, target, phase): - """Switch between two allocator boundaries while preserving their StorageImpls.""" - device = torch.cuda.current_device() - rollback_state = torch._C._cuda_getCheckpointState(device, mempool) - 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.owners, phase) - 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_release_started: - raise - try: - rollback_blocks = slot_pool_active_blocks() - rollback_owners = checkpoint_live_storage_owners( - rollback_blocks, - f"{phase} rollback", - (current.owners, target.owners), - ) - validate_checkpoint_live_storages(rollback_blocks, rollback_owners) - release_checkpoint_live_storages(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) per_callable_output_unflatten_spec = [None] * len(flatten_sample_args) @@ -2383,33 +2010,6 @@ def restore_slot_pool_boundary(current, target, phase): 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 = [] - branch_checkpoint_active = ( - branch_checkpoint is not None or branch_pre_checkpoint is not None - ) - 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: - gc.collect() - torch.cuda.synchronize() - drain_slot_pool_pending_frees() - branch_pre_checkpoint = checkpoint_slot_pool_boundary( - 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, - branch_pre_checkpoint, - 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( @@ -2422,7 +2022,6 @@ def restore_slot_pool_boundary(current, target, phase): 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] @@ -2435,21 +2034,9 @@ def restore_slot_pool_boundary(current, target, phase): with capture_saved_tensors(per_callable_fwd_idx, 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, func ) - if branch_group is not None: - record_replaced_io_storages( - original_flatten_outputs, - flatten_outputs, - branch_stale_storages, - ) - if branch_group[0] > 0: - branch_stale_storages.update( - per_callable_native_saved_storages[per_callable_fwd_idx] - ) - del original_flatten_outputs 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 @@ -2465,7 +2052,6 @@ def restore_slot_pool_boundary(current, target, phase): 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 @@ -2561,17 +2147,9 @@ def restore_slot_pool_boundary(current, target, phase): 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 @@ -2600,8 +2178,6 @@ def restore_slot_pool_boundary(current, target, phase): 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 static output and gradient objects after their capture lifetime. # Their backing storage remains alive either in the graph pool or an explicit @@ -2612,10 +2188,7 @@ def restore_slot_pool_boundary(current, target, phase): 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) + clear_native_io_target_rows((per_callable_bwd_idx,), clear_outputs=True) # Weak ref the static grad inputs of the previous backward pass within the # same chunk. @@ -2624,10 +2197,7 @@ def restore_slot_pool_boundary(current, target, phase): 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) + clear_native_io_target_rows((idx,), clear_grads=True) previous_per_callable_bwd_idx = per_callable_bwd_idx # Weak ref the static grad inputs of the previous chunk's last backward @@ -2641,70 +2211,12 @@ def restore_slot_pool_boundary(current, target, phase): 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) + clear_native_io_target_rows((idx,), clear_grads=True) 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: - branch_checkpoint = checkpoint_slot_pool_boundary( - 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, - branch_pre_checkpoint, - f"backward branch 1/{branch_group[1]} restore pre-boundary", - ) - else: - current_checkpoint = checkpoint_slot_pool_boundary( - 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.owners, - branch_checkpoint.owners, - ), - save_state=False, - ) - restore_slot_pool_boundary( - current_checkpoint, - branch_checkpoint, - f"{'forward' if c_id > 0 else 'backward'} branch " - f"{branch_group[0] + 1}/{branch_group[1]} restore post-boundary", - ) - 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 = None - branch_pre_checkpoint = None - gc.collect() - torch.cuda.synchronize() - drain_slot_pool_pending_frees() - else: # Capture forward graphs per_callable_static_outputs = [] @@ -2778,9 +2290,6 @@ def restore_slot_pool_boundary(current, target, phase): 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 is not None or branch_pre_checkpoint is not None: - raise RuntimeError("CUDA graph capture ended inside a slot checkpoint group.") - if allocator_settings_to_restore is not None: _allocator_settings_guard.restore() @@ -3183,8 +2692,9 @@ def make_graphed_callables( 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. When ``sample_args`` is a mutable list, entries whose leading input is - rebound to a staging surface are updated in place. + to the graph pool. Mutually exclusive variants must appear in ``_order`` as complete PP/VPP + schedules. When ``sample_args`` is a mutable list, entries whose leading input is rebound to + a staging surface are updated in place. pre_warmup_hook: callable, default = None A hook function that will be called before the warmup iterations. post_warmup_hook: callable, default = None From 8a9f58c2340deb86e24f9ee1d9ce0a95f97b0704 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 18:59:32 +0900 Subject: [PATCH 27/30] Avoid duplicate CUDA graph saved-tensor spill Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 33 +++++++++++++++++++++++++++++ transformer_engine/pytorch/graph.py | 16 ++++++++++---- 2 files changed, 45 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index d2c94cc8e3..4493bce0d4 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1955,6 +1955,39 @@ def forward(self, inp): reset_graphs(graphed) +def test_slot_memory_does_not_duplicate_output_backed_saves() -> None: + """A saved public output must not also consume spill space in its arena.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.sigmoid() + + elements = 4096 + module = Module().cuda() + samples = ((torch.ones(elements, device="cuda", requires_grad=True),),) + graphed = make_graphed_callables( + (module,), + samples, + 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),), + ) + + try: + arenas = graphed[0]._te_cuda_graph_saved_arenas + assert len(arenas) == 1 + assert next(iter(arenas.values())).numel() == elements * samples[0][0].element_size() + + inp = torch.randn(elements, device="cuda", requires_grad=True) + graphed[0](inp).sum().backward() + expected = inp.detach().sigmoid() + torch.testing.assert_close(inp.grad, expected * (1.0 - expected)) + finally: + reset_graphs(graphed) + + def test_slot_memory_honors_user_grad_liveness_groups() -> None: """The private plan may keep adjacent asynchronous gradient consumers disjoint.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index c8dbded65a..75eb457679 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -16,7 +16,6 @@ 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 @@ -1514,11 +1513,15 @@ 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) - def native_saved_storage_components(plan): + def native_saved_storage_components(plan, included_saved_indices=None): """Return overlapping byte-range components for native saved storages.""" records_by_storage_group = {} for saved_idx, spec in enumerate(plan): - if spec[0] != "native" or spec[7] == 0: + if ( + spec[0] != "native" + or spec[7] == 0 + or (included_saved_indices is not None and saved_idx not in included_saved_indices) + ): continue records_by_storage_group.setdefault(spec[8], []).append( (spec[9], spec[9] + spec[7], saved_idx) @@ -1598,8 +1601,13 @@ def plan_native_saved_alias_targets( raise RuntimeError(f"Native saved tensor {saved_idx} has no backing-storage plan.") packed_components = [] + unassigned_saved_indices = { + saved_idx + for saved_idx, spec in enumerate(plan) + if spec[0] == "native" and target_views[saved_idx] is None + } for component_start, component_end, component_records in native_saved_storage_components( - plan + plan, unassigned_saved_indices ): component_alignment = max( plan[saved_idx][4].itemsize for _, saved_idx in component_records From 3c9a9426ea3db08b5fc9a61b7243cf99686d309e Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 19:28:58 +0900 Subject: [PATCH 28/30] Simplify CUDA graph arena sizing Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 15 ++--- transformer_engine/pytorch/graph.py | 95 ++++++++++++----------------- 2 files changed, 45 insertions(+), 65 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 4493bce0d4..8ab8a00bac 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1613,8 +1613,8 @@ def forward(self, inp): 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.""" +def test_slot_memory_variants_reuse_native_saved_allocations() -> None: + """Additional variants must reuse native saved-tensor allocations, not grow the pool.""" class Module(torch.nn.Module): def __init__(self, canonical): @@ -1671,7 +1671,7 @@ def capture(variants): finally: reset_graphs(baseline) - graphed, forked_pool_bytes = capture(5) + graphed, variant_pool_bytes = capture(5) try: for graph_idx, graph in enumerate(graphed): inp = torch.randn(4096, device="cuda", requires_grad=True) @@ -1682,7 +1682,7 @@ def capture(variants): else 2.0 * inp.detach().sin() * inp.detach().cos() ) torch.testing.assert_close(inp.grad, expected) - assert forked_pool_bytes == baseline_pool_bytes + assert variant_pool_bytes == baseline_pool_bytes finally: reset_graphs(graphed) @@ -1727,7 +1727,7 @@ def forward(self, inp): def test_slot_memory_native_io_aliases_graph_pool_storage() -> None: - """DCP slot plans default to forked graph-pool I/O aliases.""" + """DCP variants share graph-pool I/O aliases within each physical slot.""" class Module(torch.nn.Module): def forward(self, inp): @@ -1744,10 +1744,7 @@ def forward(self, inp): samples, num_warmup_iters=2, allow_unused_input=True, - _order=[ - *(variant + 1 for variant in range(variants)), - *(-variant - 1 for variant in range(variants)), - ], + _order=_variant_major_order([1, -1], variants), _num_layers_per_chunk=[1] * variants, _reuse_graph_input_output_buffers=True, _graph_memory_slots=tuple( diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 75eb457679..b067901210 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -1540,10 +1540,20 @@ def native_saved_storage_components(plan, included_saved_indices=None): components.extend(group_components) return components + def native_saved_packed_components(plan, included_saved_indices=None): + """Return arena sizes and origins for native saved-storage components.""" + packed_components = [] + for component_start, component_end, component_records in native_saved_storage_components( + plan, included_saved_indices + ): + alignment = max(plan[saved_idx][4].itemsize for _, saved_idx in component_records) + origin = component_start // alignment * alignment + packed_components.append((component_end - origin, origin, component_records)) + return packed_components + def plan_native_saved_alias_targets( plan, canonical_targets, - measure_spill=False, protected_storage_ranges=None, preassigned_targets=None, ): @@ -1600,54 +1610,24 @@ def plan_native_saved_alias_targets( if storage_group is None or storage_offset_bytes is None: raise RuntimeError(f"Native saved tensor {saved_idx} has no backing-storage plan.") - packed_components = [] unassigned_saved_indices = { saved_idx for saved_idx, spec in enumerate(plan) if spec[0] == "native" and target_views[saved_idx] is None } - for component_start, component_end, component_records in native_saved_storage_components( - plan, unassigned_saved_indices - ): - 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, - ) - ) + packed_components = native_saved_packed_components(plan, unassigned_saved_indices) - 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) + banks = storage_banks 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 " @@ -1660,9 +1640,6 @@ def plan_native_saved_alias_targets( 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 @@ -1676,9 +1653,6 @@ def plan_native_saved_alias_targets( 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) @@ -1688,11 +1662,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(func_idx, outputs): - """Map boundary-backed saves onto the boundary address used at replay.""" + def semantic_boundary_alias_components(func_idx): + """Yield saved-storage components fully covered by one graph boundary.""" plan = per_callable_saved_tensor_plans[func_idx] aliases = per_callable_saved_tensor_boundary_aliases[func_idx] - targets = [None] * len(plan) for component_start, component_end, component_records in native_saved_storage_components( plan ): @@ -1727,6 +1700,15 @@ def semantic_boundary_alias_targets(func_idx, outputs): if not component_aliases: continue + yield component_records, component_aliases + + 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] + targets = [None] * len(plan) + for component_records, component_aliases in semantic_boundary_alias_components(func_idx): + component_saved_indices = [saved_idx for _, saved_idx in component_records] + anchor_storage = None anchor_shift = None for saved_idx, alias in component_aliases: @@ -1828,23 +1810,24 @@ def slot_tensor_targets(plan, arena=None): 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, + preassigned_saved_indices = { + saved_idx + for component_records, _ in semantic_boundary_alias_components(func_idx) + for _, saved_idx in component_records + } + unassigned_saved_indices = { + saved_idx + for saved_idx, spec in enumerate(plan) + if spec[0] == "native" and saved_idx not in preassigned_saved_indices + } + spill_bytes = sum( + _align_up(component_size) + for component_size, _, _ in native_saved_packed_components( + plan, unassigned_saved_indices + ) ) 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 with torch.cuda.use_mem_pool(slot_allocator_pool): slot_saved_arenas = { From 86ba0607add8710453ed75505f15a9866fc96010 Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Sun, 23 Aug 2026 21:34:07 +0900 Subject: [PATCH 29/30] Simplify CUDA graph slot arena placement Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 4 +- transformer_engine/pytorch/graph.py | 171 ++++++++++------------------ 2 files changed, 64 insertions(+), 111 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 8ab8a00bac..080bf4db6d 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1726,8 +1726,8 @@ def forward(self, inp): reset_graphs(graphed) -def test_slot_memory_native_io_aliases_graph_pool_storage() -> None: - """DCP variants share graph-pool I/O aliases within each physical slot.""" +def test_slot_memory_explicit_outputs_share_slot_address() -> None: + """DCP variants share their explicit output-arena address within a physical slot.""" class Module(torch.nn.Module): def forward(self, inp): diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index b067901210..c6be95279f 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -1290,13 +1290,13 @@ def match_staging_group(group_idx, seen_storages): static_input_surface[0] = input_target per_callable_static_input_surfaces[func_idx] = tuple(static_input_surface) - def native_io_family(func_idx): + def slot_io_family(func_idx): """Return the slot/liveness family shared by mutually exclusive CP branches.""" - arena_id, _ = slot_io_memory_alias_groups[func_idx] - return (arena_id, *slot_io_liveness_groups[func_idx]) + slot_id, _ = slot_io_memory_alias_groups[func_idx] + return (slot_id, *slot_io_liveness_groups[func_idx]) - def prepare_native_io_targets(per_callable_plans, kind): - """Validate same-slot CP branches and create their lazy alias targets.""" + def prepare_slot_io_targets(per_callable_plans, kind): + """Validate same-slot CP branches and reserve their explicit target rows.""" if per_callable_plans is None: return None @@ -1322,7 +1322,7 @@ def storage_alias_topology(plan): plans_by_family = {} for func_idx, plan in enumerate(per_callable_plans): _, branch_id = slot_io_memory_alias_groups[func_idx] - family = native_io_family(func_idx) + family = slot_io_family(func_idx) branch_plans = plans_by_family.setdefault(family, {}) if branch_id in branch_plans: raise RuntimeError( @@ -1366,60 +1366,58 @@ def storage_alias_topology(plan): return [[None] * len(plan) for plan in per_callable_plans] - per_callable_output_tensor_targets = prepare_native_io_targets( + per_callable_output_tensor_targets = prepare_slot_io_targets( per_callable_output_tensor_plans, "output" ) - per_callable_user_grad_tensor_targets = prepare_native_io_targets( + per_callable_user_grad_tensor_targets = prepare_slot_io_targets( per_callable_user_grad_tensor_plans, "user_grad" ) - native_io_family_sizes = {} + param_grad_family_sizes = {} if use_slot_memory: for func_idx in range(len(flatten_sample_args)): - family = native_io_family(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.""" - family = native_io_family(func_idx) + family = slot_io_family(func_idx) + param_grad_family_sizes[family] = param_grad_family_sizes.get(family, 0) + 1 + param_grad_anchors = {} + param_grad_capture_counts = {} + release_slot_io_targets = use_slot_memory + + def param_grad_alias_target(func_idx, tensor_idx, tensor, spec): + """Alias one CP branch onto the first branch's graph-pool parameter gradient.""" + family = slot_io_family(func_idx) key = (*family, tensor_idx) - anchors = native_io_anchors[kind] - counts = native_io_capture_counts[kind] - anchor = anchors.get(key) + anchor = param_grad_anchors.get(key) if anchor is None: target = tensor - anchors[key] = tensor + param_grad_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"CUDA graph parameter-gradient 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] + captured = param_grad_capture_counts.get(key, 0) + 1 + expected = param_grad_family_sizes[family] if captured > expected: raise RuntimeError( - f"CUDA graph native {kind} alias {key} captured {captured} of " + f"CUDA graph parameter-gradient alias {key} captured {captured} of " f"{expected} CP branches." ) if captured == expected: - anchors.pop(key) - counts.pop(key, None) + param_grad_anchors.pop(key) + param_grad_capture_counts.pop(key, None) else: - counts[key] = captured + param_grad_capture_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: + def clear_slot_io_target_rows(func_indices, clear_outputs=False, clear_grads=False): + """Drop explicit I/O views after the corresponding TE value dies.""" + if not release_slot_io_targets: return for func_idx in func_indices: if clear_outputs: @@ -1455,8 +1453,13 @@ def copy_outputs_to_slot_arena(func_idx, flatten_outputs, func): 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 spec[7] == 0: + copied_outputs.append(output) + continue + raise RuntimeError( + f"CUDA graph input {func_idx} has no slot-arena output target " + f"at position {tensor_idx}." + ) if target is not output: target.copy_(output) copied_outputs.append(target) @@ -1483,11 +1486,11 @@ def copy_user_grads_to_slot_arena(func_idx, static_input_surface, grad_inputs): 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" + if target is None and spec is not None and spec[7] != 0: + raise RuntimeError( + f"CUDA graph input {func_idx} has no slot-arena user-gradient target " + f"at position {input_idx}." ) - targets[input_idx] = target if target is not None: if target is not grad_input: with torch.no_grad(): @@ -1502,9 +1505,7 @@ def copy_user_grads_to_slot_arena(func_idx, static_input_surface, grad_inputs): ) 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" - ) + target = param_grad_alias_target(func_idx, input_idx, grad_input, spec) per_callable_param_grad_tensor_targets[func_idx][input_idx] = target if target is not grad_input: with torch.no_grad(): @@ -1553,48 +1554,17 @@ def native_saved_packed_components(plan, included_saved_indices=None): def plan_native_saved_alias_targets( plan, - canonical_targets, - protected_storage_ranges=None, + arena, + start_offset=0, preassigned_targets=None, ): - """Pack one CP branch's native saved tensors into canonical live storages.""" - protected_storage_ranges = protected_storage_ranges or {} + """Pack one CP branch's native saved tensors after its arena outputs.""" 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, - } - ) for saved_idx, spec in enumerate(plan): if spec[0] != "native": continue @@ -1617,41 +1587,30 @@ def plan_native_saved_alias_targets( } packed_components = native_saved_packed_components(plan, unassigned_saved_indices) - banks = storage_banks + arena_bytes = 0 if arena is None else arena.numel() * arena.element_size() + offset = start_offset 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): - offset = _align_up(bank["cursor"]) - if offset + component_size <= bank["capacity"]: - candidates.append( - (bank["capacity"] - offset - component_size, bank_idx, offset) - ) - if not candidates: + offset = _align_up(offset) + if arena is None or offset + component_size > arena_bytes: raise RuntimeError( - "CUDA graph CP branch native saved tensors do not fit in canonical " - "live storages: " + "CUDA graph CP branch native saved tensors do not fit in its slot arena: " 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)}." + f"offset={offset}, arena_bytes={arena_bytes}." ) - _, 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: spec = plan[saved_idx] - target_offset = ( - bank["base_offset"] + component_target_offset + source_offset - component_origin - ) + target_offset = 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 = _storage_view(bank["storage"], target_offset, spec) + target = _arena_view(arena, target_offset, spec) target.requires_grad_(spec[6]) target_views[saved_idx] = target + offset += component_size missing = [ saved_idx @@ -1847,15 +1806,11 @@ def slot_tensor_targets(plan, arena=None): ) 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, + arena, + start_offset=output_bytes, preassigned_targets=preassigned_targets, ) ) @@ -2179,7 +2134,7 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): per_callable_static_outputs[per_callable_bwd_idx] = make_weak_ref( static_outputs ) - clear_native_io_target_rows((per_callable_bwd_idx,), clear_outputs=True) + clear_slot_io_target_rows((per_callable_bwd_idx,), clear_outputs=True) # Weak ref the static grad inputs of the previous backward pass within the # same chunk. @@ -2188,7 +2143,7 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): per_callable_static_grad_inputs[idx] = make_weak_ref( per_callable_static_grad_inputs[idx] ) - clear_native_io_target_rows((idx,), clear_grads=True) + clear_slot_io_target_rows((idx,), clear_grads=True) previous_per_callable_bwd_idx = per_callable_bwd_idx # Weak ref the static grad inputs of the previous chunk's last backward @@ -2202,7 +2157,7 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): per_callable_static_grad_inputs[idx] = make_weak_ref( per_callable_static_grad_inputs[idx] ) - clear_native_io_target_rows((idx,), clear_grads=True) + clear_slot_io_target_rows((idx,), clear_grads=True) previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx del static_outputs if ceil(c_id) == c_id: @@ -2284,12 +2239,10 @@ def validate_captured_module_grads(func_idx, static_grad_inputs): if allocator_settings_to_restore is not None: _allocator_settings_guard.restore() - if use_slot_memory and ( - any(native_io_anchors.values()) or any(native_io_capture_counts.values()) - ): + if use_slot_memory and (param_grad_anchors or param_grad_capture_counts): raise RuntimeError( - "CUDA graph capture ended with incomplete native I/O aliases: " - f"anchors={native_io_anchors}, counts={native_io_capture_counts}." + "CUDA graph capture ended with incomplete parameter-gradient aliases: " + f"anchors={param_grad_anchors}, counts={param_grad_capture_counts}." ) # Now for every per_callable list, per_callable_*[i] holds the stuff for the ith callable. From e2fc9563472cc6dc004e05ba8a586ada0ab3eeca Mon Sep 17 00:00:00 2001 From: Tailai Ma Date: Mon, 24 Aug 2026 17:55:15 +0900 Subject: [PATCH 30/30] Move CP P2P transport override to a standalone change Signed-off-by: Tailai Ma --- tests/pytorch/test_cuda_graphs.py | 32 -------------- .../dot_product_attention/context_parallel.py | 42 ++++--------------- 2 files changed, 8 insertions(+), 66 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 080bf4db6d..98df47cb16 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -4,7 +4,6 @@ import contextlib import gc -import weakref from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest @@ -25,10 +24,6 @@ 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 import transformer_engine.pytorch.graph as te_graph from transformer_engine.common import recipe @@ -47,33 +42,6 @@ } -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 test_slot_memory_arena_view_uses_typed_storage_offset(): backing = torch.empty(16, dtype=torch.float32, device="cuda") arena = backing[4:] 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 b3b3145106..afaee671f0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -5,7 +5,6 @@ """Context Parallelism.""" import itertools import os -import weakref from typing import List, Tuple, Union import torch import transformer_engine_torch as tex @@ -55,27 +54,6 @@ _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" @@ -85,39 +63,35 @@ 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, transport_group + torch.distributed.isend, send_tensor, send_dst, cp_group ) recv_op = torch.distributed.P2POp( - torch.distributed.irecv, recv_tensor, recv_src, transport_group + torch.distributed.irecv, recv_tensor, recv_src, cp_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, transport_group + torch.distributed.irecv, recv_tensor, recv_src, cp_group ) send_op = torch.distributed.P2POp( - torch.distributed.isend, send_tensor, send_dst, transport_group + torch.distributed.isend, send_tensor, send_dst, cp_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, transport_group) - recv_op = torch.distributed.irecv(recv_tensor, recv_src, transport_group) + send_op = torch.distributed.isend(send_tensor, send_dst, cp_group) + recv_op = torch.distributed.irecv(recv_tensor, recv_src, cp_group) send_recv_ops.append(send_op) send_recv_ops.append(recv_op) else: - recv_op = torch.distributed.irecv(recv_tensor, recv_src, transport_group) - send_op = torch.distributed.isend(send_tensor, send_dst, transport_group) + recv_op = torch.distributed.irecv(recv_tensor, recv_src, cp_group) + send_op = torch.distributed.isend(send_tensor, send_dst, cp_group) send_recv_ops.append(recv_op) send_recv_ops.append(send_op) send_recv_reqs = send_recv_ops