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..11cd15ba98 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,50 @@ 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() - ) + self._previous_recompute_enabled = _FP8_ACTIVATION_RECOMPUTE_ENABLED + self._previous_recompute_phase = _FP8_ACTIVATION_RECOMPUTE_PHASE + fp8_enabled = FP8GlobalStateManager.is_fp8_enabled() + _FP8_ACTIVATION_RECOMPUTE_ENABLED = self.activation_recompute and fp8_enabled _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase qstate = FP8GlobalStateManager.quantization_state + self._previous_is_first_fp8_module = qstate.is_first_fp8_module if self.activation_recompute and not self.recompute_phase: - activation_recompute_forward._is_first_fp8_module.append(qstate.is_first_fp8_module) + 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 if self.activation_recompute and self.recompute_phase: - qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0) + if self.state.is_first_fp8_module is None: + raise RuntimeError("FP8 recompute state was not captured during forward") + qstate.is_first_fp8_module = self.state.is_first_fp8_module + 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 + if self.activation_recompute and self.recompute_phase: + # Only recompute restores outer ownership. After the original forward, it must remain + # consumed or reserved so later checkpoint frames cannot claim it. + qstate = FP8GlobalStateManager.quantization_state + 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 def is_fp8_activation_recompute_enabled() -> bool: @@ -369,8 +396,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 +428,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 +470,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 +640,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 +836,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..45083e5726 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,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