[PyTorch] Preserve FP8 recompute state for inner autocast

qa/L0_pytorch_unittest/test.sh

@@ -62,6 +62,7 @@ if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then
     python3 $TE_PATH/tests/pytorch/test_checkpoint.py --save-checkpoint all || error_exit "Failed to generate checkpoint files"
 fi
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_checkpoint.xml $TE_PATH/tests/pytorch/test_checkpoint.py || test_fail "test_checkpoint.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fp8_activation_recompute.xml $TE_PATH/tests/pytorch/test_fp8_activation_recompute.py || test_fail "test_fp8_activation_recompute.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_router.xml $TE_PATH/tests/pytorch/test_fused_router.py || test_fail "test_fused_router.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml $TE_PATH/tests/pytorch/test_partial_cast.py || test_fail "test_partial_cast.py"
 # Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision

tests/pytorch/test_fp8_activation_recompute.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+from transformer_engine.common import recipe
+from transformer_engine.pytorch import Linear, autocast, checkpoint
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
+
+
+fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)
+
+
+def _make_input():
+    return torch.randn(
+        16,
+        16,
+        device="cuda",
+        dtype=torch.bfloat16,
+        requires_grad=True,
+    )
+
+
+def _assert_finite_loss_and_grads(loss, inp, *layers):
+    assert torch.isfinite(loss)
+    assert inp.grad is not None
+    assert torch.isfinite(inp.grad).all()
+    for layer in layers:
+        assert layer.weight.grad is not None
+        assert torch.isfinite(layer.weight.grad).all()
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("use_reentrant", [True, False])
+def test_fp8_checkpoint_with_inner_autocast(use_reentrant):
+    """Delayed-scaling metadata is preserved when FP8 starts inside the checkpoint."""
+    FP8GlobalStateManager.reset()
+    fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+    layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+    inp = _make_input()
+
+    def checkpointed_body(value):
+        with autocast(enabled=True, recipe=fp8_recipe):
+            return layer(value)
+
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = checkpoint(checkpointed_body, inp, use_reentrant=use_reentrant)
+        loss = out.float().sum()
+    loss.backward()
+    torch.cuda.synchronize()
+
+    _assert_finite_loss_and_grads(loss, inp, layer)
+    assert "global_fp8_buffer_pos_fwd_recompute" in layer.fp8_meta
+
+
+@pytest.mark.parametrize("use_reentrant", [True, False])
+def test_checkpoint_without_fp8_does_not_save_fp8_recompute_state(use_reentrant):
+    """A checkpointed non-FP8 module does not save FP8 recompute metadata."""
+    FP8GlobalStateManager.reset()
+    layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+    inp = _make_input()
+
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = checkpoint(layer, inp, use_reentrant=use_reentrant)
+        loss = out.float().sum()
+    loss.backward()
+    torch.cuda.synchronize()
+
+    _assert_finite_loss_and_grads(loss, inp, layer)
+    assert "global_fp8_buffer_pos_fwd_recompute" not in layer.fp8_meta
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("use_reentrant", [True, False])
+def test_checkpoint_with_mixed_fp8_regions_saves_only_fp8_recompute_state(use_reentrant):
+    """Only the inner FP8 region of a mixed checkpoint saves recompute metadata."""
+    FP8GlobalStateManager.reset()
+    fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+    non_fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+    fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+    inp = _make_input()
+
+    def checkpointed_body(value):
+        value = non_fp8_layer(value)
+        with autocast(enabled=True, recipe=fp8_recipe):
+            return fp8_layer(value)
+
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = checkpoint(checkpointed_body, inp, use_reentrant=use_reentrant)
+        loss = out.float().sum()
+    loss.backward()
+    torch.cuda.synchronize()
+
+    _assert_finite_loss_and_grads(loss, inp, non_fp8_layer, fp8_layer)
+    assert "global_fp8_buffer_pos_fwd_recompute" not in non_fp8_layer.fp8_meta
+    assert "global_fp8_buffer_pos_fwd_recompute" in fp8_layer.fp8_meta

transformer_engine/pytorch/distributed.py

@@ -256,9 +256,12 @@ def __init__(self, activation_recompute: bool = False, recompute_phase: bool = F
 
     def __enter__(self):
         global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
-        _FP8_ACTIVATION_RECOMPUTE_ENABLED = (
-            self.activation_recompute and FP8GlobalStateManager.is_fp8_enabled()
-        )
+        # Track the checkpoint region independently of the FP8 state at entry.
+        # A checkpointed callable may open its own FP8 autocast context (for
+        # example, to select precision per layer). Delayed-scaling modules in
+        # that inner context must still save their scale and amax metadata for
+        # the recompute forward.
+        _FP8_ACTIVATION_RECOMPUTE_ENABLED = self.activation_recompute
         _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase
 
         qstate = FP8GlobalStateManager.quantization_state
@@ -275,7 +278,7 @@ def __exit__(self, *exc_details):
 
 def is_fp8_activation_recompute_enabled() -> bool:
     """Return global boolean"""
-    return _FP8_ACTIVATION_RECOMPUTE_ENABLED
+    return _FP8_ACTIVATION_RECOMPUTE_ENABLED and FP8GlobalStateManager.is_fp8_enabled()
 
 
 def in_fp8_activation_recompute_phase() -> bool: