diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 46aab13075..5d51b936bb 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -141,6 +141,61 @@ def reset_global_fp8_state(): FP8GlobalStateManager.reset() +def test_warmup_forward_failure_cleanup_and_retry(): + """A failed warmup must leave the module reusable for capture and training.""" + from transformer_engine.pytorch import graph as te_graph + + class FailDuringWarmup(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(32, 32, dtype=torch.bfloat16, device="cuda") + self.dropout = torch.nn.Dropout(0.1) + self.fail = True + + def forward(self, inp): + output = self.dropout(self.linear(inp)) + if self.fail: + raise RuntimeError("controlled warmup forward failure") + return output + + model = FailDuringWarmup() + inp = torch.randn(4, 32, device="cuda", dtype=torch.bfloat16) + original_call = type(model).__call__ + hooks = [(module, dict(module._forward_hooks)) for module in model.modules()] + cpu_rng = torch.get_rng_state() + cuda_rng = torch.cuda.get_rng_state() + graphed = None + try: + with pytest.raises(RuntimeError, match="controlled warmup forward failure"): + make_graphed_callables(model, (inp,), num_warmup_iters=1) + for module, original_hooks in hooks: + assert module._forward_hooks == original_hooks + assert type(model).__call__ is original_call + assert not te_graph.is_graph_capturing() + torch.testing.assert_close(torch.get_rng_state(), cpu_rng, rtol=0, atol=0) + torch.testing.assert_close(torch.cuda.get_rng_state(), cuda_rng, rtol=0, atol=0) + + model.fail = False + reference = copy.deepcopy(model) + reference_input = inp.clone() + rng = torch.cuda.get_rng_state() + expected = reference(reference_input) + expected.sum().backward() + expected_parameter_grads = [p.grad.clone() for p in reference.parameters()] + torch.cuda.set_rng_state(rng) + graphed = make_graphed_callables(model, (inp,), num_warmup_iters=1) + actual = graphed(inp) + actual.sum().backward() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + for parameter, expected_grad in zip(model.parameters(), expected_parameter_grads): + torch.testing.assert_close(parameter.grad, expected_grad, rtol=0, atol=0) + finally: + if graphed is not None: + graphed.reset() + type(model).__call__ = original_call + te_graph.set_capture_end() + + def assert_all_equal(l1: List[torch.Tensor], l2: List[torch.Tensor], names=None) -> bool: """Check that two lists of tensors match exactly.""" assert len(l1) == len(l2), "Unequal number of outputs." diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 5362c7cc27..8a7ff292e5 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Functions for CUDA Graphs support in FP8""" + from collections.abc import Iterable import contextlib import gc @@ -576,11 +577,13 @@ def hook_fn(module, inputs, outputs, func_idx=func_idx): # pylint: disable=unus _run_capture_time_hooks(capture_time_hooks, callable_idx, "forward_pre_hooks", func) hooks = [] - for module in func.modules(): - hooks.append(module.register_forward_hook(hook_fn)) - outputs = func(*args, **kwargs) - for hook in hooks: - hook.remove() + try: + for module in func.modules(): + hooks.append(module.register_forward_hook(hook_fn)) + outputs = func(*args, **kwargs) + finally: + for hook in hooks: + hook.remove() _run_capture_time_hooks(capture_time_hooks, callable_idx, "forward_hooks", func) @@ -1479,6 +1482,10 @@ def make_graphed_callables( `original PyTorch implementation `_ for more documentation. + If a callable raises during warmup, temporary discovery hooks, saved RNG state, + module call wrappers, and the capture flag are restored before the exception is + re-raised. This cleanup does not make fatal CUDA errors recoverable. + .. warning:: Arguments 'fp8_enabled', 'fp8_calibrating', 'fp8_recipe', 'fp8_group', and 'fp8_weight_caching' are deprecated. @@ -1717,34 +1724,30 @@ def call_func(self, *args, **kwargs): ] original_rng_states = [state.get_state() for state in generators] - 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, - clone_param_grads_on_return=clone_param_grads_on_return, - pre_warmup_hook=pre_warmup_hook, - post_warmup_hook=post_warmup_hook, - capture_time_hooks=capture_time_hooks, - ) - - # Ensures warmup does not affect numerics for ops such as dropout. - for gen, state in zip(generators, original_rng_states): - gen.set_state(state) - - # 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) - - set_capture_end() + try: + 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, + clone_param_grads_on_return=clone_param_grads_on_return, + pre_warmup_hook=pre_warmup_hook, + post_warmup_hook=post_warmup_hook, + capture_time_hooks=capture_time_hooks, + ) + finally: + # Restore capture-owned state even if warmup raises. + for gen, state in zip(generators, original_rng_states): + gen.set_state(state) + for module_cls, old_call in old_call_funcs.items(): + module_cls.__call__ = old_call + restore_fp8_tensors(modules, saved_fp8_tensors) + set_capture_end() return graphed_callables