Fix duplicate FP8 state updates during activation checkpoint recomputation

tests/pytorch/test_reentrant_fp8_updates.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Regression tests for FP8 update ownership during activation recompute."""
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+from transformer_engine.common import recipe
+from transformer_engine.pytorch.distributed import (
+    in_fp8_activation_recompute_phase,
+    is_fp8_activation_recompute_enabled,
+)
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
+
+
+fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize(
+    ("checkpoint_mode", "segments", "num_layers"),
+    (
+        ("none", "single", 3),
+        ("non-reentrant", "single", 3),
+        ("non-reentrant", "per-layer", 3),
+        ("reentrant", "single", 1),
+        ("reentrant", "single", 3),
+        ("reentrant", "per-layer", 3),
+        ("nested-reentrant", "nested", 3),
+    ),
+)
+def test_delayed_scaling_updates_once_per_autocast(
+    monkeypatch, checkpoint_mode, segments, num_layers
+):
+    """Activation recompute must not advance global FP8 state per module/segment."""
+
+    FP8GlobalStateManager.reset()
+    counts = {"forward": 0, "backward": 0}
+    original_update = FP8GlobalStateManager.reduce_and_update_fp8_tensors
+
+    def counted_update(_cls, forward=True):
+        counts["forward" if forward else "backward"] += 1
+        return original_update(forward=forward)
+
+    monkeypatch.setattr(
+        FP8GlobalStateManager,
+        "reduce_and_update_fp8_tensors",
+        classmethod(counted_update),
+    )
+
+    torch.manual_seed(20260715)
+    torch.cuda.manual_seed_all(20260715)
+    layers = [
+        te.Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() for _ in range(num_layers)
+    ]
+    network = torch.nn.Sequential(*layers)
+    inp = torch.randn(
+        16,
+        16,
+        device="cuda",
+        dtype=torch.bfloat16,
+        requires_grad=True,
+    )
+    fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
+        enabled=True,
+        recipe=fp8_recipe,
+    ):
+        if checkpoint_mode == "none":
+            out = network(inp)
+        elif checkpoint_mode == "nested-reentrant":
+
+            def inner(x):
+                return layers[0](x)
+
+            def outer(x):
+                x = te.checkpoint(inner, x, use_reentrant=True)
+                for layer in layers[1:]:
+                    x = layer(x)
+                return x
+
+            out = te.checkpoint(outer, inp, use_reentrant=True)
+        elif segments == "single":
+            out = te.checkpoint(
+                network,
+                inp,
+                use_reentrant=checkpoint_mode == "reentrant",
+            )
+        else:
+            out = inp
+            for layer in layers:
+                out = te.checkpoint(
+                    layer,
+                    out,
+                    use_reentrant=checkpoint_mode == "reentrant",
+                )
+        loss = out.float().sum()
+
+    loss.backward()
+    torch.cuda.synchronize()
+
+    assert torch.isfinite(loss)
+    assert inp.grad is not None
+    assert torch.isfinite(inp.grad).all()
+    assert inp.grad.abs().max() > 0
+    for layer in layers:
+        assert layer.weight.grad is not None
+        assert torch.isfinite(layer.weight.grad).all()
+        assert layer.weight.grad.abs().max() > 0
+    assert counts == {"forward": 1, "backward": 1}
+    assert FP8GlobalStateManager.quantization_state.autocast_depth == 0
+    assert not FP8GlobalStateManager.is_fp8_enabled()
+    assert not is_fp8_activation_recompute_enabled()
+    assert not in_fp8_activation_recompute_phase()
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_reentrant_checkpoint_gradients_match_uncheckpointed():
+    """Reentrant recompute should preserve the uncheckpointed FP8 numerics."""
+
+    def run(checkpoint):
+        FP8GlobalStateManager.reset()
+        torch.manual_seed(20260715)
+        torch.cuda.manual_seed_all(20260715)
+        layers = [
+            te.Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() for _ in range(3)
+        ]
+        network = torch.nn.Sequential(*layers)
+        inp = torch.randn(
+            16,
+            16,
+            device="cuda",
+            dtype=torch.bfloat16,
+            requires_grad=True,
+        )
+        fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+        with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
+            enabled=True,
+            recipe=fp8_recipe,
+        ):
+            out = te.checkpoint(network, inp, use_reentrant=True) if checkpoint else network(inp)
+            loss = out.float().sum()
+        loss.backward()
+        torch.cuda.synchronize()
+        return (
+            loss.detach(),
+            inp.grad.detach(),
+            *(layer.weight.grad.detach() for layer in layers),
+        )
+
+    reference = run(checkpoint=False)
+    checkpointed = run(checkpoint=True)
+    for actual, expected in zip(checkpointed, reference):
+        torch.testing.assert_close(actual, expected, rtol=0.05, atol=0.01)

transformer_engine/pytorch/distributed.py

@@ -238,6 +238,13 @@ def gather_split_1d_tensor(tensor: torch.Tensor, tp_group: dist_group_type) -> t
     return gathered
 
 
+@dataclass
+class _ActivationRecomputeState:
+    """FP8 state shared by a checkpoint's original forward and recompute."""
+
+    is_first_fp8_module: Optional[bool] = None
+
+
 class activation_recompute_forward(AbstractContextManager, ContextDecorator):
     """Context manager used to control the forward runtime behavior when executed
     under the `CheckpointFunction` function. For running FP8, the forward pass will
@@ -247,30 +254,50 @@ class activation_recompute_forward(AbstractContextManager, ContextDecorator):
     activations, followed by calculation of gradients using these values.
     """
 
-    _is_first_fp8_module: List = []
-
-    def __init__(self, activation_recompute: bool = False, recompute_phase: bool = False):
+    def __init__(
+        self,
+        activation_recompute: bool = False,
+        recompute_phase: bool = False,
+        state: Optional[_ActivationRecomputeState] = None,
+        reserve_first_fp8_module: bool = False,
+    ):
         super().__init__()
         self.activation_recompute = activation_recompute
         self.recompute_phase = recompute_phase
+        self.state = _ActivationRecomputeState() if state is None else state
+        self.reserve_first_fp8_module = reserve_first_fp8_module
 
     def __enter__(self):
         global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
-        _FP8_ACTIVATION_RECOMPUTE_ENABLED = (
-            self.activation_recompute and FP8GlobalStateManager.is_fp8_enabled()
-        )
+        self._previous_recompute_enabled = _FP8_ACTIVATION_RECOMPUTE_ENABLED
+        self._previous_recompute_phase = _FP8_ACTIVATION_RECOMPUTE_PHASE
+        fp8_enabled = FP8GlobalStateManager.is_fp8_enabled()
+        _FP8_ACTIVATION_RECOMPUTE_ENABLED = self.activation_recompute and fp8_enabled
         _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase
 
         qstate = FP8GlobalStateManager.quantization_state
+        self._previous_is_first_fp8_module = qstate.is_first_fp8_module
         if self.activation_recompute and not self.recompute_phase:
-            activation_recompute_forward._is_first_fp8_module.append(qstate.is_first_fp8_module)
+            self.state.is_first_fp8_module = qstate.is_first_fp8_module
+            # Reentrant checkpoint forward runs under no_grad, so no module can
+            # consume backward-update ownership. Reserve it for this frame.
+            if self.reserve_first_fp8_module and fp8_enabled:
+                qstate.is_first_fp8_module = False
         if self.activation_recompute and self.recompute_phase:
-            qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0)
+            if self.state.is_first_fp8_module is None:
+                raise RuntimeError("FP8 recompute state was not captured during forward")
+            qstate.is_first_fp8_module = self.state.is_first_fp8_module
+        return self
 
     def __exit__(self, *exc_details):
         global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
-        _FP8_ACTIVATION_RECOMPUTE_ENABLED = False
-        _FP8_ACTIVATION_RECOMPUTE_PHASE = False
+        if self.activation_recompute and self.recompute_phase:
+            # Only recompute restores outer ownership. After the original forward, it must remain
+            # consumed or reserved so later checkpoint frames cannot claim it.
+            qstate = FP8GlobalStateManager.quantization_state
+            qstate.is_first_fp8_module = self._previous_is_first_fp8_module
+        _FP8_ACTIVATION_RECOMPUTE_ENABLED = self._previous_recompute_enabled
+        _FP8_ACTIVATION_RECOMPUTE_PHASE = self._previous_recompute_phase
 
 
 def is_fp8_activation_recompute_enabled() -> bool:
@@ -369,8 +396,14 @@ def forward(
         # Preserve torch autocast context for the backward pass
         torch_gpu_amp_ctx, torch_cpu_amp_ctx = _get_active_autocast_contexts()
 
+        fp8_recompute_state = _ActivationRecomputeState()
         with torch.no_grad(), forward_ctx:
-            with activation_recompute_forward(activation_recompute=True, recompute_phase=False):
+            with activation_recompute_forward(
+                activation_recompute=True,
+                recompute_phase=False,
+                state=fp8_recompute_state,
+                reserve_first_fp8_module=True,
+            ):
                 outputs = run_function(*args, **kwargs)
 
         # Divide hidden states across model parallel group and only keep
@@ -395,6 +428,7 @@ def forward(
         ctx.torch_cpu_amp_ctx = torch_cpu_amp_ctx
         ctx.fp8 = fp8
         ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
+        ctx.fp8_recompute_state = fp8_recompute_state
         ctx.kwargs = kwargs
 
         return outputs
@@ -436,9 +470,13 @@ def backward(
         # Compute the forward pass.
         detached_inputs = detach_variable(inputs)
         with torch.enable_grad(), ctx.recompute_ctx, ctx.torch_gpu_amp_ctx, ctx.torch_cpu_amp_ctx, activation_recompute_forward(
-            activation_recompute=True, recompute_phase=True
+            activation_recompute=True,
+            recompute_phase=True,
+            state=ctx.fp8_recompute_state,
         ), autocast(
-            enabled=ctx.fp8, recipe=ctx.fp8_recipe
+            enabled=ctx.fp8,
+            recipe=ctx.fp8_recipe,
+            _recompute=True,
         ):
             outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
 
@@ -602,13 +640,16 @@ def use_reentrant_activation_recompute():
 
 def get_activation_recompute_contexts():
     """Returns context objects for the checkpointed forward pass and the forward recompute phase."""
+    state = _ActivationRecomputeState()
     forward_ctx = activation_recompute_forward(
         activation_recompute=True,
         recompute_phase=False,
+        state=state,
     )
     recompute_ctx = activation_recompute_forward(
         activation_recompute=True,
         recompute_phase=True,
+        state=state,
     )
     return forward_ctx, recompute_ctx
 
@@ -795,7 +836,9 @@ def recompute_fn(*args, **kwargs):
         with torch.autograd.enable_grad(), (
             te_recompute_ctx
         ), user_recompute_ctx, torch_gpu_amp_forward_ctx, torch_cpu_amp_forward_ctx, autocast(
-            enabled=fp8, recipe=fp8_recipe
+            enabled=fp8,
+            recipe=fp8_recipe,
+            _recompute=True,
         ):
             function(*args, **kwargs)
 

transformer_engine/pytorch/module/layernorm_linear.py

@@ -52,7 +52,6 @@
     symmetric_all_reduce,
     reduce_scatter_along_first_dim,
     gather_along_first_dim,
-    in_fp8_activation_recompute_phase,
     _fsdp_scatter_tensors,
     _fsdp_gather_tensors,
 )
@@ -565,11 +564,7 @@ def forward(
             ctx.normalization = normalization
             ctx.reduce_and_update_bwd_fp8_tensors = False
             if ctx.fp8 and requires_grad(inp, ln_weight, ln_bias, weight, bias):
-                qstate = FP8GlobalStateManager.quantization_state
-                _first_fp8_module = qstate.is_first_fp8_module
                 ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module()
-                if in_fp8_activation_recompute_phase():
-                    qstate.is_first_fp8_module = _first_fp8_module
             ctx.wgrad_store = wgrad_store
             ctx.debug = debug
 

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -58,7 +58,6 @@
     reduce_scatter_along_first_dim,
     gather_along_first_dim,
     use_reentrant_activation_recompute,
-    in_fp8_activation_recompute_phase,
     _fsdp_scatter_tensors,
     _get_cuda_rng_state,
     _set_cuda_rng_state,
@@ -908,7 +907,7 @@ def _forward(
                 qstate = FP8GlobalStateManager.quantization_state
                 _first_fp8_module = qstate.is_first_fp8_module
                 ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module()
-                if in_fp8_activation_recompute_phase() or is_recomputation:
+                if is_recomputation:
                     qstate.is_first_fp8_module = _first_fp8_module
 
             ctx.wgrad_store = wgrad_store

transformer_engine/pytorch/module/linear.py

@@ -245,12 +245,7 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None:
 
 def _check_fp8_reduce_and_update():
     """Check if this is the first FP8 module (for backward reduce-and-update)."""
-    qstate = FP8GlobalStateManager.quantization_state
-    _first_fp8_module = qstate.is_first_fp8_module
-    result = FP8GlobalStateManager.is_first_fp8_module()
-    if in_fp8_activation_recompute_phase():
-        qstate.is_first_fp8_module = _first_fp8_module
-    return result
+    return FP8GlobalStateManager.is_first_fp8_module()
 
 
 def _linear_forward_impl(