diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 1b9e11792e..9511fbfc8e 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -2,6 +2,9 @@ # # See LICENSE for license information. +import contextlib +import gc +import weakref from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest @@ -22,7 +25,12 @@ 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 from utils import ModelConfig, reset_rng_states @@ -39,6 +47,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 +775,1164 @@ 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, user_grad=None): + """Build one private graph-memory slot used by the focused tests below.""" + 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: + """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 = [] + warmup_hooks = [] + + 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["pre_warmup_hook"]() + assert warmup_hooks == ["pre"] + 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),), + 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: + """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),), + ) + + +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.""" + + 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) + + +@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.""" + + 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.""" + + 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) + + +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.""" + + 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) + + +@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: + """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: + """DCP slot plans 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_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_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): + 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, user_grad=0), + _slot(1, 1, 0, overlap=1, warmup=1, user_grad=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, 1) + 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..b3b3145106 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,26 @@ _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 +85,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..1a86caef57 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -6,6 +6,8 @@ #include "pybind.h" +#include +#include #include #include #include @@ -135,6 +137,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..573abd70b9 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -3,17 +3,21 @@ # See LICENSE for license information. """Functions for CUDA Graphs support in FP8""" + +from bisect import bisect_right 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 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 @@ -23,7 +27,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 +42,140 @@ 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: + """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 + self._settings_to_restore = restore + setter(settings) + + def restore(self) -> None: + """Restore the allocator settings saved by :meth:`apply`.""" + 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() + + +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 ( + 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()) + ): + return None + 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]) + 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 @@ -69,12 +207,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 @@ -90,10 +230,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( @@ -108,6 +263,8 @@ 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, + _allocator_settings_guard: Optional[_AllocatorSettingsGuard] = None, pre_warmup_hook: Optional[Callable] = None, post_warmup_hook: Optional[Callable] = None, ) -> SingleOrTuple[Callable]: @@ -251,10 +408,36 @@ 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: + 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: raise ValueError( @@ -264,6 +447,43 @@ def _make_graphed_callables( raise RuntimeError( "`_reuse_graph_input_output_buffers` is only available in training mode." ) + + saved_tensor_arena_ids = None + 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( + "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) != 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 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 + # 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): @@ -410,7 +630,42 @@ def _make_graphed_callables( 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 + 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()] + 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 +699,257 @@ 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 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." + ) + 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(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 + 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 + + 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 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] + 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()): @@ -454,6 +957,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 @@ -483,13 +990,52 @@ 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) - outputs, _ = _tree_flatten(func(*args, **kwargs)) - for hook in hooks: - hook.remove() + 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 + ): + 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, + 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, + observed_output_plan, + "Output", + ) + else: + outputs, _ = _tree_flatten(func(*args, **kwargs)) if is_training: inputs = tuple(i for i in static_input_surface if i.requires_grad) with _none_grad_context_wrapper(inputs): @@ -501,6 +1047,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 +1079,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 +1109,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 +1140,1168 @@ 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: + 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): + 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_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: + 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(_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 + # 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() + 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(), + 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}." + ) + 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( + 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, func): + """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." + ) + 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 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) + 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 = [{} for _ in 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(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): + 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 + # 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 + + # 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[func_idx][boundary_idx] + else: + 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." + ) + + 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={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={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 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: + 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, + ) + 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 = { + 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 + } + + per_callable_slot_saved_targets = [] + for func_idx, plan in enumerate(per_callable_saved_tensor_plans): + 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 + ) + 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) + 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 = { + arena_id: torch.empty( + (required_bytes,), dtype=torch.uint8, device=torch.cuda.current_device() + ) + 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): + 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 + 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 + 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]}." + ) + + 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(), + 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 + + 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): + 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." + ) + + 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_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 + 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_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 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() + 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 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_blocks, + current_owners, + target_state, + target_blocks, + target_owners, + 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) @@ -570,6 +2313,47 @@ 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 = [] + 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, + ) + ) + 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_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 +2366,43 @@ 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 = ( + per_callable_slot_saved_targets[per_callable_fwd_idx] + if use_slot_memory + else None + ) 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, 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 + 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 +2412,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 +2486,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 +2508,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 +2571,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 +2588,87 @@ 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: + 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, + 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_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_storage_owners = None + branch_pre_checkpoint_state = None + branch_pre_checkpoint_live_blocks = None + branch_pre_checkpoint_storage_owners = None + gc.collect() + torch.cuda.synchronize() + drain_slot_pool_pending_frees() + else: # Capture forward graphs per_callable_static_outputs = [] @@ -764,7 +2687,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 +2741,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: + _allocator_settings_guard.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 +2774,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 +3015,10 @@ 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) + 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] @@ -1143,6 +3097,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 +3138,16 @@ 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-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, 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. 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 @@ -1292,8 +3257,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): @@ -1317,8 +3280,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 = {} @@ -1344,57 +3308,94 @@ 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) + 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: - # 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, - 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) - - # Remove FP8 wrapper. - for module_cls, old_call in old_call_funcs.items(): - module_cls.__call__ = old_call - - # Restore FP8 state. - restore_fp8_tensors(modules, saved_fp8_tensors) + def run_pre_warmup_hook(): + nonlocal warmup_cleanup_pending + if pre_warmup_hook is not None: + pre_warmup_hook() + warmup_cleanup_pending = True + + 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 + 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) + + 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)] + + 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=guarded_pre_warmup_hook, + post_warmup_hook=guarded_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) + if guarded_post_warmup_hook is not None: + capture_cleanup.callback(guarded_post_warmup_hook) - set_capture_end() return graphed_callables