From 317e256b823815b227a8d35afc301b3cf034116e Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 29 Jul 2026 18:09:01 -0700 Subject: [PATCH 01/11] fix: preserve FP8 recompute state for inner autocast Signed-off-by: Nitin Vegesna --- .../pytorch/test_fp8_activation_recompute.py | 99 +++++++++++++++++++ transformer_engine/pytorch/distributed.py | 11 ++- 2 files changed, 106 insertions(+), 4 deletions(-) create mode 100644 tests/pytorch/test_fp8_activation_recompute.py diff --git a/tests/pytorch/test_fp8_activation_recompute.py b/tests/pytorch/test_fp8_activation_recompute.py new file mode 100644 index 0000000000..2836fe524f --- /dev/null +++ b/tests/pytorch/test_fp8_activation_recompute.py @@ -0,0 +1,99 @@ +# Copyright (c) 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 diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index d1525b53f0..89829fe6d4 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/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: From 4fb02b06ef4048990df33b7292dffa9361933d51 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 29 Jul 2026 23:02:11 -0700 Subject: [PATCH 02/11] test: fix activation recompute test license header Signed-off-by: Nitin Vegesna --- tests/pytorch/test_fp8_activation_recompute.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/test_fp8_activation_recompute.py b/tests/pytorch/test_fp8_activation_recompute.py index 2836fe524f..887a9de6f6 100644 --- a/tests/pytorch/test_fp8_activation_recompute.py +++ b/tests/pytorch/test_fp8_activation_recompute.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. From 34cf656cd8e01e09cd83cf52ecd3bdd2d88100c0 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 29 Jul 2026 23:08:30 -0700 Subject: [PATCH 03/11] test: run FP8 recompute coverage in PyTorch QA Signed-off-by: Nitin Vegesna --- qa/L0_pytorch_unittest/test.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 973077bf4e..726096a943 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -66,6 +66,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 From 00c727fd6f139f085c0bea75992b93d4c24668e9 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Sun, 2 Aug 2026 15:08:25 -0700 Subject: [PATCH 04/11] test: compare inner FP8 autocast recompute numerics Signed-off-by: Nitin Vegesna --- qa/L0_pytorch_unittest/test.sh | 1 - .../pytorch/test_fp8_activation_recompute.py | 99 ------------------- tests/pytorch/test_numerics.py | 65 +++++++++++- 3 files changed, 61 insertions(+), 104 deletions(-) delete mode 100644 tests/pytorch/test_fp8_activation_recompute.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 726096a943..973077bf4e 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -66,7 +66,6 @@ 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 diff --git a/tests/pytorch/test_fp8_activation_recompute.py b/tests/pytorch/test_fp8_activation_recompute.py deleted file mode 100644 index 887a9de6f6..0000000000 --- a/tests/pytorch/test_fp8_activation_recompute.py +++ /dev/null @@ -1,99 +0,0 @@ -# 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 diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 8249c7fedd..2e4aaa811b 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -648,7 +648,15 @@ def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, recipe, fp8_m def _test_e2e_full_recompute( - bs, dtype, config, fp8, recipe, fp8_model_params=False, recompute=False, use_reentrant=True + bs, + dtype, + config, + fp8, + recipe, + fp8_model_params=False, + recompute=False, + use_reentrant=True, + inner_autocast=False, ): reset_rng_states() FP8GlobalStateManager.reset() @@ -685,10 +693,17 @@ def _test_e2e_full_recompute( te_inp_hidden_states.retain_grad() te_inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q) - with autocast(enabled=fp8, recipe=recipe): + forward = block + if inner_autocast: + + def forward(*args, **kwargs): + with autocast(enabled=fp8, recipe=recipe): + return block(*args, **kwargs) + + with autocast(enabled=fp8 and not inner_autocast, recipe=recipe): if recompute: te_out = te_checkpoint( - block, + forward, te_inp_hidden_states, attention_mask=te_inp_attn_mask, checkpoint_core_attention=False, @@ -697,7 +712,7 @@ def _test_e2e_full_recompute( use_reentrant=use_reentrant, ) else: - te_out = block( + te_out = forward( te_inp_hidden_states, attention_mask=te_inp_attn_mask, checkpoint_core_attention=False, @@ -787,6 +802,48 @@ def test_gpt_full_activation_recompute( ) +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkeypatch): + """Check recompute numerics when FP8 autocast starts inside the checkpointed callable.""" + if not use_reentrant: + # Non-reentrant checkpoint becomes non-deterministic with bias+GELU fusion. + monkeypatch.setenv("NVTE_BIAS_GELU_NVFUSION", "0") + + dtype = torch.bfloat16 + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + config = model_configs["126m"] + + outputs, names = _test_e2e_full_recompute( + 1, + dtype, + config, + True, + fp8_recipe, + recompute=False, + use_reentrant=use_reentrant, + ) + outputs_recompute, _ = _test_e2e_full_recompute( + 1, + dtype, + config, + True, + fp8_recipe, + recompute=True, + use_reentrant=use_reentrant, + inner_autocast=True, + ) + + for name, ref, test in zip(names, outputs, outputs_recompute): + torch.testing.assert_close( + test, + ref, + msg=f"Mismatch in tensor {name}", + rtol=0.125, + atol=0.0675, + ) + + def _test_e2e_checkpointing_get_model(config, dtype): sigma = 0.023 init_method = init_method_normal(sigma) From 281e97f26b42ee9ecc0c5de3f88a79af023d0e16 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 4 Aug 2026 13:11:57 -0700 Subject: [PATCH 05/11] refactor: rename FP8 recompute flag to _IN_ACTIVATION_RECOMPUTE_REGION The global no longer encodes FP8 state now that the FP8 gate lives in is_fp8_activation_recompute_enabled(); it only marks the checkpoint region. Rename it to match. The public getter keeps its name since it now returns the conjunction of the region flag and the current FP8 state. Signed-off-by: Nitin Vegesna --- transformer_engine/pytorch/distributed.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 89829fe6d4..676e23885d 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -62,7 +62,7 @@ _USE_REENTRANT_ACTIVATION_RECOMPUTE = True -_FP8_ACTIVATION_RECOMPUTE_ENABLED = False +_IN_ACTIVATION_RECOMPUTE_REGION = False _FP8_ACTIVATION_RECOMPUTE_PHASE = False @@ -255,13 +255,13 @@ def __init__(self, activation_recompute: bool = False, recompute_phase: bool = F self.recompute_phase = recompute_phase def __enter__(self): - global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE + global _IN_ACTIVATION_RECOMPUTE_REGION, _FP8_ACTIVATION_RECOMPUTE_PHASE # 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 + _IN_ACTIVATION_RECOMPUTE_REGION = self.activation_recompute _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase qstate = FP8GlobalStateManager.quantization_state @@ -271,14 +271,14 @@ def __enter__(self): qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0) def __exit__(self, *exc_details): - global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE - _FP8_ACTIVATION_RECOMPUTE_ENABLED = False + global _IN_ACTIVATION_RECOMPUTE_REGION, _FP8_ACTIVATION_RECOMPUTE_PHASE + _IN_ACTIVATION_RECOMPUTE_REGION = False _FP8_ACTIVATION_RECOMPUTE_PHASE = False def is_fp8_activation_recompute_enabled() -> bool: """Return global boolean""" - return _FP8_ACTIVATION_RECOMPUTE_ENABLED and FP8GlobalStateManager.is_fp8_enabled() + return _IN_ACTIVATION_RECOMPUTE_REGION and FP8GlobalStateManager.is_fp8_enabled() def in_fp8_activation_recompute_phase() -> bool: From 199aa6fe20c08dd97c4d6e241a15eb0602e1058a Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 4 Aug 2026 13:16:34 -0700 Subject: [PATCH 06/11] refactor: rename recompute phase global for symmetry _FP8_ACTIVATION_RECOMPUTE_PHASE carries no FP8 state either - in_fp8_activation_recompute_phase() returns it ungated and callers apply their own FP8 gate. Rename to _ACTIVATION_RECOMPUTE_PHASE to match _IN_ACTIVATION_RECOMPUTE_REGION. Both public getters keep their names. Signed-off-by: Nitin Vegesna --- transformer_engine/pytorch/distributed.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 676e23885d..c9e77b9487 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -63,7 +63,7 @@ _USE_REENTRANT_ACTIVATION_RECOMPUTE = True _IN_ACTIVATION_RECOMPUTE_REGION = False -_FP8_ACTIVATION_RECOMPUTE_PHASE = False +_ACTIVATION_RECOMPUTE_PHASE = False _ALL_ACTIVE_RNG_STATES = {} @@ -255,14 +255,14 @@ def __init__(self, activation_recompute: bool = False, recompute_phase: bool = F self.recompute_phase = recompute_phase def __enter__(self): - global _IN_ACTIVATION_RECOMPUTE_REGION, _FP8_ACTIVATION_RECOMPUTE_PHASE + global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE # 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. _IN_ACTIVATION_RECOMPUTE_REGION = self.activation_recompute - _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase + _ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase qstate = FP8GlobalStateManager.quantization_state if self.activation_recompute and not self.recompute_phase: @@ -271,9 +271,9 @@ def __enter__(self): qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0) def __exit__(self, *exc_details): - global _IN_ACTIVATION_RECOMPUTE_REGION, _FP8_ACTIVATION_RECOMPUTE_PHASE + global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE _IN_ACTIVATION_RECOMPUTE_REGION = False - _FP8_ACTIVATION_RECOMPUTE_PHASE = False + _ACTIVATION_RECOMPUTE_PHASE = False def is_fp8_activation_recompute_enabled() -> bool: @@ -283,7 +283,7 @@ def is_fp8_activation_recompute_enabled() -> bool: def in_fp8_activation_recompute_phase() -> bool: """Return global boolean""" - return _FP8_ACTIVATION_RECOMPUTE_PHASE + return _ACTIVATION_RECOMPUTE_PHASE def _get_active_autocast_contexts(): From 3e4c2969abca2708e029026d210182b1f779142b Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 4 Aug 2026 13:39:28 -0700 Subject: [PATCH 07/11] test: restore dropped FP8 recompute coverage and assert stash bookkeeping Moving the regression into test_numerics.py silently dropped two of the three original tests. Restore both: the non-FP8 negative case (which needs no FP8 hardware and is the direct guard for making the region flag FP8-agnostic) and the mixed FP8/non-FP8 region case. The inner-autocast test could not observe a missing stash: the outer autocast is disabled, so autocast_depth never returns to 0 with FP8 enabled, reduce_and_update_fp8_tensors is never called, and the forward scale stays at 1.0 - making stashed and unstashed recompute identical. Record the stash/restore of the forward scale and assert every stash is restored exactly once, and use an inner-autocast reference so the two runs differ only by recompute. Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 92 ++++++++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 2e4aaa811b..98739ffc01 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -802,6 +802,9 @@ def test_gpt_full_activation_recompute( ) +_FP8_RECOMPUTE_KEY = "global_fp8_buffer_pos_fwd_recompute" + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("use_reentrant", all_boolean) def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkeypatch): @@ -814,6 +817,8 @@ def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkey fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) config = model_configs["126m"] + # Reference also opens the autocast inside the callable, so the only difference + # between the two runs is activation recompute. outputs, names = _test_e2e_full_recompute( 1, dtype, @@ -822,7 +827,38 @@ def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkey fp8_recipe, recompute=False, use_reentrant=use_reentrant, + inner_autocast=True, + ) + + # The outer autocast is disabled here, so the forward scale is never updated between + # the two forward phases and the comparison below cannot observe a missing stash. + # Record the stash/restore of the forward scale and assert the bookkeeping directly. + forward_key = FP8GlobalStateManager.get_meta_tensor_key(forward=True) + stashed_scales, restored_scales = [], [] + stash_fn = FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute + restore_fn = FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute + + def record_stash(fp8_meta): + stash_fn(fp8_meta) + if _FP8_RECOMPUTE_KEY in fp8_meta: + stashed_scales.append(fp8_meta[forward_key].scale.clone()) + + def record_restore(fp8_meta): + restore_fn(fp8_meta) + if _FP8_RECOMPUTE_KEY in fp8_meta: + restored_scales.append(fp8_meta[forward_key].scale.clone()) + + monkeypatch.setattr( + FP8GlobalStateManager, + "copy_forward_fp8_meta_tensors_for_recompute", + staticmethod(record_stash), + ) + monkeypatch.setattr( + FP8GlobalStateManager, + "get_old_fp8_meta_tensors_for_recompute", + staticmethod(record_restore), ) + outputs_recompute, _ = _test_e2e_full_recompute( 1, dtype, @@ -834,6 +870,15 @@ def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkey inner_autocast=True, ) + assert stashed_scales, "No FP8 module stashed a forward scale for the recompute phase" + assert len(restored_scales) == len( + stashed_scales + ), "Every stashed forward scale must be restored exactly once in the recompute phase" + for i, (stashed, restored) in enumerate(zip(stashed_scales, restored_scales)): + torch.testing.assert_close( + restored, stashed, rtol=0.0, atol=0.0, msg=f"Recompute scale differs for module {i}" + ) + for name, ref, test in zip(names, outputs, outputs_recompute): torch.testing.assert_close( test, @@ -844,6 +889,53 @@ def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkey ) +def _checkpointed_linear_backward(body, use_reentrant, *layers): + """Run a checkpointed callable end to end and check the gradients are finite.""" + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + out = te_checkpoint(body, inp, use_reentrant=use_reentrant) + loss = out.float().sum() + loss.backward() + torch.cuda.synchronize() + + assert torch.isfinite(loss) + assert inp.grad is not None and 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.parametrize("use_reentrant", all_boolean) +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() + + _checkpointed_linear_backward(layer, use_reentrant, layer) + + assert _FP8_RECOMPUTE_KEY not in layer.fp8_meta + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +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() + + def body(value): + value = non_fp8_layer(value) + with autocast(enabled=True, recipe=fp8_recipe): + return fp8_layer(value) + + _checkpointed_linear_backward(body, use_reentrant, non_fp8_layer, fp8_layer) + + assert _FP8_RECOMPUTE_KEY not in non_fp8_layer.fp8_meta + assert _FP8_RECOMPUTE_KEY in fp8_layer.fp8_meta + + def _test_e2e_checkpointing_get_model(config, dtype): sigma = 0.023 init_method = init_method_normal(sigma) From 2ce745f17f8bd4b0071fac426691f36d6c2e7884 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 4 Aug 2026 13:55:54 -0700 Subject: [PATCH 08/11] test: make recompute stash assertions meaningful, guard bf16 The scale equality assertion was vacuous: record_restore read the scale after the real restore had already copied the stashed value into it, so it compared a tensor with itself, and every scale is 1.0 here anyway. Count stash and restore per module instead, and check a module was stashed before the real restore runs so a regression reports that rather than a KeyError from the recompute buffer lookup. Also skip the non-FP8 checkpoint test when bf16 is unavailable (it has no FP8 skipif, so it would otherwise run on pre-Ampere), and correct the is_fp8_activation_recompute_enabled docstring, which still claimed to return a bare global. Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 27 +++++++++-------------- transformer_engine/pytorch/distributed.py | 2 +- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 98739ffc01..dec16d3f69 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -830,23 +830,23 @@ def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkey inner_autocast=True, ) - # The outer autocast is disabled here, so the forward scale is never updated between - # the two forward phases and the comparison below cannot observe a missing stash. - # Record the stash/restore of the forward scale and assert the bookkeeping directly. - forward_key = FP8GlobalStateManager.get_meta_tensor_key(forward=True) - stashed_scales, restored_scales = [], [] + # The outer autocast is disabled here, so the forward scale is never updated between the + # two forward phases: a missing stash cannot change any number, it can only crash. Count + # the stash and restore per module instead, and check a module was stashed before the real + # restore runs, so a regression reports that rather than a KeyError on the buffer lookup. + stash_counts, restore_counts = {}, {} stash_fn = FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute restore_fn = FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute def record_stash(fp8_meta): stash_fn(fp8_meta) if _FP8_RECOMPUTE_KEY in fp8_meta: - stashed_scales.append(fp8_meta[forward_key].scale.clone()) + stash_counts[id(fp8_meta)] = stash_counts.get(id(fp8_meta), 0) + 1 def record_restore(fp8_meta): + assert id(fp8_meta) in stash_counts, "Recompute restored a scale that was never stashed" + restore_counts[id(fp8_meta)] = restore_counts.get(id(fp8_meta), 0) + 1 restore_fn(fp8_meta) - if _FP8_RECOMPUTE_KEY in fp8_meta: - restored_scales.append(fp8_meta[forward_key].scale.clone()) monkeypatch.setattr( FP8GlobalStateManager, @@ -870,14 +870,8 @@ def record_restore(fp8_meta): inner_autocast=True, ) - assert stashed_scales, "No FP8 module stashed a forward scale for the recompute phase" - assert len(restored_scales) == len( - stashed_scales - ), "Every stashed forward scale must be restored exactly once in the recompute phase" - for i, (stashed, restored) in enumerate(zip(stashed_scales, restored_scales)): - torch.testing.assert_close( - restored, stashed, rtol=0.0, atol=0.0, msg=f"Recompute scale differs for module {i}" - ) + assert stash_counts, "No FP8 module stashed a forward scale for the recompute phase" + assert restore_counts == stash_counts, "Stash and restore of forward scales are unbalanced" for name, ref, test in zip(names, outputs, outputs_recompute): torch.testing.assert_close( @@ -905,6 +899,7 @@ def _checkpointed_linear_backward(body, use_reentrant, *layers): assert torch.isfinite(layer.weight.grad).all() +@pytest.mark.skipif(not is_bf16_available(), reason="bf16 requires sm_80 or higher") @pytest.mark.parametrize("use_reentrant", all_boolean) def test_checkpoint_without_fp8_does_not_save_fp8_recompute_state(use_reentrant): """A checkpointed non-FP8 module does not save FP8 recompute metadata.""" diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index c9e77b9487..2dddb6587e 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -277,7 +277,7 @@ def __exit__(self, *exc_details): def is_fp8_activation_recompute_enabled() -> bool: - """Return global boolean""" + """Whether we are in an activation recompute region with FP8 currently enabled""" return _IN_ACTIVATION_RECOMPUTE_REGION and FP8GlobalStateManager.is_fp8_enabled() From 9902b535ed425189f589376d05cbd20410570d1f Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 4 Aug 2026 14:23:31 -0700 Subject: [PATCH 09/11] test: assert the recompute reinstalls the stashed FP8 state Counting stashes and restores catches a missing stash but not a wrong one. The outputs cannot catch either: the forward scale never moves here (autocast_depth stays >= 1 so reduce_and_update_fp8_tensors never fires), and a wrong scale would only perturb FP8 rounding regardless, since scale_inv is derived from the same scale at cast time. Compare the state each restore installs against a clone captured at stash time, for both the scale and the amax history, and assert at least one restore actually moved the live state so the comparison is not vacuous. This is exact rather than tolerance-based and works for both reentrant and non-reentrant checkpointing. Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 40 ++++++++++++++++++++++++++-------- 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index dec16d3f69..633bf8a1e3 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -830,23 +830,42 @@ def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkey inner_autocast=True, ) - # The outer autocast is disabled here, so the forward scale is never updated between the - # two forward phases: a missing stash cannot change any number, it can only crash. Count - # the stash and restore per module instead, and check a module was stashed before the real - # restore runs, so a regression reports that rather than a KeyError on the buffer lookup. - stash_counts, restore_counts = {}, {} + # The outputs cannot see a bad stash: the outer autocast is disabled, so the forward scale is + # never updated between the two forward phases, and a wrong scale would only perturb FP8 + # rounding anyway because scale_inv is derived from the same scale at cast time. Assert on the + # stashed state instead - each restore must reinstall exactly what was stashed before the + # phase-1 forward - and check a module was stashed before the real restore runs, so a missing + # stash reports that rather than a KeyError on the buffer lookup. + forward_key = FP8GlobalStateManager.get_meta_tensor_key(forward=True) + stash_counts, restore_counts, stashed_state, restore_moved = {}, {}, {}, [] stash_fn = FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute restore_fn = FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute def record_stash(fp8_meta): stash_fn(fp8_meta) - if _FP8_RECOMPUTE_KEY in fp8_meta: - stash_counts[id(fp8_meta)] = stash_counts.get(id(fp8_meta), 0) + 1 + if _FP8_RECOMPUTE_KEY not in fp8_meta: + return + key = id(fp8_meta) + stash_counts[key] = stash_counts.get(key, 0) + 1 + scaling = fp8_meta[forward_key] + stashed_state.setdefault(key, []).append( + (scaling.scale.clone(), scaling.amax_history.clone()) + ) def record_restore(fp8_meta): - assert id(fp8_meta) in stash_counts, "Recompute restored a scale that was never stashed" - restore_counts[id(fp8_meta)] = restore_counts.get(id(fp8_meta), 0) + 1 + key = id(fp8_meta) + assert key in stash_counts, "Recompute restored a scale that was never stashed" + restore_counts[key] = restore_counts.get(key, 0) + 1 + scaling = fp8_meta[forward_key] + amax_before = scaling.amax_history.clone() restore_fn(fp8_meta) + want_scale, want_amax = stashed_state[key].pop(0) + assert torch.equal(scaling.scale, want_scale), "Recompute got a different forward scale" + assert torch.equal( + scaling.amax_history, want_amax + ), "Recompute got a different amax history" + # Phase 1 writes amax, so a real restore has to move the live state back. + restore_moved.append(not torch.equal(amax_before, want_amax)) monkeypatch.setattr( FP8GlobalStateManager, @@ -872,6 +891,9 @@ def record_restore(fp8_meta): assert stash_counts, "No FP8 module stashed a forward scale for the recompute phase" assert restore_counts == stash_counts, "Stash and restore of forward scales are unbalanced" + assert any( + restore_moved + ), "No restore moved the live FP8 state, so the checks above are vacuous" for name, ref, test in zip(names, outputs, outputs_recompute): torch.testing.assert_close( From 0d2e9b2af4ab7e19e1b64d2714b8d41081cd4476 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 4 Aug 2026 14:37:03 -0700 Subject: [PATCH 10/11] test: drop tautological recompute state assertions Comparing the restored state against a clone taken in the stash wrapper could not fail: the real stash clones the same tensors at the same instant, and the restore copies that entry back, so both sides derive from one snapshot. It exercised the stash/restore plumbing rather than anything this change touches. Keep the checks that do have teeth - that some module stashed, and that stashes and restores balance per module - and skip the bookkeeping for modules whose recipe is not delayed scaling, since the restore site is not gated on it while the stash site is. Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 36 ++++++++++------------------------ 1 file changed, 10 insertions(+), 26 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 633bf8a1e3..0c5f5112d7 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -832,40 +832,27 @@ def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkey # The outputs cannot see a bad stash: the outer autocast is disabled, so the forward scale is # never updated between the two forward phases, and a wrong scale would only perturb FP8 - # rounding anyway because scale_inv is derived from the same scale at cast time. Assert on the - # stashed state instead - each restore must reinstall exactly what was stashed before the - # phase-1 forward - and check a module was stashed before the real restore runs, so a missing - # stash reports that rather than a KeyError on the buffer lookup. - forward_key = FP8GlobalStateManager.get_meta_tensor_key(forward=True) - stash_counts, restore_counts, stashed_state, restore_moved = {}, {}, {}, [] + # rounding anyway because scale_inv is derived from the same scale at cast time. What is worth + # asserting is that the stash happened at all - before the fix, phase 1 skipped it while the + # recompute still restored, which surfaced as a KeyError on the recompute buffer lookup. + stash_counts, restore_counts = {}, {} stash_fn = FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute restore_fn = FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute def record_stash(fp8_meta): stash_fn(fp8_meta) - if _FP8_RECOMPUTE_KEY not in fp8_meta: - return - key = id(fp8_meta) - stash_counts[key] = stash_counts.get(key, 0) + 1 - scaling = fp8_meta[forward_key] - stashed_state.setdefault(key, []).append( - (scaling.scale.clone(), scaling.amax_history.clone()) - ) + if _FP8_RECOMPUTE_KEY in fp8_meta: + stash_counts[id(fp8_meta)] = stash_counts.get(id(fp8_meta), 0) + 1 def record_restore(fp8_meta): + # The restore site is not gated on delayed scaling, but only delayed scaling stashes. + if not fp8_meta["recipe"].delayed(): + restore_fn(fp8_meta) + return key = id(fp8_meta) assert key in stash_counts, "Recompute restored a scale that was never stashed" restore_counts[key] = restore_counts.get(key, 0) + 1 - scaling = fp8_meta[forward_key] - amax_before = scaling.amax_history.clone() restore_fn(fp8_meta) - want_scale, want_amax = stashed_state[key].pop(0) - assert torch.equal(scaling.scale, want_scale), "Recompute got a different forward scale" - assert torch.equal( - scaling.amax_history, want_amax - ), "Recompute got a different amax history" - # Phase 1 writes amax, so a real restore has to move the live state back. - restore_moved.append(not torch.equal(amax_before, want_amax)) monkeypatch.setattr( FP8GlobalStateManager, @@ -891,9 +878,6 @@ def record_restore(fp8_meta): assert stash_counts, "No FP8 module stashed a forward scale for the recompute phase" assert restore_counts == stash_counts, "Stash and restore of forward scales are unbalanced" - assert any( - restore_moved - ), "No restore moved the live FP8 state, so the checks above are vacuous" for name, ref, test in zip(names, outputs, outputs_recompute): torch.testing.assert_close( From 4c031dcd07dfe342d90d78e7b6473a6cf85d8243 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 4 Aug 2026 14:51:09 -0700 Subject: [PATCH 11/11] test: pin the changed predicate directly Replace test_checkpoint_without_fp8_does_not_save_fp8_recompute_state, which could not fail: with no autocast self.fp8 is False, so the stash site is unreachable before the region flag is ever consulted, and the assertion holds no matter what the getter returns. The mixed-region test already makes the same negative claim in an FP8-enabled session. Assert instead on is_fp8_activation_recompute_enabled() itself from inside a checkpointed callable that opens its own autocast: True in both phases inside the autocast, False outside it. The first half fails before the fix; the second pins the FP8 term against a later over-correction that drops it. Also move the fp8_meta key constant up to the other module constants. Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 44 ++++++++++++++++++++++++---------- 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 0c5f5112d7..e6a83d92bc 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -42,6 +42,10 @@ is_nvfp4_available, ) from transformer_engine.pytorch import checkpoint as te_checkpoint +from transformer_engine.pytorch.distributed import ( + is_fp8_activation_recompute_enabled, + in_fp8_activation_recompute_phase, +) from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.common import recipe from transformer_engine.pytorch import DType @@ -82,6 +86,9 @@ all_boolean = [True, False] +# fp8_meta key written by FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute +_FP8_RECOMPUTE_KEY = "global_fp8_buffer_pos_fwd_recompute" + all_activations = [ "gelu", "geglu", @@ -802,9 +809,6 @@ def test_gpt_full_activation_recompute( ) -_FP8_RECOMPUTE_KEY = "global_fp8_buffer_pos_fwd_recompute" - - @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("use_reentrant", all_boolean) def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkeypatch): @@ -830,11 +834,8 @@ def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkey inner_autocast=True, ) - # The outputs cannot see a bad stash: the outer autocast is disabled, so the forward scale is - # never updated between the two forward phases, and a wrong scale would only perturb FP8 - # rounding anyway because scale_inv is derived from the same scale at cast time. What is worth - # asserting is that the stash happened at all - before the fix, phase 1 skipped it while the - # recompute still restored, which surfaced as a KeyError on the recompute buffer lookup. + # Before the fix, phase 1 skipped the stash while the recompute still restored it, which + # surfaced as a KeyError on the recompute buffer lookup. Count both sides to pin that down. stash_counts, restore_counts = {}, {} stash_fn = FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute restore_fn = FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute @@ -905,16 +906,33 @@ def _checkpointed_linear_backward(body, use_reentrant, *layers): assert torch.isfinite(layer.weight.grad).all() -@pytest.mark.skipif(not is_bf16_available(), reason="bf16 requires sm_80 or higher") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("use_reentrant", all_boolean) -def test_checkpoint_without_fp8_does_not_save_fp8_recompute_state(use_reentrant): - """A checkpointed non-FP8 module does not save FP8 recompute metadata.""" +def test_checkpoint_inner_autocast_is_an_fp8_recompute_region(use_reentrant): + """An FP8 autocast opened inside a checkpointed callable is an FP8 recompute region.""" FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() - _checkpointed_linear_backward(layer, use_reentrant, layer) + observed = [] + + def body(value): + outside = is_fp8_activation_recompute_enabled() + with autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + outside, + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + ) + ) + return layer(value) + + _checkpointed_linear_backward(body, use_reentrant, layer) - assert _FP8_RECOMPUTE_KEY not in layer.fp8_meta + # One entry for the checkpointed forward, one for the recompute during backward. The + # query is only an FP8 recompute region inside the autocast, in both phases. + assert observed == [(False, True, False), (False, True, True)] @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)