Fix duplicate FP8 state updates during activation checkpoint recomputation

tests/pytorch/test_fp8_checkpoint_state_edges.py

@@ -0,0 +1,222 @@
+# 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)

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)