diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 5a848dc0e8..8f85f57f32 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -5,6 +5,8 @@ from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest import copy +import gc +import weakref import torch from transformer_engine.pytorch import ( @@ -994,6 +996,206 @@ def hook(module: torch.nn.Module) -> None: ] +def test_ordered_warmup_releases_consumed_outputs() -> None: + """Ordered warmup should only retain outputs until their corresponding backward.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.previous_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + is_warmup = not torch.cuda.is_current_stream_capturing() + if is_warmup and self.previous_output is not None: + assert self.previous_output() is None + output = input_ * 2 + if is_warmup: + self.previous_output = weakref.ref(output) + return output + + module = OutputLifetimeModule() + sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in range(2)) + graphed_callables = make_graphed_callables( + (module,), + sample_args, + num_warmup_iters=2, + _order=[1, -1, 1, -1], + _num_layers_per_chunk=[1], + ) + assert module.previous_output is not None + assert module.previous_output() is None + reset_graphs(graphed_callables) + + +def test_unordered_warmup_releases_consumed_outputs() -> None: + """Unordered warmup should release each output after its corresponding backward.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self, output_refs: list, module_idx: int) -> None: + super().__init__() + self.output_refs = output_refs + self.module_idx = module_idx + self.capture_started = False + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.capture_started = True + else: + self.output_refs[self.module_idx] = weakref.ref(output) + return output + + output_refs = [None, None] + modules = tuple(OutputLifetimeModule(output_refs, module_idx) for module_idx in range(2)) + + def first_module_backward_pre_hook(_module: torch.nn.Module) -> None: + if not modules[0].capture_started: + assert output_refs[1] is not None + assert output_refs[1]() is None + + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in modules), + num_warmup_iters=2, + capture_time_hooks=[ + {"backward_pre_hooks": {0: first_module_backward_pre_hook}}, + None, + ], + ) + assert all(output_ref is not None and output_ref() is None for output_ref in output_refs) + reset_graphs(graphed_callables) + + +def test_inference_warmup_does_not_retain_outputs() -> None: + """Inference warmup should release outputs as soon as each forward returns.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self, previous_output: list) -> None: + super().__init__() + self.previous_output = previous_output + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + is_warmup = not torch.cuda.is_current_stream_capturing() + if is_warmup and self.previous_output[0] is not None: + assert self.previous_output[0]() is None + output = input_ * 2 + if is_warmup: + self.previous_output[0] = weakref.ref(output) + return output + + previous_output = [None] + modules = tuple(OutputLifetimeModule(previous_output).eval() for _ in range(2)) + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, 8, device="cuda"),) for _ in modules), + num_warmup_iters=2, + ) + assert previous_output[0] is not None + assert previous_output[0]() is None + reset_graphs(graphed_callables) + + +def test_reused_capture_buffers_release_outputs_after_backward() -> None: + """Capture locals must not keep weak-refed output buffers alive.""" + + class OutputLifetimeModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.previous_capture_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + if ( + torch.cuda.is_current_stream_capturing() + and self.previous_capture_output is not None + ): + assert self.previous_capture_output() is None + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.previous_capture_output = weakref.ref(output) + return output + + module = OutputLifetimeModule() + sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in range(2)) + graphed_callables = make_graphed_callables( + (module,), + sample_args, + _order=[1, -1, 1, -1], + _num_layers_per_chunk=[1], + _reuse_graph_input_output_buffers=True, + ) + assert module.previous_capture_output is not None + assert module.previous_capture_output() is None + reset_graphs(graphed_callables) + + +def test_reset_releases_only_the_selected_callable() -> None: + """Reset releases one callable's graph state without retaining its peers.""" + + class CaptureOutputModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.capture_output = None + + def forward(self, input_: torch.Tensor) -> torch.Tensor: + output = input_ * 2 + if torch.cuda.is_current_stream_capturing(): + self.capture_output = weakref.ref(output) + return output + + modules = tuple(CaptureOutputModule().cuda() for _ in range(2)) + graphed_callables = make_graphed_callables( + modules, + tuple((torch.ones(4, device="cuda", requires_grad=True),) for _ in modules), + ) + capture_outputs = tuple(module.capture_output for module in modules) + assert all(output is not None and output() is not None for output in capture_outputs) + + graphed_callables[0].reset() + graphed_callables[0].reset() + gc.collect() + assert capture_outputs[0]() is None + assert capture_outputs[1]() is not None + + output = graphed_callables[1](torch.randn(4, device="cuda", requires_grad=True)) + output.sum().backward() + del output + graphed_callables[1].reset() + gc.collect() + assert capture_outputs[1]() is None + + +@pytest.mark.parametrize("with_order", (False, True)) +def test_reset_rejects_all_replay_entry_points(with_order: bool) -> None: + """Reset is idempotent and terminal for forward and backward replay.""" + + class TestModule(torch.nn.Module): + def forward(self, input_: torch.Tensor) -> torch.Tensor: + return input_ * 2 + + module = TestModule().cuda() + sample_input = torch.ones(4, device="cuda", requires_grad=True) + graph_options = {} + if with_order: + graph_options = {"_order": [1, -1], "_num_layers_per_chunk": [1]} + graphed_callable = make_graphed_callables(module, (sample_input,), **graph_options) + output = graphed_callable(torch.randn_like(sample_input, requires_grad=True)) + torch.cuda.synchronize() + + graphed_callable.reset() + graphed_callable.reset() + if not with_order: + # The eager fallback for a different training state is invalid after reset too. + graphed_callable.eval() + + error = "has been reset and can no longer be used" + with pytest.raises(RuntimeError, match=error): + graphed_callable(torch.randn_like(sample_input, requires_grad=True)) + with pytest.raises(RuntimeError, match=error): + graphed_callable.backward_dw() + with pytest.raises(RuntimeError, match=error): + output.sum().backward() + + @pytest.mark.parametrize("with_order", (False, True)) def test_make_graphed_callables_with_capture_time_hooks(with_order: bool) -> None: """Test capture-time hooks around warmup and graph capture.""" diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index b298b3d8ff..a5fb392ddc 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -633,13 +633,17 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): warmup_outputs = [] for func_idx, func in zip(warmup_func_idx, warmup_func): outputs = _run_warmup_forward(func_idx, func, func_idx) - warmup_outputs.append((func_idx, func, outputs)) - if is_training: - for func_idx, func, outputs in reversed(warmup_outputs): - _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx) + if is_training: + warmup_outputs.append((func_idx, func, outputs)) + else: + del outputs + while warmup_outputs: + func_idx, func, outputs = warmup_outputs.pop() + _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx) + del outputs else: # Follow _order exactly, mirroring the capture phase. - per_fwd_outputs = {} # per_callable_fwd_idx -> flattened outputs + per_fwd_outputs = {} # per_callable_fwd_idx -> outstanding flattened outputs fwd_idx = [0] * num_model_chunks bwd_idx = [0] * num_model_chunks for c_id in _order: @@ -653,7 +657,10 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): ) + (fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) func = callables[callable_idx] outputs = _run_warmup_forward(per_callable_fwd_idx, func, callable_idx) - per_fwd_outputs[per_callable_fwd_idx] = outputs + if is_training: + per_fwd_outputs[per_callable_fwd_idx] = outputs + else: + del outputs fwd_idx[m_chunk] += 1 elif ceil(c_id) == c_id: # Backward pass for chunk -c_id. @@ -665,10 +672,11 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): _prefix_num_layers[m_chunk] * num_microbatches ) + (bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no) func = callables[callable_idx] - outputs = per_fwd_outputs[per_callable_bwd_idx] + outputs = per_fwd_outputs.pop(per_callable_bwd_idx) _run_warmup_backward( per_callable_bwd_idx, func, outputs, warmup_iter, callable_idx ) + del outputs bwd_idx[m_chunk] += 1 if post_warmup_hook is not None: @@ -729,6 +737,7 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): 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 + del outputs, flatten_outputs fwd_idx[m_chunk] += 1 else: # Capture backward graph for model chunk c_id, microbatch bwd_idx[-c_id-1] @@ -917,6 +926,11 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx): per_callable_static_grad_inputs[idx] ) previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx + + # The per-callable containers now own all tensors that must survive + # capture. Drop local strong references so weak-refed graph buffers can + # be returned to the shared CUDA graph pool before the next capture. + del static_outputs, static_grad_inputs, grad_inputs if ceil(c_id) == c_id: bwd_idx[m_chunk] += 1 else: @@ -1028,12 +1042,22 @@ def make_graphed_autograd_function( static_grad_inputs, returned_param_grad_clone_slots, ): + is_reset = False + + def ensure_not_reset(): + """Reject replay after this callable's graph state has been released.""" + if is_reset: + raise RuntimeError( + "This graphed callable has been reset and can no longer be used." + ) + 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): # pylint: disable=missing-function-docstring + ensure_not_reset() # Set flag for whether to update FP8 weight updates ctx.is_first_module = FP8GlobalStateManager.is_first_fp8_module() @@ -1071,6 +1095,7 @@ def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *i @torch.autograd.function.once_differentiable def backward(ctx, *grads): # pylint: disable=missing-function-docstring + ensure_not_reset() # Replay backward graph if len(grads) != len(static_grad_outputs): @@ -1119,6 +1144,7 @@ def backward(ctx, *grads): return (None, None, None) + tuple(grad_inputs) def functionalized(*user_args, **user_kwargs): + ensure_not_reset() # Decide whether to update FP8 weights skip_fp8_weight_update = None @@ -1170,16 +1196,39 @@ def functionalized(*user_args, **user_kwargs): ) return _tree_unflatten(out, output_unflatten_spec) - return functionalized - - def make_graphed_attribute_functions(graph_idx): - # Get te modules for current graph + def release_static_state(): + """Release per-callable state captured by replay closures.""" + nonlocal fwd_graph, bwd_graph, is_reset + nonlocal module_params + nonlocal static_input_surface, static_outputs + nonlocal static_grad_outputs, static_grad_inputs + + is_reset = True + + # Drop the per-callable references that can own graph-pool storage. + fwd_graph = None + bwd_graph = None + module_params = () + static_input_surface = () + static_outputs = () + static_grad_outputs = () + static_grad_inputs = () + + return functionalized, release_static_state, ensure_not_reset + + def make_graphed_attribute_functions(graph_idx, release_static_state, ensure_not_reset): + # Snapshot per-callable state so returned closures do not retain the outer lists. + fwd_graph = fwd_graphs[graph_idx] + bwd_graph = bwd_graphs[graph_idx] + bwd_dw_graph = bwd_dw_graphs[graph_idx] + need_bwd_dw = need_bwd_dw_graph.get(graph_idx, False) te_modules = visited_te_modules.get(graph_idx, set()) # Attach backward_dw as an attribute to the graphed callable. def backward_dw(): - if need_bwd_dw_graph.get(graph_idx, False): - bwd_dw_graphs[graph_idx].replay() + ensure_not_reset() + if need_bwd_dw: + bwd_dw_graph.replay() # Trigger the grad accumulation hook for wgrad graphs. for module in te_modules: @@ -1191,16 +1240,24 @@ def backward_dw(): # Attach reset as an attribute to the graphed callable. def reset(): - fwd_graphs[graph_idx].reset() - bwd_graphs[graph_idx].reset() - bwd_dw_graphs[graph_idx].reset() + nonlocal fwd_graph, bwd_graph, bwd_dw_graph, te_modules + + for graph in (fwd_graph, bwd_graph, bwd_dw_graph): + if graph is not None: + graph.reset() + + fwd_graph = None + bwd_graph = None + bwd_dw_graph = None + te_modules = () + release_static_state() return backward_dw, reset # Put together the final graphed callables ret = [] for i in range(len(sample_args)): - graphed = make_graphed_autograd_function( + graphed, release_static_state, ensure_not_reset = make_graphed_autograd_function( fwd_graphs[i], bwd_graphs[i], per_callable_module_params[i], @@ -1218,8 +1275,17 @@ def reset(): te_modules = visited_te_modules.get(i, set()) if isinstance(func, torch.nn.Module): - def make_graphed_forward(func, graph_training_state, graphed, orig_fwd, te_modules): + def make_graphed_forward( + func, + graph_training_state, + graphed, + orig_fwd, + te_modules, + ensure_not_reset, + ): def new_fwd(*user_args, **user_kwargs): + ensure_not_reset() + # If the module's training-or-eval state matches what we graphed, # run the graph, otherwise run the original forward method if func.training == graph_training_state: @@ -1264,7 +1330,14 @@ def new_fwd(*user_args, **user_kwargs): return new_fwd - forward = make_graphed_forward(func, func.training, graphed, func.forward, te_modules) + forward = make_graphed_forward( + func, + func.training, + graphed, + func.forward, + te_modules, + ensure_not_reset, + ) if _order is None: func.forward = forward ret.append(func) @@ -1273,7 +1346,11 @@ def new_fwd(*user_args, **user_kwargs): else: ret.append(graphed) - backward_dw_func, reset_func = make_graphed_attribute_functions(i) + backward_dw_func, reset_func = make_graphed_attribute_functions( + i, + release_static_state, + ensure_not_reset, + ) setattr(ret[-1], "backward_dw", backward_dw_func) setattr(ret[-1], "reset", reset_func)