diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index e6a83d92bc..693f4a0d84 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -4,6 +4,7 @@ import math import os +from contextlib import nullcontext from typing import Dict, List, Tuple, Optional import pytest @@ -955,6 +956,208 @@ def body(value): assert _FP8_RECOMPUTE_KEY in fp8_layer.fp8_meta +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +@pytest.mark.parametrize("training", all_boolean) +def test_checkpoint_with_eval_module_preserves_fp8_recompute_state(training, use_reentrant): + """An eval module participating in a checkpoint must stash like a training module.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + layer.train(training) + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + return layer(value) + + _checkpointed_linear_backward(body, use_reentrant, layer) + + assert _FP8_RECOMPUTE_KEY in layer.fp8_meta + assert "updated_scale_fwd" in layer.fp8_meta + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert len(recompute_buffer) == 1 + assert all(len(stashed) == 0 for stashed in recompute_buffer) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +@pytest.mark.parametrize("switch_to_eval", all_boolean) +def test_checkpoint_mode_change_between_phases_does_not_leak_recompute_stash( + switch_to_eval, use_reentrant +): + """The replay consumes the stash created by the original forward despite a mode change.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + observed = [] + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + layer.training, + ) + ) + return layer(value) + + stash_lengths = [] + for i in range(3): + layer.train() + inp = (torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) * (2.0**i)).requires_grad_() + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = te_checkpoint(body, inp, use_reentrant=use_reentrant).float().sum() + + if switch_to_eval: + layer.eval() + loss.backward() + torch.cuda.synchronize() + + assert torch.isfinite(loss) + assert inp.grad is not None and torch.isfinite(inp.grad).all() + assert layer.weight.grad is not None and torch.isfinite(layer.weight.grad).all() + assert _FP8_RECOMPUTE_KEY in layer.fp8_meta + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert len(recompute_buffer) == 1 + stash_lengths.append(len(recompute_buffer[layer.fp8_meta[_FP8_RECOMPUTE_KEY]])) + assert "updated_scale_fwd" in layer.fp8_meta + assert torch.equal(layer.fp8_meta["scaling_fwd"].scale, layer.fp8_meta["updated_scale_fwd"]) + assert observed[-2:] == [ + (True, False, True), + (True, True, not switch_to_eval), + ] + + assert stash_lengths == [0, 0, 0], f"Recompute stash leaked: {stash_lengths}" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_checkpoint_eval_intermediate_stashes_under_reentrant_no_grad(use_reentrant): + """An eval module with an intermediate input stashes even under reentrant `no_grad`.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + control = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + toggled = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda().eval() + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + return toggled(control(value)) + + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = te_checkpoint(body, inp, use_reentrant=use_reentrant).float().sum() + + toggled.train() + loss.backward() + torch.cuda.synchronize() + + assert inp.grad is not None and torch.isfinite(inp.grad).all() + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert len(recompute_buffer) == 2 + for layer in (control, toggled): + assert layer.weight.grad is not None and torch.isfinite(layer.weight.grad).all() + assert _FP8_RECOMPUTE_KEY in layer.fp8_meta + assert len(recompute_buffer[layer.fp8_meta[_FP8_RECOMPUTE_KEY]]) == 0 + assert "updated_scale_fwd" in layer.fp8_meta + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +@pytest.mark.parametrize("training", all_boolean) +def test_checkpoint_without_backward_does_not_accumulate_recompute_stashes(training, use_reentrant): + """Forwards with autograd disabled must not leave unreachable recompute state.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + layer.train(training) + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + + observed = [] + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + ) + ) + return layer(value) + + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + for _ in range(3): + out = te_checkpoint(body, inp, use_reentrant=use_reentrant) + assert torch.isfinite(out).all() + + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert all(len(stashed) == 0 for stashed in recompute_buffer) + assert observed == [(False, False)] * 3 + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +@pytest.mark.parametrize("enable_grad_in", ["context_fn", "function"]) +def test_checkpoint_outer_no_grad_preserves_explicit_inner_grad(enable_grad_in, use_reentrant): + """The no-grad bypass preserves an explicit nested request for autograd.""" + weight = torch.randn(16, 16, device="cuda", dtype=torch.float32) + input_data = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + + def run(checkpointed): + FP8GlobalStateManager.reset() + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda().eval() + with torch.no_grad(): + layer.weight.copy_(weight) + inp = input_data.clone().requires_grad_() + observed = [] + + def body(value): + grad_ctx = torch.enable_grad() if enable_grad_in == "function" else nullcontext() + with grad_ctx, autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + torch.is_grad_enabled(), + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + ) + ) + return layer(value) + + checkpoint_kwargs = {"use_reentrant": use_reentrant} + forward_ctx = nullcontext() + if enable_grad_in == "context_fn": + checkpoint_kwargs["context_fn"] = lambda: (torch.enable_grad(), nullcontext()) + if not checkpointed: + # Match the checkpoint forward context in the direct reference without + # changing grad state at the checkpoint call site itself. + forward_ctx = torch.enable_grad() + + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16), forward_ctx: + if checkpointed: + out = te_checkpoint(body, inp, **checkpoint_kwargs) + else: + out = body(inp) + + assert out.requires_grad + out.float().sum().backward() + torch.cuda.synchronize() + assert inp.grad is not None and torch.isfinite(inp.grad).all() + assert layer.weight.grad is not None and torch.isfinite(layer.weight.grad).all() + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert all(len(stashed) == 0 for stashed in recompute_buffer) + return out.detach(), inp.grad.detach(), layer.weight.grad.detach(), observed + + ref_out, ref_dgrad, ref_wgrad, ref_observed = run(checkpointed=False) + out, dgrad, wgrad, observed = run(checkpointed=True) + + torch.testing.assert_close(out, ref_out) + torch.testing.assert_close(dgrad, ref_dgrad) + torch.testing.assert_close(wgrad, ref_wgrad) + assert ref_observed == [(True, False, False)] + assert observed == ref_observed + + def _test_e2e_checkpointing_get_model(config, dtype): sigma = 0.023 init_method = init_method_normal(sigma) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 8605a4746b..595601f553 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -729,6 +729,7 @@ def checkpoint( context_fn = kwargs.pop("context_fn", noop_context_fn) determinism_check = kwargs.pop("determinism_check", "default") debug = kwargs.pop("debug", False) + if not has_te_modules(function): return torch.utils.checkpoint.checkpoint( function, @@ -740,6 +741,14 @@ def checkpoint( **kwargs, ) + # There will be no backward recompute when checkpoint is called with autograd + # disabled. Run the forward directly so FP8 modules do not save recompute state + # that can never be consumed. Preserve the user-provided forward context. + if not torch.is_grad_enabled(): + forward_ctx, _ = context_fn() + with forward_ctx: + return function(*args, **kwargs) + from .module.base import TransformerEngineBaseModule if isinstance(function, TransformerEngineBaseModule): diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index e9a65c3648..ad3883c54a 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1632,7 +1632,10 @@ def prepare_forward( FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta) # Activation recomputation is used and this is the first forward phase. - if self.training and is_fp8_activation_recompute_enabled(): + # Every delayed-scaling module in the first checkpoint phase must stash. + # Checkpoint phase, rather than module training mode, determines whether + # the matching recompute forward will need the original metadata. + if is_fp8_activation_recompute_enabled(): FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta) nvtx_range_push(self.__class__.__name__ + " forward")