diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index e6a83d92bc..2d630959a3 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,98 @@ 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("recompute_training", all_boolean) +def test_checkpoint_eval_module_balances_fp8_recompute_state(recompute_training, use_reentrant): + """An eval module must stash metadata for the recompute forward.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda().eval() + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + # Exercise an intermediate input whose grad state differs between the + # reentrant and non-reentrant checkpoint forward implementations. + return layer(value * 2) + + 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() + + layer.train(recompute_training) + loss.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() + assert _FP8_RECOMPUTE_KEY in layer.fp8_meta + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert len(recompute_buffer) == 1 + assert len(recompute_buffer[layer.fp8_meta[_FP8_RECOMPUTE_KEY]]) == 0 + assert "updated_scale_fwd" in layer.fp8_meta + assert torch.equal(layer.fp8_meta["scaling_fwd"].scale, layer.fp8_meta["updated_scale_fwd"]) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_checkpoint_without_autograd_does_not_accumulate_recompute_stashes(): + """Checkpoint calls without autograd must not leave unreachable metadata.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + 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) + assert torch.isfinite(out).all() + + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert _FP8_RECOMPUTE_KEY not in layer.fp8_meta + 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) +def test_checkpoint_without_autograd_preserves_forward_context(): + """The direct path must preserve a context that explicitly enables gradients.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda().eval() + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + return layer(value) + + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + out = te_checkpoint( + body, + inp, + context_fn=lambda: (torch.enable_grad(), nullcontext()), + ) + + 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 _FP8_RECOMPUTE_KEY not in layer.fp8_meta + assert all(len(stashed) == 0 for stashed in recompute_buffer) + + 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..4e397405a2 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -740,6 +740,13 @@ def checkpoint( **kwargs, ) + # When checkpoint is entered with autograd disabled, run the forward directly + # to avoid unreachable FP8 recompute state. Preserve the 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..71db9f04bf 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1632,7 +1632,7 @@ 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(): + if is_fp8_activation_recompute_enabled(): FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta) nvtx_range_push(self.__class__.__name__ + " forward")