[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,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)

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

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