From 5b6f598f9578334f967ff71f3a5d780b0828292f Mon Sep 17 00:00:00 2001
From: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
Date: Wed, 15 Jul 2026 11:43:17 +0800
Subject: [PATCH 1/5] test: reproduce duplicate FP8 updates during activation
 checkpoint recompute

Signed-off-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
---
 tests/pytorch/test_reentrant_fp8_updates.py | 159 ++++++++++++++++++++
 1 file changed, 159 insertions(+)
 create mode 100644 tests/pytorch/test_reentrant_fp8_updates.py

diff --git a/tests/pytorch/test_reentrant_fp8_updates.py b/tests/pytorch/test_reentrant_fp8_updates.py
new file mode 100644
index 0000000000..46d40235a2
--- /dev/null
+++ b/tests/pytorch/test_reentrant_fp8_updates.py
@@ -0,0 +1,159 @@
+# 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):
+        del cls
+        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.125, atol=0.0675)

From 396948571937518bd8f2d99044aba80070bec09a Mon Sep 17 00:00:00 2001
From: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
Date: Wed, 15 Jul 2026 11:51:03 +0800
Subject: [PATCH 2/5] fix: preserve FP8 update ownership across recompute

Signed-off-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
---
 transformer_engine/pytorch/distributed.py     | 69 +++++++++++++++----
 .../pytorch/module/layernorm_linear.py        |  5 --
 .../pytorch/module/layernorm_mlp.py           |  3 +-
 transformer_engine/pytorch/module/linear.py   |  7 +-
 transformer_engine/pytorch/quantization.py    | 23 +++++--
 5 files changed, 76 insertions(+), 31 deletions(-)

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index c050f26869..af85fa0ea7 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/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,48 @@ 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:
+            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 +394,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 +426,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 +468,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 +638,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 +834,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)
 
diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 43799b003c..715782e30b 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/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
 
diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py
index ad629c3f11..f614d1cae3 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/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
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 78a4d31852..f5abb5eb93 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/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(
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index f1bffb122f..45083e5726 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -724,6 +724,7 @@ def autocast_enter(
         fp8_recipe: Optional[Recipe] = None,
         fp8_group: Optional[dist_group_type] = None,
         _graph: bool = False,
+        _recompute: bool = False,
     ) -> None:
         """Set state and tracking variables for entry into FP8 region."""
 
@@ -741,7 +742,7 @@ def autocast_enter(
         qstate.fp8_distributed_group = fp8_group
         qstate.fp8_graph_capturing = _graph
 
-        if qstate.autocast_depth == 0:
+        if qstate.autocast_depth == 0 and not _recompute:
             qstate.is_first_fp8_module = True
         qstate.autocast_depth += 1
 
@@ -759,14 +760,20 @@ def autocast_enter(
                 assert nvfp4_available, reason_for_no_nvfp4
 
     @classmethod
-    def autocast_exit(cls, enabled: bool, _graph: bool) -> None:
+    def autocast_exit(cls, enabled: bool, _graph: bool, _recompute: bool = False) -> None:
         """Set state and tracking variables for exit from FP8 region."""
         qstate = cls.quantization_state
         qstate.autocast_depth -= 1
         # Reduce only the non-FP8 weight modules here.
         # FP8 weight modules are reduced at the end of the optimizer
         # step after the weight amax is populated.
-        if enabled and qstate.autocast_depth == 0 and not _graph and torch.is_grad_enabled():
+        if (
+            enabled
+            and qstate.autocast_depth == 0
+            and not _graph
+            and not _recompute
+            and torch.is_grad_enabled()
+        ):
             # delayed scaling only function, for other recipes (current scaling with any granularity),
             # this is noop for other recipes because cls.global_amax_buffer is empty list
             cls.reduce_and_update_fp8_tensors(forward=True)
@@ -1006,6 +1013,7 @@ class autocast:
         "_recipe",
         "_amax_reduction_group",
         "_graph",
+        "_recompute",
         "_fp8_state",
     )
 
@@ -1016,12 +1024,14 @@ def __init__(
         recipe: Optional["Recipe"] = None,
         amax_reduction_group: Optional["dist_group_type"] = None,
         _graph: bool = False,
+        _recompute: bool = False,
     ) -> None:
         self._enabled = enabled
         self._calibrating = calibrating
         self._recipe = recipe
         self._amax_reduction_group = amax_reduction_group
         self._graph = _graph
+        self._recompute = _recompute
         self._fp8_state = None
 
     def __enter__(self) -> "autocast":
@@ -1040,13 +1050,18 @@ def __enter__(self) -> "autocast":
             fp8_recipe=self._recipe,
             fp8_group=self._amax_reduction_group,
             _graph=self._graph,
+            _recompute=self._recompute,
         )
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         try:
             FP8GlobalStateManager.set_autocast_state(self._fp8_state)
-            FP8GlobalStateManager.autocast_exit(self._enabled, _graph=self._graph)
+            FP8GlobalStateManager.autocast_exit(
+                self._enabled,
+                _graph=self._graph,
+                _recompute=self._recompute,
+            )
         finally:
             # Clear the saved state so the instance can be entered again
             # sequentially (and so a failure inside the restore path does not

From dc4a11ddfac3a129db9287790a5bfc85646d363c Mon Sep 17 00:00:00 2001
From: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
Date: Wed, 15 Jul 2026 13:18:50 +0800
Subject: [PATCH 3/5] test: tighten FP8 recompute regression checks

Signed-off-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
---
 tests/pytorch/test_reentrant_fp8_updates.py | 5 ++---
 transformer_engine/pytorch/distributed.py   | 2 ++
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/tests/pytorch/test_reentrant_fp8_updates.py b/tests/pytorch/test_reentrant_fp8_updates.py
index 46d40235a2..a6eca97b65 100644
--- a/tests/pytorch/test_reentrant_fp8_updates.py
+++ b/tests/pytorch/test_reentrant_fp8_updates.py
@@ -41,8 +41,7 @@ def test_delayed_scaling_updates_once_per_autocast(
     counts = {"forward": 0, "backward": 0}
     original_update = FP8GlobalStateManager.reduce_and_update_fp8_tensors
 
-    def counted_update(cls, forward=True):
-        del cls
+    def counted_update(_cls, forward=True):
         counts["forward" if forward else "backward"] += 1
         return original_update(forward=forward)
 
@@ -156,4 +155,4 @@ def run(checkpoint):
     reference = run(checkpoint=False)
     checkpointed = run(checkpoint=True)
     for actual, expected in zip(checkpointed, reference):
-        torch.testing.assert_close(actual, expected, rtol=0.125, atol=0.0675)
+        torch.testing.assert_close(actual, expected, rtol=0.05, atol=0.01)
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index af85fa0ea7..11cd15ba98 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -292,6 +292,8 @@ def __enter__(self):
     def __exit__(self, *exc_details):
         global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
         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

From 2f3489e3bb12e2777a75077f55a25898954f2392 Mon Sep 17 00:00:00 2001
From: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
Date: Thu, 23 Jul 2026 01:19:07 +0800
Subject: [PATCH 4/5] fix: handle FP8 checkpoint ownership edge cases

---
 .../test_fp8_checkpoint_state_edges.py        | 227 ++++++++++++++++++
 transformer_engine/pytorch/distributed.py     |  68 ++++--
 transformer_engine/pytorch/quantization.py    |  11 +
 3 files changed, 284 insertions(+), 22 deletions(-)
 create mode 100644 tests/pytorch/test_fp8_checkpoint_state_edges.py

diff --git a/tests/pytorch/test_fp8_checkpoint_state_edges.py b/tests/pytorch/test_fp8_checkpoint_state_edges.py
new file mode 100644
index 0000000000..16ea12d73a
--- /dev/null
+++ b/tests/pytorch/test_fp8_checkpoint_state_edges.py
@@ -0,0 +1,227 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Targeted FP8 checkpoint ownership and exception-state regressions."""
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+from transformer_engine.common.recipe import DelayedScaling, Format
+from transformer_engine.pytorch.distributed import (
+    _ActivationRecomputeState,
+    activation_recompute_forward,
+    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)
+pytestmark = pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+
+WIDTH = 16
+SEED = 20260722
+
+
+@dataclass
+class UpdateCounter:
+    forward: int = 0
+    backward: int = 0
+
+    def snapshot(self) -> tuple[int, int]:
+        return self.forward, self.backward
+
+    def delta(self, before: tuple[int, int]) -> tuple[int, int]:
+        return self.forward - before[0], self.backward - before[1]
+
+
+@pytest.fixture(autouse=True)
+def clean_fp8_state():
+    FP8GlobalStateManager.reset()
+    yield
+    qstate = FP8GlobalStateManager.quantization_state
+    assert qstate.autocast_depth == 0
+    assert not FP8GlobalStateManager.is_fp8_enabled()
+    assert not is_fp8_activation_recompute_enabled()
+    assert not in_fp8_activation_recompute_phase()
+
+
+@pytest.fixture
+def update_counter(monkeypatch) -> UpdateCounter:
+    counter = UpdateCounter()
+    original = FP8GlobalStateManager.reduce_and_update_fp8_tensors
+
+    def counted(_cls, forward=True):
+        if forward:
+            counter.forward += 1
+        else:
+            counter.backward += 1
+        return original(forward=forward)
+
+    monkeypatch.setattr(
+        FP8GlobalStateManager,
+        "reduce_and_update_fp8_tensors",
+        classmethod(counted),
+    )
+    return counter
+
+
+def make_linear() -> te.Linear:
+    return te.Linear(
+        WIDTH,
+        WIDTH,
+        bias=False,
+        params_dtype=torch.float32,
+        init_method=lambda tensor: torch.nn.init.normal_(tensor, mean=0.0, std=0.1),
+    ).cuda()
+
+
+def make_input() -> torch.Tensor:
+    return torch.randn(
+        WIDTH,
+        WIDTH,
+        dtype=torch.bfloat16,
+        device="cuda",
+        requires_grad=True,
+    )
+
+
+def assert_finite_nonzero(loss, inp, module) -> None:
+    assert torch.isfinite(loss)
+    assert inp.grad is not None
+    assert torch.isfinite(inp.grad).all()
+    assert torch.count_nonzero(inp.grad) > 0
+    for parameter in module.parameters():
+        assert parameter.grad is not None
+        assert torch.isfinite(parameter.grad).all()
+        assert torch.count_nonzero(parameter.grad) > 0
+
+
+def assert_global_state_restored() -> None:
+    qstate = FP8GlobalStateManager.quantization_state
+    assert qstate.autocast_depth == 0
+    assert not is_fp8_activation_recompute_enabled()
+    assert not in_fp8_activation_recompute_phase()
+
+
+def run_recovery(update_counter: UpdateCounter) -> None:
+    before = update_counter.snapshot()
+    layer = make_linear()
+    inp = make_input()
+    recipe = DelayedScaling(fp8_format=Format.HYBRID)
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
+        enabled=True, recipe=recipe
+    ):
+        loss = layer(inp).float().square().mean()
+    loss.backward()
+    torch.cuda.synchronize()
+    assert_finite_nonzero(loss, inp, layer)
+    assert update_counter.delta(before) == (1, 1)
+    assert_global_state_restored()
+
+
+def test_nested_autocast_does_not_revive_consumed_owner(update_counter):
+    """A nested exit must not make first-module ownership available again."""
+    torch.manual_seed(SEED)
+    layers = torch.nn.ModuleList([make_linear(), make_linear()])
+    inp = make_input()
+    recipe = DelayedScaling(fp8_format=Format.HYBRID)
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
+        enabled=True, recipe=recipe
+    ):
+        with te.autocast(enabled=True, recipe=recipe):
+            out = layers[0](inp)
+        loss = layers[1](out).float().square().mean()
+    loss.backward()
+    torch.cuda.synchronize()
+    assert_finite_nonzero(loss, inp, layers)
+    assert update_counter.snapshot() == (1, 1)
+
+
+@pytest.mark.parametrize("use_reentrant", (True, False))
+def test_nested_autocast_inside_checkpoint_has_one_owner(update_counter, use_reentrant):
+    torch.manual_seed(SEED)
+    layers = torch.nn.ModuleList([make_linear(), make_linear()])
+    inp = make_input()
+    recipe = DelayedScaling(fp8_format=Format.HYBRID)
+
+    def body(value):
+        with te.autocast(enabled=True, recipe=recipe):
+            value = layers[0](value)
+        return layers[1](value)
+
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
+        enabled=True, recipe=recipe
+    ):
+        loss = te.checkpoint(body, inp, use_reentrant=use_reentrant).float().square().mean()
+    loss.backward()
+    torch.cuda.synchronize()
+    assert_finite_nonzero(loss, inp, layers)
+    assert update_counter.snapshot() == (1, 1)
+
+
+def test_original_forward_exception_restores_owner(update_counter):
+    """A failed checkpoint frame must return its reservation to the outer scope."""
+    recovery = make_linear()
+    inp = make_input()
+    recipe = DelayedScaling(fp8_format=Format.HYBRID)
+
+    def fail(_value):
+        raise RuntimeError("intentional original-forward failure")
+
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
+        enabled=True, recipe=recipe
+    ):
+        with pytest.raises(RuntimeError, match="intentional original-forward failure"):
+            te.checkpoint(fail, inp, use_reentrant=True)
+        loss = recovery(inp).float().square().mean()
+    loss.backward()
+    torch.cuda.synchronize()
+    assert_finite_nonzero(loss, inp, recovery)
+    assert update_counter.snapshot() == (1, 1)
+
+
+def test_recompute_enter_failure_does_not_leak_state(update_counter):
+    """Validation must happen before process-global recompute state is changed."""
+    qstate = FP8GlobalStateManager.quantization_state
+    qstate.is_first_fp8_module = True
+    with pytest.raises(RuntimeError, match="was not captured"):
+        with activation_recompute_forward(
+            activation_recompute=True,
+            recompute_phase=True,
+            state=_ActivationRecomputeState(),
+        ):
+            pass
+    assert qstate.is_first_fp8_module
+    assert_global_state_restored()
+    run_recovery(update_counter)
+
+
+def test_recompute_body_exception_restores_state(update_counter):
+    """A recompute exception must not poison a subsequent normal FP8 scope."""
+    failing = make_linear()
+    failed_input = make_input()
+    recipe = DelayedScaling(fp8_format=Format.HYBRID)
+
+    def fail_during_recompute(value):
+        result = failing(value)
+        if torch.is_grad_enabled():
+            raise RuntimeError("intentional recompute failure")
+        return result
+
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
+        enabled=True, recipe=recipe
+    ):
+        loss = te.checkpoint(
+            fail_during_recompute,
+            failed_input,
+            use_reentrant=True,
+        ).float().square().mean()
+    with pytest.raises(RuntimeError, match="intentional recompute failure"):
+        loss.backward()
+    assert_global_state_restored()
+    run_recovery(update_counter)
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 11cd15ba98..50ace56f50 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -269,35 +269,59 @@ def __init__(
 
     def __enter__(self):
         global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
-        self._previous_recompute_enabled = _FP8_ACTIVATION_RECOMPUTE_ENABLED
-        self._previous_recompute_phase = _FP8_ACTIVATION_RECOMPUTE_PHASE
+        qstate = FP8GlobalStateManager.quantization_state
         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
+        # Validate before mutating process-global state.  A failed __enter__ is not
+        # followed by __exit__, so validation after committing would leak flags.
+        if (
+            self.activation_recompute
+            and self.recompute_phase
+            and self.state.is_first_fp8_module is None
+        ):
+            raise RuntimeError("FP8 recompute state was not captured during forward")
+
+        self._previous_recompute_enabled = _FP8_ACTIVATION_RECOMPUTE_ENABLED
+        self._previous_recompute_phase = _FP8_ACTIVATION_RECOMPUTE_PHASE
         self._previous_is_first_fp8_module = qstate.is_first_fp8_module
-        if self.activation_recompute and not self.recompute_phase:
-            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:
-            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
+        try:
+            _FP8_ACTIVATION_RECOMPUTE_ENABLED = self.activation_recompute and fp8_enabled
+            _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase
+
+            if self.activation_recompute and not self.recompute_phase:
+                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
+            elif self.activation_recompute and self.recompute_phase:
+                qstate.is_first_fp8_module = self.state.is_first_fp8_module
+        except BaseException:
+            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
+            raise
         return self
 
     def __exit__(self, *exc_details):
         global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
-        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
+        qstate = FP8GlobalStateManager.quantization_state
+        try:
+            if self.activation_recompute and self.recompute_phase:
+                # Consumption during recompute is monotonic.  Merge it with the
+                # outer frame instead of restoring a stale available snapshot.
+                qstate.is_first_fp8_module = (
+                    self._previous_is_first_fp8_module and qstate.is_first_fp8_module
+                )
+            elif self.activation_recompute and exc_details and exc_details[0] is not None:
+                # A failed original forward creates no usable checkpoint frame.  Roll
+                # back ownership so a recovery module in the same outer autocast can
+                # become the first backward-update owner.
+                qstate.is_first_fp8_module = self._previous_is_first_fp8_module
+                self.state.is_first_fp8_module = None
+        finally:
+            _FP8_ACTIVATION_RECOMPUTE_ENABLED = self._previous_recompute_enabled
+            _FP8_ACTIVATION_RECOMPUTE_PHASE = self._previous_recompute_phase
 
 
 def is_fp8_activation_recompute_enabled() -> bool:
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 45083e5726..75c51d7b36 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -1056,7 +1056,18 @@ def __enter__(self) -> "autocast":
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
         try:
+            # Restoring a nested autocast's configuration must not revive backward-update
+            # ownership that was consumed while the nested scope was active.  Ownership is
+            # monotonic within the logical outer autocast: available -> consumed.
+            qstate = FP8GlobalStateManager.quantization_state
+            nested_autocast = qstate.autocast_depth > 1
+            live_is_first_fp8_module = qstate.is_first_fp8_module
             FP8GlobalStateManager.set_autocast_state(self._fp8_state)
+            if nested_autocast:
+                qstate = FP8GlobalStateManager.quantization_state
+                qstate.is_first_fp8_module = (
+                    qstate.is_first_fp8_module and live_is_first_fp8_module
+                )
             FP8GlobalStateManager.autocast_exit(
                 self._enabled,
                 _graph=self._graph,

From 026036c7be41b6590cbaa5a5ac0a36be417d496d Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 22 Jul 2026 17:20:20 +0000
Subject: [PATCH 5/5] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../test_fp8_checkpoint_state_edges.py        | 35 ++++++++-----------
 transformer_engine/pytorch/quantization.py    |  4 +--
 2 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/tests/pytorch/test_fp8_checkpoint_state_edges.py b/tests/pytorch/test_fp8_checkpoint_state_edges.py
index 16ea12d73a..a72164fee4 100644
--- a/tests/pytorch/test_fp8_checkpoint_state_edges.py
+++ b/tests/pytorch/test_fp8_checkpoint_state_edges.py
@@ -113,9 +113,7 @@ def run_recovery(update_counter: UpdateCounter) -> None:
     layer = make_linear()
     inp = make_input()
     recipe = DelayedScaling(fp8_format=Format.HYBRID)
-    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
-        enabled=True, recipe=recipe
-    ):
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(enabled=True, recipe=recipe):
         loss = layer(inp).float().square().mean()
     loss.backward()
     torch.cuda.synchronize()
@@ -130,9 +128,7 @@ def test_nested_autocast_does_not_revive_consumed_owner(update_counter):
     layers = torch.nn.ModuleList([make_linear(), make_linear()])
     inp = make_input()
     recipe = DelayedScaling(fp8_format=Format.HYBRID)
-    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
-        enabled=True, recipe=recipe
-    ):
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(enabled=True, recipe=recipe):
         with te.autocast(enabled=True, recipe=recipe):
             out = layers[0](inp)
         loss = layers[1](out).float().square().mean()
@@ -154,9 +150,7 @@ def body(value):
             value = layers[0](value)
         return layers[1](value)
 
-    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
-        enabled=True, recipe=recipe
-    ):
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(enabled=True, recipe=recipe):
         loss = te.checkpoint(body, inp, use_reentrant=use_reentrant).float().square().mean()
     loss.backward()
     torch.cuda.synchronize()
@@ -173,9 +167,7 @@ def test_original_forward_exception_restores_owner(update_counter):
     def fail(_value):
         raise RuntimeError("intentional original-forward failure")
 
-    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
-        enabled=True, recipe=recipe
-    ):
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(enabled=True, recipe=recipe):
         with pytest.raises(RuntimeError, match="intentional original-forward failure"):
             te.checkpoint(fail, inp, use_reentrant=True)
         loss = recovery(inp).float().square().mean()
@@ -213,14 +205,17 @@ def fail_during_recompute(value):
             raise RuntimeError("intentional recompute failure")
         return result
 
-    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(
-        enabled=True, recipe=recipe
-    ):
-        loss = te.checkpoint(
-            fail_during_recompute,
-            failed_input,
-            use_reentrant=True,
-        ).float().square().mean()
+    with torch.autocast("cuda", dtype=torch.bfloat16), te.autocast(enabled=True, recipe=recipe):
+        loss = (
+            te.checkpoint(
+                fail_during_recompute,
+                failed_input,
+                use_reentrant=True,
+            )
+            .float()
+            .square()
+            .mean()
+        )
     with pytest.raises(RuntimeError, match="intentional recompute failure"):
         loss.backward()
     assert_global_state_restored()
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 75c51d7b36..29aa59d692 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -1065,9 +1065,7 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> None:
             FP8GlobalStateManager.set_autocast_state(self._fp8_state)
             if nested_autocast:
                 qstate = FP8GlobalStateManager.quantization_state
-                qstate.is_first_fp8_module = (
-                    qstate.is_first_fp8_module and live_is_first_fp8_module
-                )
+                qstate.is_first_fp8_module = qstate.is_first_fp8_module and live_is_first_fp8_module
             FP8GlobalStateManager.autocast_exit(
                 self._enabled,
                 _graph=self._graph,