From 711e78b998c31a47514e8a9bbaba4d12d89df90a Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 4 Aug 2026 21:14:53 -0700 Subject: [PATCH 1/4] Pair FP8 recompute restores with their stashes The delayed-scaling stash and its two restore sites made independent decisions. Module mode changes between the original forward and checkpoint replay could therefore leak a stash or restore one that was never created. Track pending stashes per module and record whether each prepare_forward call swapped one in, so end_forward performs exactly the matching restore. Stash every delayed-scaling FP8 module encountered in checkpoint phase 1: in reentrant checkpointing the original forward runs under no_grad, so an eval module receiving an intermediate tensor has no module-local autograd signal even though backward will replay it. Tests cover training and eval modules, both checkpoint implementations, mode changes in both directions, repeated iterations, multi-module reentrant replay, and exact FIFO drainage. Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 106 ++++++++++++++++++++++ transformer_engine/pytorch/module/base.py | 31 ++++++- 2 files changed, 132 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index e6a83d92bc..e01eb629b8 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -955,6 +955,112 @@ def body(value): assert _FP8_RECOMPUTE_KEY in fp8_layer.fp8_meta +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +@pytest.mark.parametrize("training", all_boolean) +def test_checkpoint_with_eval_module_preserves_fp8_recompute_state(training, use_reentrant): + """An eval module participating in a checkpoint must stash like a training module.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + layer.train(training) + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + return layer(value) + + _checkpointed_linear_backward(body, use_reentrant, layer) + + assert _FP8_RECOMPUTE_KEY in layer.fp8_meta + assert "updated_scale_fwd" in layer.fp8_meta + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert len(recompute_buffer) == 1 + assert all(len(stashed) == 0 for stashed in recompute_buffer) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +@pytest.mark.parametrize("switch_to_eval", all_boolean) +def test_checkpoint_mode_change_between_phases_does_not_leak_recompute_stash( + switch_to_eval, use_reentrant +): + """The replay consumes the stash created by the original forward despite a mode change.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + observed = [] + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + layer.training, + ) + ) + return layer(value) + + stash_lengths = [] + for i in range(3): + layer.train() + inp = (torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) * (2.0**i)).requires_grad_() + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = te_checkpoint(body, inp, use_reentrant=use_reentrant).float().sum() + + if switch_to_eval: + layer.eval() + loss.backward() + torch.cuda.synchronize() + + assert torch.isfinite(loss) + assert inp.grad is not None and torch.isfinite(inp.grad).all() + assert layer.weight.grad is not None and torch.isfinite(layer.weight.grad).all() + assert _FP8_RECOMPUTE_KEY in layer.fp8_meta + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert len(recompute_buffer) == 1 + stash_lengths.append(len(recompute_buffer[layer.fp8_meta[_FP8_RECOMPUTE_KEY]])) + assert "updated_scale_fwd" in layer.fp8_meta + assert torch.equal(layer.fp8_meta["scaling_fwd"].scale, layer.fp8_meta["updated_scale_fwd"]) + assert observed[-2:] == [ + (True, False, True), + (True, True, not switch_to_eval), + ] + + assert stash_lengths == [0, 0, 0], f"Recompute stash leaked: {stash_lengths}" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +def test_checkpoint_eval_intermediate_stashes_under_reentrant_no_grad(use_reentrant): + """An eval module with an intermediate input stashes even under reentrant `no_grad`.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + control = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + toggled = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda().eval() + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + return toggled(control(value)) + + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True) + with torch.autocast("cuda", dtype=torch.bfloat16): + loss = te_checkpoint(body, inp, use_reentrant=use_reentrant).float().sum() + + toggled.train() + loss.backward() + torch.cuda.synchronize() + + assert inp.grad is not None and torch.isfinite(inp.grad).all() + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert len(recompute_buffer) == 2 + for layer in (control, toggled): + assert layer.weight.grad is not None and torch.isfinite(layer.weight.grad).all() + assert _FP8_RECOMPUTE_KEY in layer.fp8_meta + assert len(recompute_buffer[layer.fp8_meta[_FP8_RECOMPUTE_KEY]]) == 0 + assert "updated_scale_fwd" in layer.fp8_meta + + def _test_e2e_checkpointing_get_model(config, dtype): sigma = 0.023 init_method = init_method_normal(sigma) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index e9a65c3648..999b0d7db9 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -903,6 +903,9 @@ def __init__(self, name: Optional[str] = None) -> None: self.fp8_meta["fp8_checkpoint"] = False self.fp8_meta["fp8_group"] = None self.fp8_meta_tensors_initialized = False + # Pending FP8 recompute stashes, and whether this forward call swapped them in. + self.fp8_recompute_stashes = 0 + self.fp8_recompute_meta_restored = False self.quantizers = {"scaling_fwd": [], "scaling_bwd": []} self.tp_group = None self.tp_size = 1 @@ -1458,6 +1461,8 @@ def set_extra_state(self, state: torch.Tensor) -> None: self.fp8_meta["recipe"] = state["recipe"] if "global_fp8_buffer_pos_fwd_recompute" in self.fp8_meta: del self.fp8_meta["global_fp8_buffer_pos_fwd_recompute"] + # Dropping the position orphans any pending stashes; do not try to pop them. + self.fast_setattr("fp8_recompute_stashes", 0) # Initialize before loading self.init_fp8_meta_tensors(self.fp8_meta["recipe"]) @@ -1597,10 +1602,17 @@ def prepare_forward( ) self.fast_setattr("forwarded_at_least_once", True) + # Whether this call swapped in the stashed recompute state; `end_forward` pairs on it. + restored = False + # Activation recomputation is used and this is the second forward phase. if self.fp8 and in_fp8_activation_recompute_phase(): - delayed_scaling_recipe = _has_delayed_scaling_state(self.fp8_meta) - FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) + # Restore only what this module stashed. `self.training` is re-read here in a + # different autograd phase than the stash, so the two can disagree; the count cannot. + if self.fp8_recompute_stashes > 0: + FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) + self.fast_setattr("fp8_recompute_stashes", self.fp8_recompute_stashes - 1) + restored = True else: if not inp.is_cuda: raise RuntimeError( @@ -1632,8 +1644,15 @@ def prepare_forward( FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta) # Activation recomputation is used and this is the first forward phase. - if self.training and is_fp8_activation_recompute_enabled(): + # Every module in the first checkpoint phase must stash. In reentrant + # checkpointing that phase runs under `no_grad`, so an eval module receiving + # an intermediate tensor has no module-local autograd signal even though it + # will be replayed during backward. + if is_fp8_activation_recompute_enabled(): FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta) + self.fast_setattr("fp8_recompute_stashes", self.fp8_recompute_stashes + 1) + + self.fast_setattr("fp8_recompute_meta_restored", restored) nvtx_range_push(self.__class__.__name__ + " forward") if not allow_non_contiguous and not inp.is_contiguous(): @@ -1645,8 +1664,10 @@ def end_forward(self): Required to be called at the end of the forward function to properly handle DelayedScaling metadata handling and the NVTX ranges. """ - delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta) - if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase(): + # Pairs one-to-one with the restore in `prepare_forward`. Re-deriving the condition + # here could disagree with it and reinstate stale `updated_*_fwd` values. + if self.fp8_recompute_meta_restored: + self.fast_setattr("fp8_recompute_meta_restored", False) FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta) nvtx_range_pop() From 61a36f8c035b019d8166044c07da3a69e609f894 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 18 Aug 2026 08:52:04 -0700 Subject: [PATCH 2/4] Test minimal FP8 eval recompute stash fix Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 22 ++++++++++++++++ transformer_engine/pytorch/module/base.py | 32 +++++------------------ 2 files changed, 29 insertions(+), 25 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index e01eb629b8..330c2323fa 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1061,6 +1061,28 @@ def body(value): assert "updated_scale_fwd" 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_eval_without_backward_does_not_accumulate_recompute_stashes(use_reentrant): + """Eval forwards with autograd disabled must not leave unreachable recompute state.""" + FP8GlobalStateManager.reset() + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda().eval() + inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + + def body(value): + with autocast(enabled=True, recipe=fp8_recipe): + return layer(value) + + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): + for _ in range(3): + out = te_checkpoint(body, inp, use_reentrant=use_reentrant) + assert torch.isfinite(out).all() + + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert all(len(stashed) == 0 for stashed in recompute_buffer) + + def _test_e2e_checkpointing_get_model(config, dtype): sigma = 0.023 init_method = init_method_normal(sigma) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 999b0d7db9..ad3883c54a 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -903,9 +903,6 @@ def __init__(self, name: Optional[str] = None) -> None: self.fp8_meta["fp8_checkpoint"] = False self.fp8_meta["fp8_group"] = None self.fp8_meta_tensors_initialized = False - # Pending FP8 recompute stashes, and whether this forward call swapped them in. - self.fp8_recompute_stashes = 0 - self.fp8_recompute_meta_restored = False self.quantizers = {"scaling_fwd": [], "scaling_bwd": []} self.tp_group = None self.tp_size = 1 @@ -1461,8 +1458,6 @@ def set_extra_state(self, state: torch.Tensor) -> None: self.fp8_meta["recipe"] = state["recipe"] if "global_fp8_buffer_pos_fwd_recompute" in self.fp8_meta: del self.fp8_meta["global_fp8_buffer_pos_fwd_recompute"] - # Dropping the position orphans any pending stashes; do not try to pop them. - self.fast_setattr("fp8_recompute_stashes", 0) # Initialize before loading self.init_fp8_meta_tensors(self.fp8_meta["recipe"]) @@ -1602,17 +1597,10 @@ def prepare_forward( ) self.fast_setattr("forwarded_at_least_once", True) - # Whether this call swapped in the stashed recompute state; `end_forward` pairs on it. - restored = False - # Activation recomputation is used and this is the second forward phase. if self.fp8 and in_fp8_activation_recompute_phase(): - # Restore only what this module stashed. `self.training` is re-read here in a - # different autograd phase than the stash, so the two can disagree; the count cannot. - if self.fp8_recompute_stashes > 0: - FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) - self.fast_setattr("fp8_recompute_stashes", self.fp8_recompute_stashes - 1) - restored = True + delayed_scaling_recipe = _has_delayed_scaling_state(self.fp8_meta) + FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) else: if not inp.is_cuda: raise RuntimeError( @@ -1644,15 +1632,11 @@ def prepare_forward( FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta) # Activation recomputation is used and this is the first forward phase. - # Every module in the first checkpoint phase must stash. In reentrant - # checkpointing that phase runs under `no_grad`, so an eval module receiving - # an intermediate tensor has no module-local autograd signal even though it - # will be replayed during backward. + # Every delayed-scaling module in the first checkpoint phase must stash. + # Checkpoint phase, rather than module training mode, determines whether + # the matching recompute forward will need the original metadata. if is_fp8_activation_recompute_enabled(): FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta) - self.fast_setattr("fp8_recompute_stashes", self.fp8_recompute_stashes + 1) - - self.fast_setattr("fp8_recompute_meta_restored", restored) nvtx_range_push(self.__class__.__name__ + " forward") if not allow_non_contiguous and not inp.is_contiguous(): @@ -1664,10 +1648,8 @@ def end_forward(self): Required to be called at the end of the forward function to properly handle DelayedScaling metadata handling and the NVTX ranges. """ - # Pairs one-to-one with the restore in `prepare_forward`. Re-deriving the condition - # here could disagree with it and reinstate stale `updated_*_fwd` values. - if self.fp8_recompute_meta_restored: - self.fast_setattr("fp8_recompute_meta_restored", False) + delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta) + if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase(): FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta) nvtx_range_pop() From 71a9b8da2a1c83c0c160024b9af790d888264438 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 18 Aug 2026 09:44:00 -0700 Subject: [PATCH 3/4] Bypass FP8 recompute bookkeeping without autograd Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 17 ++++++++++++++--- transformer_engine/pytorch/distributed.py | 9 +++++++++ 2 files changed, 23 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 330c2323fa..bd7c16e8ca 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1063,15 +1063,25 @@ def body(value): @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("use_reentrant", all_boolean) -def test_checkpoint_eval_without_backward_does_not_accumulate_recompute_stashes(use_reentrant): - """Eval forwards with autograd disabled must not leave unreachable recompute state.""" +@pytest.mark.parametrize("training", all_boolean) +def test_checkpoint_without_backward_does_not_accumulate_recompute_stashes(training, use_reentrant): + """Forwards with autograd disabled must not leave unreachable recompute state.""" FP8GlobalStateManager.reset() fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) - layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda().eval() + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda() + layer.train(training) inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + observed = [] + def body(value): with autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + ) + ) return layer(value) with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16): @@ -1081,6 +1091,7 @@ def body(value): recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer assert all(len(stashed) == 0 for stashed in recompute_buffer) + assert observed == [(False, False)] * 3 def _test_e2e_checkpointing_get_model(config, dtype): diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 8605a4746b..595601f553 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -729,6 +729,7 @@ def checkpoint( context_fn = kwargs.pop("context_fn", noop_context_fn) determinism_check = kwargs.pop("determinism_check", "default") debug = kwargs.pop("debug", False) + if not has_te_modules(function): return torch.utils.checkpoint.checkpoint( function, @@ -740,6 +741,14 @@ def checkpoint( **kwargs, ) + # There will be no backward recompute when checkpoint is called with autograd + # disabled. Run the forward directly so FP8 modules do not save recompute state + # that can never be consumed. Preserve the user-provided forward context. + if not torch.is_grad_enabled(): + forward_ctx, _ = context_fn() + with forward_ctx: + return function(*args, **kwargs) + from .module.base import TransformerEngineBaseModule if isinstance(function, TransformerEngineBaseModule): From 966baaa23672445cb7c75f57bc9e05f8c62666ca Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 18 Aug 2026 10:49:44 -0700 Subject: [PATCH 4/4] Test explicit grad inside no-grad checkpoint Signed-off-by: Nitin Vegesna --- tests/pytorch/test_numerics.py | 64 ++++++++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index bd7c16e8ca..693f4a0d84 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -4,6 +4,7 @@ import math import os +from contextlib import nullcontext from typing import Dict, List, Tuple, Optional import pytest @@ -1094,6 +1095,69 @@ def body(value): assert observed == [(False, False)] * 3 +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("use_reentrant", all_boolean) +@pytest.mark.parametrize("enable_grad_in", ["context_fn", "function"]) +def test_checkpoint_outer_no_grad_preserves_explicit_inner_grad(enable_grad_in, use_reentrant): + """The no-grad bypass preserves an explicit nested request for autograd.""" + weight = torch.randn(16, 16, device="cuda", dtype=torch.float32) + input_data = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16) + fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID) + + def run(checkpointed): + FP8GlobalStateManager.reset() + layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda().eval() + with torch.no_grad(): + layer.weight.copy_(weight) + inp = input_data.clone().requires_grad_() + observed = [] + + def body(value): + grad_ctx = torch.enable_grad() if enable_grad_in == "function" else nullcontext() + with grad_ctx, autocast(enabled=True, recipe=fp8_recipe): + observed.append( + ( + torch.is_grad_enabled(), + is_fp8_activation_recompute_enabled(), + in_fp8_activation_recompute_phase(), + ) + ) + return layer(value) + + checkpoint_kwargs = {"use_reentrant": use_reentrant} + forward_ctx = nullcontext() + if enable_grad_in == "context_fn": + checkpoint_kwargs["context_fn"] = lambda: (torch.enable_grad(), nullcontext()) + if not checkpointed: + # Match the checkpoint forward context in the direct reference without + # changing grad state at the checkpoint call site itself. + forward_ctx = torch.enable_grad() + + with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16), forward_ctx: + if checkpointed: + out = te_checkpoint(body, inp, **checkpoint_kwargs) + else: + out = body(inp) + + assert out.requires_grad + out.float().sum().backward() + torch.cuda.synchronize() + assert inp.grad is not None and torch.isfinite(inp.grad).all() + assert layer.weight.grad is not None and torch.isfinite(layer.weight.grad).all() + recompute_buffer = FP8GlobalStateManager.quantization_state.fp8_tensors_recompute_buffer + assert all(len(stashed) == 0 for stashed in recompute_buffer) + return out.detach(), inp.grad.detach(), layer.weight.grad.detach(), observed + + ref_out, ref_dgrad, ref_wgrad, ref_observed = run(checkpointed=False) + out, dgrad, wgrad, observed = run(checkpointed=True) + + torch.testing.assert_close(out, ref_out) + torch.testing.assert_close(dgrad, ref_dgrad) + torch.testing.assert_close(wgrad, ref_wgrad) + assert ref_observed == [(True, False, False)] + assert observed == ref_observed + + def _test_e2e_checkpointing_get_model(config, dtype): sigma = 0.023 init_method = init_method_normal(sigma)