diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 1b9e11792e..98df47cb16 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -2,6 +2,8 @@ # # See LICENSE for license information. +import contextlib +import gc from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest @@ -23,6 +25,7 @@ ) from transformer_engine.pytorch.quantization import FP8GlobalStateManager 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 +42,16 @@ } +def test_slot_memory_arena_view_uses_typed_storage_offset(): + backing = torch.empty(16, dtype=torch.float32, device="cuda") + arena = backing[4:] + spec = ("output", None, (1,), (1,), torch.float32, backing.device, False, 4) + + view = te_graph._arena_view(arena, 4, spec) + + assert view.data_ptr() == backing.data_ptr() + 5 * backing.element_size() + + def nvfp4_vanilla(): nvfp4_recipe = recipe.NVFP4BlockScaling() nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() @@ -740,3 +753,1328 @@ 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 _variant_major_order(base_order, variants, model_chunks=1): + """Repeat one complete PP/VPP schedule for each mutually exclusive variant.""" + return [ + (1 if chunk > 0 else -1) * (abs(chunk) + variant * model_chunks) + for variant in range(variants) + for chunk in base_order + ] + + +def test_graph_capture_contexts_restore_process_state_on_error(monkeypatch) -> None: + """Capture failures must restore GC and input gradients.""" + gc_was_enabled = gc.isenabled() + 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)] + + +@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( + "first_mode,second_mode", + (("independent", "offset8"), ("offset8", "independent"), ("offset4", "offset8")), +) +def test_slot_memory_rejects_branch_output_alias_mismatch(first_mode, second_mode) -> None: + """CP branches sharing an output slot must expose the same storage aliases.""" + + class Module(torch.nn.Module): + def __init__(self, mode): + super().__init__() + self.mode = mode + + def forward(self, inp): + output = inp.square() + first = output[:8] + if self.mode == "independent": + return first, output[8:].clone() + offset = 4 if self.mode == "offset4" else 8 + return first, output[offset : offset + 8] + + modules = (Module(first_mode).cuda(), Module(second_mode).cuda()) + samples = tuple((torch.ones(16, device="cuda", requires_grad=True),) for _ in modules) + graphed = None + try: + with pytest.raises(RuntimeError, match="incompatible output storage aliases"): + graphed = make_graphed_callables( + modules, + samples, + num_warmup_iters=2, + _order=[1, 2, -1, -2], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0), _slot(0, 1, 0, warmup=1)), + ) + assert not te_graph.is_graph_capturing() + finally: + if graphed is not None: + reset_graphs(graphed) + + +@pytest.mark.parametrize("state_kind", ("parameter", "buffer")) +def test_slot_memory_preserves_public_outputs_aliased_to_module_state(state_kind) -> None: + """Persistent module-state views remain outside the slot allocator pool.""" + + 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 graph capture starts.""" + + 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_rebinds_mutable_sample_args_to_staging_surfaces() -> None: + """Mutable sample banks expose the exact static inputs selected by slot staging.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.square() + + module = Module().cuda() + shared = torch.ones(4096, device="cuda", requires_grad=True) + samples = [(shared,), (shared,)] + graphed = make_graphed_callables( + (module,), + samples, + num_warmup_iters=2, + _order=[1, 1, -1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=( + _slot(0, 0, 0, warmup=0), + _slot(1, 1, 1, warmup=0), + ), + ) + + try: + assert any(tensor is shared for tensor in (samples[0][0], samples[1][0])) + assert samples[0][0].data_ptr() != samples[1][0].data_ptr() + + inp0 = torch.full((4096,), 2.0, device="cuda", requires_grad=True) + inp1 = torch.full((4096,), 3.0, device="cuda", requires_grad=True) + out0 = graphed[0](inp0) + out1 = graphed[1](inp1) + out0.sum().backward() + out1.sum().backward() + torch.testing.assert_close(inp0.grad, 2.0 * inp0.detach()) + torch.testing.assert_close(inp1.grad, 2.0 * inp1.detach()) + finally: + reset_graphs(graphed) + + +def test_slot_memory_rejects_warmup_input_alias_merge() -> None: + """Warmup-plan aliases must preserve the full input-storage alias topology.""" + + class Module(torch.nn.Module): + def forward(self, left, right): + return left.square() + right.square() + + module = Module().cuda() + left = torch.ones(4096, device="cuda", requires_grad=True) + right = torch.ones(4096, device="cuda", requires_grad=True) + merged = torch.ones(4096, device="cuda", requires_grad=True) + + with pytest.raises(RuntimeError, match="merge distinct input storages"): + make_graphed_callables( + (module,), + ((left, right), (merged, merged)), + num_warmup_iters=2, + _order=[1, 1, -1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=( + _slot(0, 0, 0, warmup=0), + _slot(1, 1, 1, warmup=0), + ), + ) + assert not te_graph.is_graph_capturing() + + +def test_slot_memory_staging_preserves_nonleading_input_aliases() -> None: + """Staging must not choose a leading-input address used by another argument.""" + + class Module(torch.nn.Module): + def forward(self, left, right): + return left.square() + right.square() + + module = Module().cuda() + branch0_left = torch.ones(4096, device="cuda", requires_grad=True) + branch0_right = torch.ones(4096, device="cuda", requires_grad=True) + branch1_left = torch.ones(4096, device="cuda", requires_grad=True) + samples = [ + (branch0_left, branch0_right), + (branch1_left, branch0_left), + ] + graphed = make_graphed_callables( + (module, module), + samples, + num_warmup_iters=2, + _order=[1, 2, -1, -2], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=( + _slot(0, 0, 0, warmup=0), + _slot(0, 1, 0, warmup=1), + ), + ) + + try: + assert samples[1][0].data_ptr() != samples[1][1].data_ptr() + left = torch.full((4096,), 2.0, device="cuda", requires_grad=True) + right = torch.full((4096,), 3.0, device="cuda", requires_grad=True) + output = graphed[1](left, right) + torch.testing.assert_close(output, left.detach().square() + right.detach().square()) + output.sum().backward() + torch.testing.assert_close(left.grad, 2.0 * left.detach()) + torch.testing.assert_close(right.grad, 2.0 * right.detach()) + finally: + reset_graphs(graphed) + + +@pytest.mark.parametrize("state_kind", ("parameter", "buffer")) +def test_slot_memory_rejects_user_inputs_aliased_to_module_state(state_kind) -> None: + """Runtime input staging must not overwrite parameters or buffers.""" + + class Module(torch.nn.Module): + def __init__(self): + super().__init__() + state = torch.ones(16, device="cuda") + if state_kind == "parameter": + self.state = torch.nn.Parameter(state) + else: + self.register_buffer("state", state) + + def forward(self, inp): + return 2.0 * inp + self.state + + module = Module() + sample = module.state.detach().view_as(module.state).requires_grad_(True) + with pytest.raises(RuntimeError, match="sharing storage with a module parameter or buffer"): + make_graphed_callables( + module, + (sample,), + num_warmup_iters=2, + _order=[1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0),), + ) + assert not te_graph.is_graph_capturing() + + +@pytest.mark.parametrize("state_kind", ("parameter", "buffer")) +def test_slot_memory_rejects_user_inputs_aliased_to_other_callable_state(state_kind) -> None: + """Staging candidates must not alias state owned by another graph callable.""" + + class StatelessModule(torch.nn.Module): + def forward(self, inp): + return inp.square() + + class StatefulModule(torch.nn.Module): + def __init__(self): + super().__init__() + state = torch.ones(16, device="cuda") + if state_kind == "parameter": + self.state = torch.nn.Parameter(state) + else: + self.register_buffer("state", state) + + def forward(self, inp): + return 2.0 * inp + self.state + + stateless = StatelessModule() + stateful = StatefulModule() + cross_callable_alias = stateful.state.detach().view_as(stateful.state).requires_grad_(True) + independent = torch.ones_like(stateful.state, requires_grad=True) + with pytest.raises(RuntimeError, match="sharing storage with a module parameter or buffer"): + make_graphed_callables( + (stateless, stateful), + ((cross_callable_alias,), (independent,)), + num_warmup_iters=2, + _order=[1, 2, -1, -2], + _num_layers_per_chunk=[1, 1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0), _slot(0, 1, 0, warmup=1)), + ) + assert not te_graph.is_graph_capturing() + + +def test_slot_memory_rejects_input_promoted_to_module_state_during_warmup() -> None: + """Lazy module state must be rechecked before input staging is selected.""" + + class LazyStateModule(torch.nn.Module): + def forward(self, inp): + if "state" not in self._buffers: + self.register_buffer("state", inp.detach()) + return 2.0 * inp + self.state + + module = LazyStateModule() + sample = torch.ones(16, device="cuda", requires_grad=True) + with pytest.raises(RuntimeError, match="module parameter or buffer.*after warmup"): + make_graphed_callables( + module, + (sample,), + num_warmup_iters=2, + _order=[1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0),), + ) + assert not te_graph.is_graph_capturing() + + +def test_slot_memory_rejects_non_tensor_sample_kwargs() -> None: + """Slot-memory static surfaces must reject non-Tensor kwargs explicitly.""" + + class Module(torch.nn.Module): + def forward(self, inp, use_bias=False): + output = inp.square() + if use_bias: + output = output + 1.0 + return output.sum().unsqueeze(0) + + module = Module().cuda() + sample = torch.ones(8, device="cuda", requires_grad=True) + with pytest.raises(TypeError, match="slot memory sample_kwargs must contain only Tensors"): + make_graphed_callables( + (module,), + ((sample,),), + sample_kwargs=({"use_bias": True},), + num_warmup_iters=1, + _order=[1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0),), + ) + assert not te_graph.is_graph_capturing() + + +def test_slot_memory_snapshots_shared_kwarg_across_alternate_liveness() -> None: + """A later forward must not overwrite a shared kwarg needed by backward.""" + + 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_reuses_variant_major_branches(reverse_replay) -> None: + """Complete CP schedules must reuse one slot backing across variants.""" + + class Module(torch.nn.Module): + def forward(self, inp): + 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_major_order([1, -1], 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( + "VARIANT_MAJOR_GRAD_MISMATCH", + { + "reverse_replay": reverse_replay, + "variant": variant, + "input_ptr": inp.data_ptr(), + "output_ptr": output.data_ptr(), + "grad_ptr": inp.grad.data_ptr(), + "grad_head": inp.grad[:4].tolist(), + "expected_head": expected[:4].tolist(), + "grad_head_i64": inp.grad[:4].view(torch.int64).tolist(), + }, + flush=True, + ) + torch.testing.assert_close( + inp.grad, + expected, + msg=lambda message: f"variant={variant}: {message}", + ) + finally: + reset_graphs(graphed) + + +def test_slot_memory_reclaims_retained_variant_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_major_order([1, -1], 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_tracks_native_saved_storage() -> None: + """Variant-major capture must preserve native 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_major_order([1, -1], 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_variants_reuse_native_saved_allocations() -> None: + """Additional variants must reuse native saved-tensor allocations, not grow the pool.""" + + class Module(torch.nn.Module): + def __init__(self, canonical): + 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) + ) + graphed = make_graphed_callables( + modules, + samples, + num_warmup_iters=2, + _order=_variant_major_order([1, 1, -1, -1], variants), + _num_layers_per_chunk=[1] * variants, + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=tuple( + _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, variant_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 variant_pool_bytes == baseline_pool_bytes + finally: + reset_graphs(graphed) + + +def test_slot_memory_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_major_order([1, -1], 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_explicit_outputs_share_slot_address() -> None: + """DCP variants share their explicit output-arena address within a physical slot.""" + + class Module(torch.nn.Module): + def forward(self, inp): + 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_major_order([1, -1], 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 repeated per CP variant. + base_order = [1, 1, 2, 2, 1, -2, 1, -2, 2, -1, 2, -1, -2, -2, -1, -1] + order = _variant_major_order(base_order, variants, model_chunks) + + 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_does_not_duplicate_output_backed_saves() -> None: + """A saved public output must not also consume spill space in its arena.""" + + class Module(torch.nn.Module): + def forward(self, inp): + return inp.sigmoid() + + elements = 4096 + module = Module().cuda() + samples = ((torch.ones(elements, device="cuda", requires_grad=True),),) + graphed = make_graphed_callables( + (module,), + samples, + num_warmup_iters=2, + _order=[1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + _graph_memory_slots=(_slot(0, 0, 0),), + ) + + try: + arenas = graphed[0]._te_cuda_graph_saved_arenas + assert len(arenas) == 1 + assert next(iter(arenas.values())).numel() == elements * samples[0][0].element_size() + + inp = torch.randn(elements, device="cuda", requires_grad=True) + graphed[0](inp).sum().backward() + expected = inp.detach().sigmoid() + torch.testing.assert_close(inp.grad, expected * (1.0 - expected)) + finally: + reset_graphs(graphed) + + +def test_slot_memory_honors_user_grad_liveness_groups() -> None: + """The private plan may keep adjacent asynchronous gradient consumers disjoint.""" + + 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..afaee671f0 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,9 @@ # See LICENSE for license information. """Context Parallelism.""" -import os import itertools -from typing import List, Union, Tuple +import os +from typing import List, Tuple, Union import torch import transformer_engine_torch as tex @@ -54,7 +54,6 @@ _seq_chunk_ids_cache_for_reordering_before_attn = {} _seq_chunk_ids_cache_for_reordering_after_attn = {} _softmax_offset_chunk_ids_cache = {} - # 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,7 +63,6 @@ def flash_attn_p2p_communicate( ): """Point-to-point communications of KV and dKV in Attention with context parallelism""" send_recv_ops = [] - if batch_p2p_comm: if rank % 2 == 0: send_op = torch.distributed.P2POp( diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 86b8a4acf4..c6be95279f 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -3,12 +3,14 @@ # See LICENSE for license information. """Functions for CUDA Graphs support in FP8""" + from collections.abc import Iterable import contextlib import gc +import os import warnings from math import ceil -from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union import torch from torch.utils._pytree import tree_flatten as _tree_flatten @@ -23,7 +25,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 +40,148 @@ 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 _storage_view(storage, offset: int, spec: Tuple[Any, ...]) -> torch.Tensor: + """Materialize a typed tensor view at a byte offset in a CUDA storage.""" + target = torch.empty((0,), dtype=spec[4], device=spec[5]) + if offset % target.element_size(): + raise RuntimeError("CUDA graph arena view has an unaligned byte offset.") + return target.set_( + storage, + offset // target.element_size(), + spec[2], + spec[3], + ) + + +def _arena_view(arena: torch.Tensor, offset: int, spec: Tuple[Any, ...]) -> torch.Tensor: + """Materialize a typed tensor view at a byte offset in an arena.""" + offset += arena.storage_offset() * arena.element_size() + return _storage_view(arena.untyped_storage(), offset, spec) + + def set_capture_start() -> None: """Record beginning of `make_graphed_callables`.""" global _IS_GRAPH_CAPTURING @@ -69,12 +213,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 +236,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 +269,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 +414,8 @@ 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 + _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 +425,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): @@ -358,6 +556,8 @@ def _make_graphed_callables( "In the beta API, sample_args " + "for each callable must contain only Tensors. Other types are not allowed." ) + if use_slot_memory and not all(isinstance(arg, torch.Tensor) for arg in flatten_kwarg): + raise TypeError("CUDA graph slot memory sample_kwargs must contain only Tensors.") # If a callable is an nn.Module, its graph's full input surface is the args the user explicitly # passes to forward (ie, its sample_args) AND the module's parameter attributes. @@ -367,6 +567,7 @@ def _make_graphed_callables( # torch.cuda.make_graphed_callables. per_callable_len_user_args = [len(args) for args in flatten_sample_args] if _order is None: + per_callable_funcs = callables per_callable_module_params = [ tuple(c.parameters()) if isinstance(c, torch.nn.Module) else () for c in callables ] @@ -374,17 +575,15 @@ def _make_graphed_callables( flatten_sample_args[i] + per_callable_module_params[i] for i in range(len(callables)) ] else: + per_callable_funcs = [] per_callable_module_params = [] for m_chunk in range(num_model_chunks): for _ in range(num_microbatches): for l_no in range(_num_layers_per_chunk[m_chunk]): + func = callables[_prefix_num_layers[m_chunk] + l_no] + per_callable_funcs.append(func) per_callable_module_params.append( - tuple(callables[_prefix_num_layers[m_chunk] + l_no].parameters()) - if isinstance( - callables[_prefix_num_layers[m_chunk] + l_no], - torch.nn.Module, - ) - else () + tuple(func.parameters()) if isinstance(func, torch.nn.Module) else () ) if len(per_callable_module_params) != len(flatten_sample_args): raise ValueError( @@ -397,6 +596,34 @@ def _make_graphed_callables( for i in range(len(flatten_sample_args)) ] + def validate_slot_user_input_state_aliases(phase): + """Reject user inputs backed by any persistent module-state storage.""" + module_state_storages = set() + inspected_modules = set() + for func in per_callable_funcs: + if not isinstance(func, torch.nn.Module) or id(func) in inspected_modules: + continue + inspected_modules.add(id(func)) + for tensor in (*func.parameters(), *func.buffers()): + if tensor.is_cuda and tensor.untyped_storage().nbytes() > 0: + module_state_storages.add(_tensor_storage_ptr(tensor)) + + for func_idx, user_args in enumerate(flatten_sample_args): + for arg_idx, arg in enumerate(user_args): + if ( + arg.is_cuda + and arg.untyped_storage().nbytes() > 0 + and _tensor_storage_ptr(arg) in module_state_storages + ): + raise RuntimeError( + "CUDA graph slot memory does not support user input tensor " + f"{arg_idx} sharing storage with a module parameter or buffer in the " + f"capture bank for graph input {func_idx} {phase}." + ) + + if use_slot_memory: + validate_slot_user_input_state_aliases("before warmup") + fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] bwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] bwd_dw_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] @@ -410,7 +637,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 +706,264 @@ 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 = {} + source_storage_by_target = {} + 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}." + ) + previous_source = source_storage_by_target.setdefault(target_ptr, source_ptr) + if previous_source != source_ptr: + raise RuntimeError( + f"Warmup-plan alias {target_idx} would merge distinct input storages " + f"{previous_source} and {source_ptr} into {target_ptr}." + ) + + source_plan = per_callable_saved_tensor_plans[template_idx] + if source_plan is None: + 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 +971,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 +1004,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 +1061,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 +1093,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 +1123,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,10 +1154,793 @@ 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: + validate_slot_user_input_state_aliases("after warmup") + per_callable_param_grad_tensor_targets = [ + [None] * len(static_input_surface) + for static_input_surface in per_callable_static_input_surfaces + ] + + 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 staging_target_preserves_aliases(group_idx, storage_id): + """Check that rebinding the leading input keeps every input-storage alias.""" + for func_idx in staging_groups[group_idx]["members"]: + static_input_surface = per_callable_static_input_surfaces[func_idx] + old_storage_id = _tensor_storage_ptr(static_input_surface[0]) + for other_input in static_input_surface[1:]: + other_storage_id = _tensor_storage_ptr(other_input) + if (old_storage_id == other_storage_id) != (storage_id == other_storage_id): + return False + return True + + def match_staging_group(group_idx, seen_storages): + for storage_id, tensor in staging_groups[group_idx]["candidates"].items(): + if storage_id in seen_storages or not staging_target_preserves_aliases( + group_idx, storage_id + ): + 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: ( + staging_target_preserves_aliases(index, None), + len(staging_groups[index]["candidates"]), + ), + ): + match_staging_group(group_idx, set()) + + for group_idx, group in enumerate(staging_groups): + input_target = staging_targets.get(group_idx) + if input_target is None: + if not staging_target_preserves_aliases(group_idx, None): + raise RuntimeError( + "CUDA graph slot staging cannot preserve a leading-input storage alias " + f"for group {group_idx}." + ) + source = next(iter(group["candidates"].values())) + signature = _saved_tensor_signature(source) + storage_numel = signature[-1] // source.element_size() + 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 slot_io_family(func_idx): + """Return the slot/liveness family shared by mutually exclusive CP branches.""" + slot_id, _ = slot_io_memory_alias_groups[func_idx] + return (slot_id, *slot_io_liveness_groups[func_idx]) + + def prepare_slot_io_targets(per_callable_plans, kind): + """Validate same-slot CP branches and reserve their explicit target rows.""" + if per_callable_plans is None: + return None + + def storage_alias_topology(plan): + indexed_specs = [ + (tensor_idx, spec) + for tensor_idx, spec in enumerate(plan) + if spec is not None and spec[0] == kind + ] + if any(len(spec) < 10 for _, spec in indexed_specs): + return None + return tuple( + ( + left_idx, + right_idx, + left[8] == right[8], + right[9] - left[9] if left[8] == right[8] else None, + ) + for position, (left_idx, left) in enumerate(indexed_specs) + for right_idx, right in indexed_specs[position + 1 :] + ) + + plans_by_family = {} + for func_idx, plan in enumerate(per_callable_plans): + _, branch_id = slot_io_memory_alias_groups[func_idx] + family = slot_io_family(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." + ) + alias_topologies = {storage_alias_topology(plan) for plan in plans} + if len(alias_topologies) != 1: + raise RuntimeError( + f"CUDA graph {kind} family {family} has incompatible {kind} storage aliases." + ) + for tensor_idx, specs in enumerate(zip(*plans)): + if all(spec is None for spec in specs): + continue + 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_slot_io_targets( + per_callable_output_tensor_plans, "output" + ) + per_callable_user_grad_tensor_targets = prepare_slot_io_targets( + per_callable_user_grad_tensor_plans, "user_grad" + ) + + param_grad_family_sizes = {} + if use_slot_memory: + for func_idx in range(len(flatten_sample_args)): + family = slot_io_family(func_idx) + param_grad_family_sizes[family] = param_grad_family_sizes.get(family, 0) + 1 + param_grad_anchors = {} + param_grad_capture_counts = {} + release_slot_io_targets = use_slot_memory + + def param_grad_alias_target(func_idx, tensor_idx, tensor, spec): + """Alias one CP branch onto the first branch's graph-pool parameter gradient.""" + family = slot_io_family(func_idx) + key = (*family, tensor_idx) + anchor = param_grad_anchors.get(key) + if anchor is None: + target = tensor + param_grad_anchors[key] = tensor + else: + available_bytes = anchor.untyped_storage().nbytes() - ( + anchor.storage_offset() * anchor.element_size() + ) + if spec[7] > available_bytes: + raise RuntimeError( + f"CUDA graph parameter-gradient alias {key} needs {spec[7]} bytes, " + f"but its first CP branch exposes only {available_bytes} bytes." + ) + target = _arena_view(anchor, 0, spec) + + captured = param_grad_capture_counts.get(key, 0) + 1 + expected = param_grad_family_sizes[family] + if captured > expected: + raise RuntimeError( + f"CUDA graph parameter-gradient alias {key} captured {captured} of " + f"{expected} CP branches." + ) + if captured == expected: + param_grad_anchors.pop(key) + param_grad_capture_counts.pop(key, None) + else: + param_grad_capture_counts[key] = captured + return target + + def clear_slot_io_target_rows(func_indices, clear_outputs=False, clear_grads=False): + """Drop explicit I/O views after the corresponding TE value dies.""" + if not release_slot_io_targets: + return + for func_idx in func_indices: + if clear_outputs: + 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 + if spec[7] == 0: + copied_outputs.append(output) + continue + raise RuntimeError( + f"CUDA graph input {func_idx} has no slot-arena output target " + f"at position {tensor_idx}." + ) + if target is not output: + target.copy_(output) + copied_outputs.append(target) + 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 and spec[7] != 0: + raise RuntimeError( + f"CUDA graph input {func_idx} has no slot-arena user-gradient target " + f"at position {input_idx}." + ) + 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 = param_grad_alias_target(func_idx, input_idx, grad_input, spec) + per_callable_param_grad_tensor_targets[func_idx][input_idx] = target + if target is not grad_input: + with torch.no_grad(): + target.copy_(grad_input) + grad_input = target + copied_grad_inputs.append(grad_input) + return tuple(copied_grad_inputs) + + def native_saved_storage_components(plan, included_saved_indices=None): + """Return overlapping byte-range components for native saved storages.""" + records_by_storage_group = {} + for saved_idx, spec in enumerate(plan): + if ( + spec[0] != "native" + or spec[7] == 0 + or (included_saved_indices is not None and saved_idx not in included_saved_indices) + ): + continue + records_by_storage_group.setdefault(spec[8], []).append( + (spec[9], spec[9] + spec[7], saved_idx) + ) + + components = [] + for records in records_by_storage_group.values(): + records.sort() + group_components = [] + for start, end, saved_idx in records: + if not group_components or start >= group_components[-1][1]: + group_components.append([start, end, [(start, saved_idx)]]) + else: + group_components[-1][1] = max(group_components[-1][1], end) + group_components[-1][2].append((start, saved_idx)) + components.extend(group_components) + return components + + def native_saved_packed_components(plan, included_saved_indices=None): + """Return arena sizes and origins for native saved-storage components.""" + packed_components = [] + for component_start, component_end, component_records in native_saved_storage_components( + plan, included_saved_indices + ): + alignment = max(plan[saved_idx][4].itemsize for _, saved_idx in component_records) + origin = component_start // alignment * alignment + packed_components.append((component_end - origin, origin, component_records)) + return packed_components + + def plan_native_saved_alias_targets( + plan, + arena, + start_offset=0, + preassigned_targets=None, + ): + """Pack one CP branch's native saved tensors after its arena outputs.""" + if preassigned_targets is None: + target_views = [None] * len(plan) + else: + if len(preassigned_targets) != len(plan): + raise RuntimeError("Native saved preassignment does not match its plan.") + target_views = list(preassigned_targets) + 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.") + + unassigned_saved_indices = { + saved_idx + for saved_idx, spec in enumerate(plan) + if spec[0] == "native" and target_views[saved_idx] is None + } + packed_components = native_saved_packed_components(plan, unassigned_saved_indices) + + arena_bytes = 0 if arena is None else arena.numel() * arena.element_size() + offset = start_offset + for component_size, component_origin, component_records in sorted( + packed_components, key=lambda item: item[0], reverse=True + ): + offset = _align_up(offset) + if arena is None or offset + component_size > arena_bytes: + raise RuntimeError( + "CUDA graph CP branch native saved tensors do not fit in its slot arena: " + f"component_bytes={component_size}, " + f"offset={offset}, arena_bytes={arena_bytes}." + ) + for source_offset, saved_idx in component_records: + spec = plan[saved_idx] + target_offset = offset + source_offset - component_origin + itemsize = spec[4].itemsize + if target_offset % itemsize: + raise RuntimeError( + f"Native saved tensor {saved_idx} has an unaligned canonical offset." + ) + target = _arena_view(arena, target_offset, spec) + target.requires_grad_(spec[6]) + target_views[saved_idx] = target + offset += component_size + + 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_components(func_idx): + """Yield saved-storage components fully covered by one graph boundary.""" + plan = per_callable_saved_tensor_plans[func_idx] + aliases = per_callable_saved_tensor_boundary_aliases[func_idx] + for component_start, component_end, component_records in native_saved_storage_components( + plan + ): + component_saved_indices = [saved_idx for _, saved_idx in component_records] + candidate_aliases = [ + (saved_idx, aliases[saved_idx]) + for saved_idx in component_saved_indices + 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 + + yield component_records, component_aliases + + def semantic_boundary_alias_targets(func_idx, outputs): + """Map boundary-backed saves onto the boundary address used at replay.""" + plan = per_callable_saved_tensor_plans[func_idx] + targets = [None] * len(plan) + for component_records, component_aliases in semantic_boundary_alias_components(func_idx): + component_saved_indices = [saved_idx for _, saved_idx in component_records] + + anchor_storage = None + anchor_shift = None + for saved_idx, alias in component_aliases: + _, 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 = _storage_view(anchor_storage, target_offset, spec) + 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) + preassigned_saved_indices = { + saved_idx + for component_records, _ in semantic_boundary_alias_components(func_idx) + for _, saved_idx in component_records + } + unassigned_saved_indices = { + saved_idx + for saved_idx, spec in enumerate(plan) + if spec[0] == "native" and saved_idx not in preassigned_saved_indices + } + spill_bytes = sum( + _align_up(component_size) + for component_size, _, _ in native_saved_packed_components( + plan, unassigned_saved_indices + ) + ) + arena_id = saved_tensor_arena_ids[func_idx] + arena_sizes[arena_id] = max(arena_sizes.get(arena_id, 0), output_bytes + spill_bytes) + + 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) + per_callable_slot_saved_targets.append( + plan_native_saved_alias_targets( + plan, + arena, + start_offset=output_bytes, + 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]}." + ) + + 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 + 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. @@ -585,12 +1971,24 @@ def hook_fn( args = sample_args[per_callable_fwd_idx] kwargs = sample_kwargs[per_callable_fwd_idx] fwd_graph = fwd_graphs[per_callable_fwd_idx] + 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, saved_alias_targets): + outputs = func(*args, **kwargs) + flatten_outputs, spec = _tree_flatten(outputs) + flatten_outputs = copy_outputs_to_slot_arena( + per_callable_fwd_idx, flatten_outputs, func + ) 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] @@ -673,13 +2071,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 +2093,48 @@ def hook_fn( retain_graph=retain_graph_in_backward, ) grad_inputs = tuple(input.grad for input in inputs) + grad_inputs = copy_user_grads_to_slot_arena( + per_callable_bwd_idx, static_input_surface, 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 - # 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 ) + clear_slot_io_target_rows((per_callable_bwd_idx,), clear_outputs=True) # Weak ref the static grad inputs of the previous backward pass within the # same chunk. @@ -730,6 +2143,7 @@ def hook_fn( per_callable_static_grad_inputs[idx] = make_weak_ref( per_callable_static_grad_inputs[idx] ) + clear_slot_io_target_rows((idx,), clear_grads=True) previous_per_callable_bwd_idx = per_callable_bwd_idx # Weak ref the static grad inputs of the previous chunk's last backward @@ -743,9 +2157,12 @@ def hook_fn( per_callable_static_grad_inputs[idx] = make_weak_ref( per_callable_static_grad_inputs[idx] ) + clear_slot_io_target_rows((idx,), clear_grads=True) previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx + del static_outputs if ceil(c_id) == c_id: bwd_idx[m_chunk] += 1 + else: # Capture forward graphs per_callable_static_outputs = [] @@ -764,7 +2181,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 +2235,16 @@ 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 allocator_settings_to_restore is not None: + _allocator_settings_guard.restore() + + if use_slot_memory and (param_grad_anchors or param_grad_capture_counts): + raise RuntimeError( + "CUDA graph capture ended with incomplete parameter-gradient aliases: " + f"anchors={param_grad_anchors}, counts={param_grad_capture_counts}." + ) + # Now for every per_callable list, per_callable_*[i] holds the stuff for the ith callable. def make_graphed_autograd_function( @@ -830,7 +2263,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 +2504,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 +2586,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 +2627,18 @@ 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. Mutually exclusive variants must appear in ``_order`` as complete PP/VPP + schedules. When ``sample_args`` is a mutable list, entries whose leading input is rebound to + a staging surface are updated in place. pre_warmup_hook: callable, default = None A hook function that will be called before the warmup iterations. post_warmup_hook: callable, default = None @@ -1292,8 +2748,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 +2771,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 +2799,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) + 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: - if just_one_callable: - forward_funcs = forward_funcs[0] - else: - forward_funcs = tuple(forward_funcs) - - # Save RNG state. - if graph_safe_rng_available(): - generators = [ - torch.cuda.default_generators[torch.cuda.current_device()], - *get_all_rng_states().values(), - ] - original_rng_states = [state.get_state() for state in generators] - else: - original_rng_states = torch.cuda.get_rng_state() - - graphed_callables = _make_graphed_callables( - forward_funcs, - sample_args, - num_warmup_iters=num_warmup_iters, - allow_unused_input=allow_unused_input, - cache_quantized_params=cache_quantized_params, - sample_kwargs=sample_kwargs, - _order=_order, - _num_layers_per_chunk=_num_layers_per_chunk, - pool=pool, - retain_graph_in_backward=retain_graph_in_backward, - _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers, - 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