[PyTorch] Pair delayed-scaling FP8 recompute metadata per module

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)

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):

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")