From 317e256b823815b227a8d35afc301b3cf034116e Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 29 Jul 2026 18:09:01 -0700 Subject: [PATCH 1/4] 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 2/4] 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 3/4] 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 4/4] 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)