[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
 
@@ -906,6 +907,20 @@ def _checkpointed_linear_backward(body, use_reentrant, *layers):
         assert torch.isfinite(layer.weight.grad).all()
 
 
+def _assert_fp8_recompute_state_drained(*layers):
+    """Check that each module consumed its stash and restored its live scale."""
+    recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer
+    assert len(recompute_buffer) == len(layers)
+    for layer in layers:
+        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
+        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)
 @pytest.mark.parametrize("use_reentrant", all_boolean)
 def test_checkpoint_inner_autocast_is_an_fp8_recompute_region(use_reentrant):
@@ -955,6 +970,176 @@ 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_state_drained(layer)
+
+
+@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_reentrant_checkpoint_without_grad_input_does_not_stash_fp8_metadata():
+    """A reentrant checkpoint that cannot receive backward must not save recompute state."""
+    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(),
+                )
+            )
+            return layer(value)
+
+    inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16)
+    with torch.enable_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+        for _ in range(3):
+            out = te_checkpoint(body, inp, use_reentrant=True)
+            assert torch.isfinite(out).all()
+            assert not out.requires_grad
+
+    assert _FP8_RECOMPUTE_KEY not in layer.fp8_meta
+    assert not FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer
+    assert observed == [(False, False)] * 3
+
+    valid_inp = torch.randn(
+        16,
+        16,
+        device="cuda",
+        dtype=torch.bfloat16,
+        requires_grad=True,
+    )
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        loss = te_checkpoint(body, valid_inp, use_reentrant=True).float().sum()
+    loss.backward()
+    torch.cuda.synchronize()
+
+    assert valid_inp.grad is not None and torch.isfinite(valid_inp.grad).all()
+    assert layer.weight.grad is not None and torch.isfinite(layer.weight.grad).all()
+    assert observed[-2:] == [(True, False), (True, True)]
+    _assert_fp8_recompute_state_drained(layer)
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_ineligible_reentrant_checkpoint_preserves_outer_recompute_context():
+    """An inner checkpoint without grad inputs must inherit its outer recompute phase."""
+    FP8GlobalStateManager.reset()
+    fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+    inner = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+    tail = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+    observed = []
+
+    def inner_body(value):
+        observed.append(
+            (
+                is_fp8_activation_recompute_enabled(),
+                in_fp8_activation_recompute_phase(),
+            )
+        )
+        return inner(value)
+
+    def outer_body(value):
+        with autocast(enabled=True, recipe=fp8_recipe):
+            inner_out = te_checkpoint(inner_body, value.detach(), use_reentrant=True)
+            return tail(value + inner_out)
+
+    inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        loss = te_checkpoint(outer_body, inp, use_reentrant=False).float().sum()
+    loss.backward()
+    torch.cuda.synchronize()
+
+    assert observed == [(True, False), (True, True)]
+    assert inp.grad is not None and torch.isfinite(inp.grad).all()
+    assert inner.weight.grad is None
+    assert tail.weight.grad is not None and torch.isfinite(tail.weight.grad).all()
+    _assert_fp8_recompute_state_drained(inner, tail)
+
+
+@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

@@ -372,9 +372,14 @@ def forward(
         # Preserve torch autocast context for the backward pass
         torch_gpu_amp_ctx, torch_cpu_amp_ctx = _get_active_autocast_contexts()
 
-        with torch.no_grad(), forward_ctx:
-            with activation_recompute_forward(activation_recompute=True, recompute_phase=False):
-                outputs = run_function(*args, **kwargs)
+        # An ineligible nested checkpoint must inherit an enclosing recompute phase.
+        activation_recompute_ctx = (
+            activation_recompute_forward(activation_recompute=True, recompute_phase=False)
+            if any(ctx.needs_input_grad)
+            else nullcontext()
+        )
+        with torch.no_grad(), forward_ctx, activation_recompute_ctx:
+            outputs = run_function(*args, **kwargs)
 
         # Divide hidden states across model parallel group and only keep
         # the chunk corresponding to the current rank.
@@ -740,6 +745,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")