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 8605a4746b..a5720ede7a 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -238,42 +238,100 @@ def gather_split_1d_tensor(tensor: torch.Tensor, tp_group: dist_group_type) -> t return gathered -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 - run without storing intermediate activations. Instead, the forward pass saves - the inputs tuple and the calling function. In the backwards pass, these are - retrieved, and the forward pass is computed again while tracking the intermediate - activations, followed by calculation of gradients using these values. - """ +@dataclass +class _ActivationRecomputeState: + """Ownership shared by the original and recomputed checkpoint forwards.""" + + is_first_fp8_module: Optional[bool] = None + forward_completed: bool = False - _is_first_fp8_module: List = [] - def __init__(self, activation_recompute: bool = False, recompute_phase: bool = False): +class activation_recompute_forward(AbstractContextManager, ContextDecorator): + """Context manager for FP8 checkpoint forward/recompute bookkeeping.""" + + 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 + if state is not None: + self.state = state + elif activation_recompute and not recompute_phase: + self.state = _ActivationRecomputeState() + else: + self.state = None + self.reserve_first_fp8_module = reserve_first_fp8_module + self._previous_region = False + self._previous_phase = False + self._previous_is_first_fp8_module = False def __enter__(self): 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 - _ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase - + self._previous_region = _IN_ACTIVATION_RECOMPUTE_REGION + self._previous_phase = _ACTIVATION_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) + self._previous_is_first_fp8_module = qstate.is_first_fp8_module + fp8_enabled = FP8GlobalStateManager.is_fp8_enabled() + + try: + # Track checkpoint regions independently of whether FP8 was enabled + # when the context was entered; a checkpointed callable may open an + # inner FP8 autocast context. + _IN_ACTIVATION_RECOMPUTE_REGION = self.activation_recompute + _ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase + + if not self.activation_recompute: + return self + if self.recompute_phase: + if ( + self.state is None + or self.state.is_first_fp8_module is None + or not self.state.forward_completed + ): + raise RuntimeError( + "FP8 recompute state was not captured during a completed original forward" + ) + qstate.is_first_fp8_module = self.state.is_first_fp8_module + else: + if self.state is None: + self.state = _ActivationRecomputeState() + self.state.is_first_fp8_module = qstate.is_first_fp8_module + self.state.forward_completed = False + if self.reserve_first_fp8_module and fp8_enabled: + qstate.is_first_fp8_module = False + except BaseException: + qstate.is_first_fp8_module = self._previous_is_first_fp8_module + _IN_ACTIVATION_RECOMPUTE_REGION = self._previous_region + _ACTIVATION_RECOMPUTE_PHASE = self._previous_phase + raise + return self def __exit__(self, *exc_details): global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE - _IN_ACTIVATION_RECOMPUTE_REGION = False - _ACTIVATION_RECOMPUTE_PHASE = False + exc_type = exc_details[0] if exc_details else None + qstate = FP8GlobalStateManager.quantization_state + + # Recompute consumes a reservation from the original forward and must + # restore the caller's ownership. An exception in the original forward + # must also roll back the reservation so no stale frame survives. + if self.activation_recompute and (self.recompute_phase or exc_type is not None): + qstate.is_first_fp8_module = self._previous_is_first_fp8_module + if self.activation_recompute and not self.recompute_phase and exc_type is not None: + if self.state is not None: + self.state.is_first_fp8_module = None + self.state.forward_completed = False + elif self.activation_recompute and not self.recompute_phase: + if self.state is not None: + self.state.forward_completed = True + + _IN_ACTIVATION_RECOMPUTE_REGION = self._previous_region + _ACTIVATION_RECOMPUTE_PHASE = self._previous_phase + return False def is_fp8_activation_recompute_enabled() -> bool: @@ -372,8 +430,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 @@ -398,6 +462,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 @@ -439,9 +504,11 @@ 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) @@ -604,14 +671,17 @@ def use_reentrant_activation_recompute(): def get_activation_recompute_contexts(): - """Returns context objects for the checkpointed forward pass and the forward recompute phase.""" + """Returns contexts sharing FP8 ownership for forward and recompute.""" + 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 @@ -798,7 +868,7 @@ 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 561e813348..74eaf64f1f 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, ) @@ -590,11 +589,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 3ee0cda50c..ba87621201 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, @@ -915,7 +914,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 56622db5e6..abfa64c0cb 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -257,12 +257,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 98c67be922..c91de7d5ef 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -632,6 +632,10 @@ def get_autocast_state(cls) -> tuple: def set_autocast_state(cls, state: tuple) -> None: """Restore a previously saved autocast state snapshot.""" qstate = cls.quantization_state + # Ownership is consumed monotonically by FP8 modules. Restoring an + # outer autocast snapshot must not revive an already-consumed owner + # when a nested autocast context exits. + current_is_first_fp8_module = qstate.is_first_fp8_module ( qstate.fp8_enabled, qstate.fp8_calibration, @@ -640,6 +644,7 @@ def set_autocast_state(cls, state: tuple) -> None: qstate.is_first_fp8_module, qstate.fp8_graph_capturing, ) = state + qstate.is_first_fp8_module = current_is_first_fp8_module and qstate.is_first_fp8_module @staticmethod def reduce_tensor_across_group_op_max(tensor: torch.Tensor, group: dist_group_type) -> None: @@ -734,6 +739,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.""" @@ -751,7 +757,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 @@ -769,14 +775,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) @@ -1016,6 +1028,7 @@ class autocast: "_recipe", "_amax_reduction_group", "_graph", + "_recompute", "_fp8_state", ) @@ -1026,12 +1039,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": @@ -1050,13 +1065,18 @@ 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: FP8GlobalStateManager.set_autocast_state(self._fp8_state) - FP8GlobalStateManager.autocast_exit(self._enabled, _graph=self._graph) + 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