diff --git a/tests/pytorch/test_fp8_checkpoint_state_edges.py b/tests/pytorch/test_fp8_checkpoint_state_edges.py new file mode 100644 index 0000000000..a72164fee4 --- /dev/null +++ b/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) diff --git a/tests/pytorch/test_reentrant_fp8_updates.py b/tests/pytorch/test_reentrant_fp8_updates.py new file mode 100644 index 0000000000..a6eca97b65 --- /dev/null +++ b/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) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index c050f26869..50ace56f50 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -238,6 +238,13 @@ def gather_split_1d_tensor(tensor: torch.Tensor, tp_group: dist_group_type) -> t return gathered +@dataclass +class _ActivationRecomputeState: + """FP8 state shared by a checkpoint's original forward and recompute.""" + + is_first_fp8_module: Optional[bool] = None + + class activation_recompute_forward(AbstractContextManager, ContextDecorator): """Context manager used to control the forward runtime behavior when executed under the `CheckpointFunction` function. For running FP8, the forward pass will @@ -247,30 +254,74 @@ class activation_recompute_forward(AbstractContextManager, ContextDecorator): activations, followed by calculation of gradients using these values. """ - _is_first_fp8_module: List = [] - - def __init__(self, activation_recompute: bool = False, recompute_phase: bool = False): + def __init__( + self, + activation_recompute: bool = False, + recompute_phase: bool = False, + state: Optional[_ActivationRecomputeState] = None, + reserve_first_fp8_module: bool = False, + ): super().__init__() self.activation_recompute = activation_recompute self.recompute_phase = recompute_phase + self.state = _ActivationRecomputeState() if state is None else state + self.reserve_first_fp8_module = reserve_first_fp8_module def __enter__(self): global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE - _FP8_ACTIVATION_RECOMPUTE_ENABLED = ( - self.activation_recompute and FP8GlobalStateManager.is_fp8_enabled() - ) - _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase - qstate = FP8GlobalStateManager.quantization_state - if self.activation_recompute and not self.recompute_phase: - activation_recompute_forward._is_first_fp8_module.append(qstate.is_first_fp8_module) - if self.activation_recompute and self.recompute_phase: - qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0) + fp8_enabled = FP8GlobalStateManager.is_fp8_enabled() + + # Validate before mutating process-global state. A failed __enter__ is not + # followed by __exit__, so validation after committing would leak flags. + if ( + self.activation_recompute + and self.recompute_phase + and self.state.is_first_fp8_module is None + ): + raise RuntimeError("FP8 recompute state was not captured during forward") + + self._previous_recompute_enabled = _FP8_ACTIVATION_RECOMPUTE_ENABLED + self._previous_recompute_phase = _FP8_ACTIVATION_RECOMPUTE_PHASE + self._previous_is_first_fp8_module = qstate.is_first_fp8_module + try: + _FP8_ACTIVATION_RECOMPUTE_ENABLED = self.activation_recompute and fp8_enabled + _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase + + if self.activation_recompute and not self.recompute_phase: + self.state.is_first_fp8_module = qstate.is_first_fp8_module + # Reentrant checkpoint forward runs under no_grad, so no module can + # consume backward-update ownership. Reserve it for this frame. + if self.reserve_first_fp8_module and fp8_enabled: + qstate.is_first_fp8_module = False + elif self.activation_recompute and self.recompute_phase: + qstate.is_first_fp8_module = self.state.is_first_fp8_module + except BaseException: + qstate.is_first_fp8_module = self._previous_is_first_fp8_module + _FP8_ACTIVATION_RECOMPUTE_ENABLED = self._previous_recompute_enabled + _FP8_ACTIVATION_RECOMPUTE_PHASE = self._previous_recompute_phase + raise + return self def __exit__(self, *exc_details): global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE - _FP8_ACTIVATION_RECOMPUTE_ENABLED = False - _FP8_ACTIVATION_RECOMPUTE_PHASE = False + qstate = FP8GlobalStateManager.quantization_state + try: + if self.activation_recompute and self.recompute_phase: + # Consumption during recompute is monotonic. Merge it with the + # outer frame instead of restoring a stale available snapshot. + qstate.is_first_fp8_module = ( + self._previous_is_first_fp8_module and qstate.is_first_fp8_module + ) + elif self.activation_recompute and exc_details and exc_details[0] is not None: + # A failed original forward creates no usable checkpoint frame. Roll + # back ownership so a recovery module in the same outer autocast can + # become the first backward-update owner. + qstate.is_first_fp8_module = self._previous_is_first_fp8_module + self.state.is_first_fp8_module = None + finally: + _FP8_ACTIVATION_RECOMPUTE_ENABLED = self._previous_recompute_enabled + _FP8_ACTIVATION_RECOMPUTE_PHASE = self._previous_recompute_phase def is_fp8_activation_recompute_enabled() -> bool: @@ -369,8 +420,14 @@ def forward( # Preserve torch autocast context for the backward pass torch_gpu_amp_ctx, torch_cpu_amp_ctx = _get_active_autocast_contexts() + fp8_recompute_state = _ActivationRecomputeState() with torch.no_grad(), forward_ctx: - with activation_recompute_forward(activation_recompute=True, recompute_phase=False): + with activation_recompute_forward( + activation_recompute=True, + recompute_phase=False, + state=fp8_recompute_state, + reserve_first_fp8_module=True, + ): outputs = run_function(*args, **kwargs) # Divide hidden states across model parallel group and only keep @@ -395,6 +452,7 @@ def forward( ctx.torch_cpu_amp_ctx = torch_cpu_amp_ctx ctx.fp8 = fp8 ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + ctx.fp8_recompute_state = fp8_recompute_state ctx.kwargs = kwargs return outputs @@ -436,9 +494,13 @@ def backward( # Compute the forward pass. detached_inputs = detach_variable(inputs) with torch.enable_grad(), ctx.recompute_ctx, ctx.torch_gpu_amp_ctx, ctx.torch_cpu_amp_ctx, activation_recompute_forward( - activation_recompute=True, recompute_phase=True + activation_recompute=True, + recompute_phase=True, + state=ctx.fp8_recompute_state, ), autocast( - enabled=ctx.fp8, recipe=ctx.fp8_recipe + enabled=ctx.fp8, + recipe=ctx.fp8_recipe, + _recompute=True, ): outputs = ctx.run_function(*detached_inputs, **ctx.kwargs) @@ -602,13 +664,16 @@ def use_reentrant_activation_recompute(): def get_activation_recompute_contexts(): """Returns context objects for the checkpointed forward pass and the forward recompute phase.""" + state = _ActivationRecomputeState() forward_ctx = activation_recompute_forward( activation_recompute=True, recompute_phase=False, + state=state, ) recompute_ctx = activation_recompute_forward( activation_recompute=True, recompute_phase=True, + state=state, ) return forward_ctx, recompute_ctx @@ -795,7 +860,9 @@ def recompute_fn(*args, **kwargs): with torch.autograd.enable_grad(), ( te_recompute_ctx ), user_recompute_ctx, torch_gpu_amp_forward_ctx, torch_cpu_amp_forward_ctx, autocast( - enabled=fp8, recipe=fp8_recipe + enabled=fp8, + recipe=fp8_recipe, + _recompute=True, ): function(*args, **kwargs) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 43799b003c..715782e30b 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -52,7 +52,6 @@ symmetric_all_reduce, reduce_scatter_along_first_dim, gather_along_first_dim, - in_fp8_activation_recompute_phase, _fsdp_scatter_tensors, _fsdp_gather_tensors, ) @@ -565,11 +564,7 @@ def forward( ctx.normalization = normalization ctx.reduce_and_update_bwd_fp8_tensors = False if ctx.fp8 and requires_grad(inp, ln_weight, ln_bias, weight, bias): - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): - qstate.is_first_fp8_module = _first_fp8_module ctx.wgrad_store = wgrad_store ctx.debug = debug diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ad629c3f11..f614d1cae3 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -58,7 +58,6 @@ reduce_scatter_along_first_dim, gather_along_first_dim, use_reentrant_activation_recompute, - in_fp8_activation_recompute_phase, _fsdp_scatter_tensors, _get_cuda_rng_state, _set_cuda_rng_state, @@ -908,7 +907,7 @@ def _forward( qstate = FP8GlobalStateManager.quantization_state _first_fp8_module = qstate.is_first_fp8_module ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase() or is_recomputation: + if is_recomputation: qstate.is_first_fp8_module = _first_fp8_module ctx.wgrad_store = wgrad_store diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 78a4d31852..f5abb5eb93 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -245,12 +245,7 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: def _check_fp8_reduce_and_update(): """Check if this is the first FP8 module (for backward reduce-and-update).""" - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module - result = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): - qstate.is_first_fp8_module = _first_fp8_module - return result + return FP8GlobalStateManager.is_first_fp8_module() def _linear_forward_impl( diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index f1bffb122f..29aa59d692 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -724,6 +724,7 @@ def autocast_enter( fp8_recipe: Optional[Recipe] = None, fp8_group: Optional[dist_group_type] = None, _graph: bool = False, + _recompute: bool = False, ) -> None: """Set state and tracking variables for entry into FP8 region.""" @@ -741,7 +742,7 @@ def autocast_enter( qstate.fp8_distributed_group = fp8_group qstate.fp8_graph_capturing = _graph - if qstate.autocast_depth == 0: + if qstate.autocast_depth == 0 and not _recompute: qstate.is_first_fp8_module = True qstate.autocast_depth += 1 @@ -759,14 +760,20 @@ def autocast_enter( assert nvfp4_available, reason_for_no_nvfp4 @classmethod - def autocast_exit(cls, enabled: bool, _graph: bool) -> None: + def autocast_exit(cls, enabled: bool, _graph: bool, _recompute: bool = False) -> None: """Set state and tracking variables for exit from FP8 region.""" qstate = cls.quantization_state qstate.autocast_depth -= 1 # Reduce only the non-FP8 weight modules here. # FP8 weight modules are reduced at the end of the optimizer # step after the weight amax is populated. - if enabled and qstate.autocast_depth == 0 and not _graph and torch.is_grad_enabled(): + if ( + enabled + and qstate.autocast_depth == 0 + and not _graph + and not _recompute + and torch.is_grad_enabled() + ): # delayed scaling only function, for other recipes (current scaling with any granularity), # this is noop for other recipes because cls.global_amax_buffer is empty list cls.reduce_and_update_fp8_tensors(forward=True) @@ -1006,6 +1013,7 @@ class autocast: "_recipe", "_amax_reduction_group", "_graph", + "_recompute", "_fp8_state", ) @@ -1016,12 +1024,14 @@ def __init__( recipe: Optional["Recipe"] = None, amax_reduction_group: Optional["dist_group_type"] = None, _graph: bool = False, + _recompute: bool = False, ) -> None: self._enabled = enabled self._calibrating = calibrating self._recipe = recipe self._amax_reduction_group = amax_reduction_group self._graph = _graph + self._recompute = _recompute self._fp8_state = None def __enter__(self) -> "autocast": @@ -1040,13 +1050,27 @@ def __enter__(self) -> "autocast": fp8_recipe=self._recipe, fp8_group=self._amax_reduction_group, _graph=self._graph, + _recompute=self._recompute, ) return self def __exit__(self, exc_type, exc_val, exc_tb) -> None: try: + # Restoring a nested autocast's configuration must not revive backward-update + # ownership that was consumed while the nested scope was active. Ownership is + # monotonic within the logical outer autocast: available -> consumed. + qstate = FP8GlobalStateManager.quantization_state + nested_autocast = qstate.autocast_depth > 1 + live_is_first_fp8_module = qstate.is_first_fp8_module FP8GlobalStateManager.set_autocast_state(self._fp8_state) - FP8GlobalStateManager.autocast_exit(self._enabled, _graph=self._graph) + if nested_autocast: + qstate = FP8GlobalStateManager.quantization_state + qstate.is_first_fp8_module = qstate.is_first_fp8_module and live_is_first_fp8_module + FP8GlobalStateManager.autocast_exit( + self._enabled, + _graph=self._graph, + _recompute=self._recompute, + ) finally: # Clear the saved state so the instance can be entered again # sequentially (and so a failure inside the restore path does not