From 5a531701bab853e960972caf7f4a77778ce0385c Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 17 Aug 2026 12:31:09 +0200 Subject: [PATCH 1/7] [PyTorch] Replace first-module ownership with deferred backward quantization state update The backward amax reduction for delayed scaling was triggered from the backward of the 'first FP8 module' of an autocast, assuming its backward runs last. That assumption breaks with activation recompute (duplicate forward/backward updates per step) and with branched autograd graphs (ownership can land in a frame whose backward never runs). Instead, module backwards now only request the update (schedule_backward_quantization_update); it is flushed exactly once at the next top-level autocast entry, after the backward pass is complete. Forward updates at autocast exit are skipped during the recompute phase. Removes the is_first_fp8_module save/restore bookkeeping from checkpoint contexts and modules; renames reduce_and_update_fp8_tensors to reduce_and_update_quantization_state (old name kept as alias). Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_backward_override.py | 8 +- .../test_quantization_state_updates.py | 270 ++++++++++++++++++ transformer_engine/pytorch/distributed.py | 29 +- .../pytorch/module/grouped_linear.py | 29 +- .../pytorch/module/layernorm_linear.py | 19 +- .../pytorch/module/layernorm_mlp.py | 15 +- transformer_engine/pytorch/module/linear.py | 26 +- transformer_engine/pytorch/ops/fuser.py | 15 +- transformer_engine/pytorch/quantization.py | 56 +++- 9 files changed, 374 insertions(+), 93 deletions(-) create mode 100644 tests/pytorch/test_quantization_state_updates.py diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index c0acf2e6b3..feff7bccdb 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -419,7 +419,7 @@ def _snapshot_backward_ctx_state( "backward_override", "fp8", "grad_output_quantizer", - "reduce_and_update_bwd_fp8_tensors", + "schedule_backward_quantization_update", ) missing_attrs = [attr for attr in required_attrs if not hasattr(state_holder, attr)] if missing_attrs: @@ -430,7 +430,7 @@ def _snapshot_backward_ctx_state( getattr(state_holder, "backward_override"), bool(getattr(state_holder, "fp8")), getattr(state_holder, "grad_output_quantizer"), - bool(getattr(state_holder, "reduce_and_update_bwd_fp8_tensors")), + bool(getattr(state_holder, "schedule_backward_quantization_update")), ) @@ -816,7 +816,7 @@ def _run_grouped_linear_single_step_with_ctx_state( required_attrs = ( "backward_override", "fp8", - "reduce_and_update_bwd_fp8_tensors", + "schedule_backward_quantization_update", ) missing_attrs = [attr for attr in required_attrs if not hasattr(y.grad_fn, attr)] if missing_attrs: @@ -827,7 +827,7 @@ def _run_grouped_linear_single_step_with_ctx_state( ctx_state = ( getattr(y.grad_fn, "backward_override"), bool(getattr(y.grad_fn, "fp8")), - bool(getattr(y.grad_fn, "reduce_and_update_bwd_fp8_tensors")), + bool(getattr(y.grad_fn, "schedule_backward_quantization_update")), ) y.backward(dy) assert x_run.grad is not None diff --git a/tests/pytorch/test_quantization_state_updates.py b/tests/pytorch/test_quantization_state_updates.py new file mode 100644 index 0000000000..d088b0b7b7 --- /dev/null +++ b/tests/pytorch/test_quantization_state_updates.py @@ -0,0 +1,270 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for delayed-scaling quantization state update scheduling. + +The backward amax reduction must run exactly once per training step, +regardless of activation checkpointing (which re-runs forwards) or the +topology of the autograd graph (branches, unused outputs). +""" + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import DelayedScaling +from transformer_engine.pytorch.distributed import checkpoint as te_checkpoint +from transformer_engine.pytorch.quantization import FP8GlobalStateManager + +fp8_available, reason_for_no_fp8 = FP8GlobalStateManager.is_fp8_available() + +HIDDEN = 128 +BATCH = 32 +STEPS = 3 + + +class _UpdateCounter: + """Counts reduce_and_update_quantization_state calls by direction.""" + + def __init__(self): + self.forward = 0 + self.backward = 0 + self._original = None + + def __enter__(self): + self._original = FP8GlobalStateManager.reduce_and_update_quantization_state.__func__ + original = self._original + counter = self + + def counted(cls, forward=True): + if forward: + counter.forward += 1 + else: + counter.backward += 1 + return original(cls, forward=forward) + + FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(counted) + return self + + def __exit__(self, *exc): + FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(self._original) + + +def _make_model(num_layers=3, seed=1234): + torch.manual_seed(seed) + return torch.nn.ModuleList( + [te.Linear(HIDDEN, HIDDEN, bias=True).cuda() for _ in range(num_layers)] + ) + + +def _run_layers(layers, x): + for layer in layers: + x = layer(x) + return x + + +def _train_step(model, x, forward_fn, recipe): + with te.autocast(enabled=True, recipe=recipe): + out = forward_fn(model, x) + loss = out.float().sum() + loss.backward() + return loss + + +def _forward_plain(model, x): + return _run_layers(model, x) + + +def _forward_reentrant(model, x): + return te_checkpoint(_run_layers, model, x, use_reentrant=True) + + +def _forward_non_reentrant(model, x): + return te_checkpoint(_run_layers, model, x, use_reentrant=False) + + +def _forward_per_layer_reentrant(model, x): + for layer in model: + x = te_checkpoint(layer, x, use_reentrant=True) + return x + + +def _forward_per_layer_non_reentrant(model, x): + for layer in model: + x = te_checkpoint(layer, x, use_reentrant=False) + return x + + +def _forward_nested(model, x): + def inner(x): + return te_checkpoint(model[1], x, use_reentrant=True) + + def outer(x): + x = model[0](x) + x = inner(x) + return model[2](x) + + return te_checkpoint(outer, x, use_reentrant=True) + + +def _forward_torch_checkpoint(model, x): + return torch.utils.checkpoint.checkpoint(_run_layers, model, x, use_reentrant=False) + + +FORWARD_FNS = { + "plain": _forward_plain, + "reentrant": _forward_reentrant, + "non_reentrant": _forward_non_reentrant, + "per_layer_reentrant": _forward_per_layer_reentrant, + "per_layer_non_reentrant": _forward_per_layer_non_reentrant, + "nested": _forward_nested, + "torch_checkpoint": _forward_torch_checkpoint, +} + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("mode", FORWARD_FNS.keys()) +def test_single_update_per_step(mode): + """Exactly one forward and one backward state update per training step.""" + forward_fn = FORWARD_FNS[mode] + model = _make_model() + recipe = DelayedScaling() + + with _UpdateCounter() as counter: + for step in range(STEPS): + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + _train_step(model, x, forward_fn, recipe) + # The backward update of this step is flushed at the next + # top-level autocast entry; flush explicitly to count it here. + FP8GlobalStateManager.flush_backward_quantization_update() + assert counter.forward == step + 1, f"{mode}: duplicate/missing forward update" + assert counter.backward == step + 1, f"{mode}: duplicate/missing backward update" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_flush_at_next_autocast_entry(): + """Without an explicit flush, the update runs when the next autocast begins.""" + model = _make_model() + recipe = DelayedScaling() + + with _UpdateCounter() as counter: + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + _train_step(model, x, _forward_plain, recipe) + assert counter.backward == 0 + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + _train_step(model, x, _forward_plain, recipe) + assert counter.backward == 1 + FP8GlobalStateManager.flush_backward_quantization_update() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("checkpoint_first_branch", [True, False]) +def test_branched_graph(checkpoint_first_branch): + """Two checkpointed sibling branches merging into one loss.""" + branch_a = _make_model(num_layers=2, seed=1) + branch_b = _make_model(num_layers=2, seed=2) + recipe = DelayedScaling() + + with _UpdateCounter() as counter: + for step in range(STEPS): + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + with te.autocast(enabled=True, recipe=recipe): + if checkpoint_first_branch: + ya = te_checkpoint(_run_layers, branch_a, x, use_reentrant=True) + else: + ya = _run_layers(branch_a, x) + yb = te_checkpoint(_run_layers, branch_b, x, use_reentrant=True) + out = ya + yb + loss = out.float().sum() + loss.backward() + FP8GlobalStateManager.flush_backward_quantization_update() + assert counter.forward == step + 1 + assert counter.backward == step + 1 + assert x.grad is not None and torch.isfinite(x.grad).all() + x.grad = None + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_checkpointed_branch_without_backward(): + """A checkpointed branch whose output never receives a gradient must not + strand the backward update (regression test for first-module ownership).""" + used = _make_model(num_layers=2, seed=1) + unused = _make_model(num_layers=2, seed=2) + recipe = DelayedScaling() + + with _UpdateCounter() as counter: + for step in range(STEPS): + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + with te.autocast(enabled=True, recipe=recipe): + # The unused branch runs first, so any "first module owns the + # backward update" scheme would assign ownership to a frame + # whose backward never executes. + y_unused = te_checkpoint(_run_layers, unused, x, use_reentrant=True) + y = _run_layers(used, x) + loss = y.float().sum() + loss.backward() + del y_unused + FP8GlobalStateManager.flush_backward_quantization_update() + assert counter.backward == step + 1, "backward update was lost" + x.grad = None + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize( + "mode", ["reentrant", "non_reentrant", "per_layer_reentrant", "nested", "torch_checkpoint"] +) +def test_checkpointing_matches_plain_numerics(mode): + """Delayed-scaling state must evolve identically with and without + activation checkpointing (duplicate updates would advance it faster).""" + forward_fn = FORWARD_FNS[mode] + recipe = DelayedScaling() + + def train(forward): + model = _make_model(seed=99) + losses = [] + torch.manual_seed(777) + for _ in range(STEPS + 1): + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + losses.append(_train_step(model, x, forward, recipe).detach()) + FP8GlobalStateManager.flush_backward_quantization_update() + state = [] + for layer in model: + for key in ("scaling_fwd", "scaling_bwd"): + meta = layer.fp8_meta[key] + state.append((meta.scale.clone(), meta.amax_history.clone())) + return losses, state + + losses_ref, state_ref = train(_forward_plain) + losses_ckpt, state_ckpt = train(forward_fn) + + for step, (l_ref, l_ckpt) in enumerate(zip(losses_ref, losses_ckpt)): + torch.testing.assert_close(l_ckpt, l_ref, rtol=0, atol=0, msg=f"loss diverged @ {step}") + for (scale_ref, hist_ref), (scale_ckpt, hist_ckpt) in zip(state_ref, state_ckpt): + assert torch.equal(scale_ckpt, scale_ref), "scale diverged under checkpointing" + assert torch.equal(hist_ckpt, hist_ref), "amax history diverged under checkpointing" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_exception_in_checkpointed_forward(): + """A failing checkpointed forward must not corrupt update scheduling.""" + model = _make_model() + recipe = DelayedScaling() + + def failing(model, x): + def fn(x): + model[0](x) + raise RuntimeError("boom") + + return te_checkpoint(fn, x, use_reentrant=True) + + with _UpdateCounter() as counter: + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + with pytest.raises(RuntimeError, match="boom"): + _train_step(model, x, failing, recipe) + + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + _train_step(model, x, _forward_reentrant, recipe) + FP8GlobalStateManager.flush_backward_quantization_update() + assert counter.backward == 1 + assert torch.isfinite(x.grad).all() diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index d1525b53f0..da1b05cf05 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -62,9 +62,6 @@ _USE_REENTRANT_ACTIVATION_RECOMPUTE = True -_FP8_ACTIVATION_RECOMPUTE_ENABLED = False -_FP8_ACTIVATION_RECOMPUTE_PHASE = False - _ALL_ACTIVE_RNG_STATES = {} @@ -247,40 +244,34 @@ 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): super().__init__() self.activation_recompute = activation_recompute self.recompute_phase = recompute_phase def __enter__(self): - global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE - _FP8_ACTIVATION_RECOMPUTE_ENABLED = ( + qstate = FP8GlobalStateManager.quantization_state + self._prev_enabled = qstate.activation_recompute_enabled + self._prev_phase = qstate.activation_recompute_phase + qstate.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) + qstate.activation_recompute_phase = self.recompute_phase 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 + qstate.activation_recompute_enabled = self._prev_enabled + qstate.activation_recompute_phase = self._prev_phase def is_fp8_activation_recompute_enabled() -> bool: """Return global boolean""" - return _FP8_ACTIVATION_RECOMPUTE_ENABLED + return FP8GlobalStateManager.quantization_state.activation_recompute_enabled def in_fp8_activation_recompute_phase() -> bool: """Return global boolean""" - return _FP8_ACTIVATION_RECOMPUTE_PHASE + return FP8GlobalStateManager.quantization_state.activation_recompute_phase def _get_active_autocast_contexts(): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 12497d0cb0..2b0e562f71 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -49,6 +49,7 @@ is_fp8_activation_recompute_enabled, in_fp8_activation_recompute_phase, ) +from ..graph import is_graph_capturing from ..distributed_weight import ( is_distributed_weight, materialize_weight_for_forward, @@ -833,12 +834,9 @@ def _forward_grouped_tensor( ctx.use_bias = use_bias ctx.inp_shape = inp.shape ctx.requires_dgrad = inp.requires_grad - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad(inp, weights[0], biases[0]): - ctx.reduce_and_update_bwd_fp8_tensors = ( - ctx.reduce_and_update_bwd_fp8_tensors - or FP8GlobalStateManager.is_first_fp8_module() - ) + ctx.schedule_backward_quantization_update = ctx.fp8 and requires_grad( + inp, weights[0], biases[0] + ) ctx.wgrad_store = wgrad_store ctx.debug = False ctx.save_original_input = False @@ -1177,12 +1175,9 @@ def forward( ctx.sequence_parallel = sequence_parallel ctx.inp_shape = inp.shape ctx.requires_dgrad = inp.requires_grad - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad(inp, weights[0], biases[0]): - ctx.reduce_and_update_bwd_fp8_tensors = ( - ctx.reduce_and_update_bwd_fp8_tensors - or FP8GlobalStateManager.is_first_fp8_module() - ) + ctx.schedule_backward_quantization_update = ctx.fp8 and requires_grad( + inp, weights[0], biases[0] + ) ctx.wgrad_store = wgrad_store ctx.debug = debug ctx.save_original_input = save_original_input @@ -1200,7 +1195,7 @@ def forward( ctx.grad_input_quantizers = [None] * num_gemms ctx.grad_weight_quantizers = [None] * num_gemms ctx.grad_output_quantizers = [None] * num_gemms - ctx.reduce_and_update_bwd_fp8_tensors = False + ctx.schedule_backward_quantization_update = False # [*, in_features] -> [*, out_features] except first dimension changes for SP return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @@ -1388,8 +1383,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): if not ctx.use_bias: grad_biases = [None] * N - if ctx.reduce_and_update_bwd_fp8_tensors: - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) + if ctx.schedule_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.schedule_backward_quantization_update() return ( dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits @@ -1652,8 +1647,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): ): grad_biases = [None] * ctx.num_gemms - if ctx.reduce_and_update_bwd_fp8_tensors: - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) + if ctx.schedule_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.schedule_backward_quantization_update() return ( dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 561e813348..5229464121 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, ) @@ -588,13 +587,9 @@ def forward( ctx.ub_name = ub_name ctx.requires_dgrad = inp_requires_grad 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.schedule_backward_quantization_update = ctx.fp8 and requires_grad( + inp, ln_weight, ln_bias, weight, bias + ) ctx.wgrad_store = wgrad_store ctx.debug = debug @@ -609,7 +604,7 @@ def forward( ctx.grad_input_quantizer = None ctx.grad_weight_quantizer = None ctx.grad_output_quantizer = None - ctx.reduce_and_update_bwd_fp8_tensors = False + ctx.schedule_backward_quantization_update = False # ------------------------------------------------------ # Cached state for backward pass is ready... @@ -1184,10 +1179,8 @@ def wgrad_gemm( else: wgrad = None - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): - nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors") - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) - nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors") + if ctx.schedule_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.schedule_backward_quantization_update() # Scatter fp8 weight buffers # if ctx.fp8 and not isinstance(weight, QuantizedTensorStorage): diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 3ee0cda50c..478df76e13 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,15 +907,9 @@ def _forward( inp.requires_grad or ln_weight.requires_grad or ln_bias.requires_grad ) ctx.normalization = normalization - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.fp8 and requires_grad( + ctx.schedule_backward_quantization_update = ctx.fp8 and requires_grad( inp, ln_weight, ln_bias, fc1_weight, fc2_weight, fc1_bias, fc2_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() or is_recomputation: - qstate.is_first_fp8_module = _first_fp8_module + ) ctx.wgrad_store = wgrad_store if is_recomputation: # return the recomputed tensors @@ -1806,8 +1799,8 @@ def fc1_wgrad_gemm( else: fc2_wgrad = None - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) + if ctx.schedule_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.schedule_backward_quantization_update() # FIX THIS # Scatter Fp8 tranposed-weight buffers diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 56622db5e6..daad50ffbe 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -233,8 +233,8 @@ class LinearBwdArgs: origin_weight_overwrites_main_grad: bool = False main_grad_func: Optional[Callable[[], torch.Tensor]] = None - # --- FP8 reduce-and-update bookkeeping --- - reduce_and_update_bwd_fp8_tensors: bool = False + # --- Quantization state update bookkeeping --- + schedule_backward_quantization_update: bool = False # --- Misc --- cpu_offloading: bool = False @@ -255,16 +255,6 @@ def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: ) # pylint: disable=unbalanced-tuple-unpacking -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 - - def _linear_forward_impl( args: LinearFwdArgs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], None, Optional[Dict]]: @@ -1439,9 +1429,9 @@ def forward( or fwd_args.weight_requires_grad or fwd_args.bias_requires_grad ): - bwd_args.reduce_and_update_bwd_fp8_tensors = _check_fp8_reduce_and_update() + bwd_args.schedule_backward_quantization_update = True if fwd_args.backward_override is not None: - bwd_args.reduce_and_update_bwd_fp8_tensors = False + bwd_args.schedule_backward_quantization_update = False return out, new_weight_workspace @@ -1459,15 +1449,13 @@ def backward( if bwd_args.ub_name is not None: nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" result = _linear_backward(bwd_args) + (None,) # fwd_args grad slot - reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors + schedule_backward_quantization_update = bwd_args.schedule_backward_quantization_update # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, # main_grad closure) so they don't outlive backward via ctx under retain_graph. ctx.backward_objects = None del bwd_args - if reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): - nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors") - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) - nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors") + if schedule_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.schedule_backward_quantization_update() return result diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 09ffb004dd..92e81c838a 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -176,10 +176,11 @@ def forward( func_ctx.save_for_backward(*tensors_to_save) func_ctx.tensor_objects = tensor_objects - # Whether to perform recipe update in backward pass - is_first_module = False - if fuser.first_op_requiring_backward < fuser._num_basic_ops: - is_first_module = FP8GlobalStateManager.is_first_fp8_module() + # Whether to request a recipe update in backward pass + schedule_backward_quantization_update = ( + fuser.first_op_requiring_backward < fuser._num_basic_ops + and FP8GlobalStateManager.is_fp8_enabled() + ) # Other context func_ctx.backward_ops = fuser._backward_ops @@ -188,7 +189,7 @@ def forward( func_ctx.basic_op_num_params = fuser._basic_op_num_params func_ctx.num_extra_inputs = fuser.num_extra_inputs func_ctx.num_extra_outputs = len(extra_outputs_flat) - func_ctx.is_first_module = is_first_module + func_ctx.schedule_backward_quantization_update = schedule_backward_quantization_update # Mark output tensors as not deletable in backward for tensor in [x] + extra_outputs_flat: @@ -291,8 +292,8 @@ def backward( grad_extra_inputs_flat.extend(dxs) # Update FP8 scaling factors - if func_ctx.is_first_module and not _is_graph_capturing(): - FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) + if func_ctx.schedule_backward_quantization_update and not _is_graph_capturing(): + FP8GlobalStateManager.schedule_backward_quantization_update() return ( dx, # input_ diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 07a3b80483..be61364af8 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -400,6 +400,9 @@ class FP8GlobalState: fp8_parameters: bool = False high_precision_init_val: bool = False is_first_fp8_module: bool = False + pending_backward_quantization_update: bool = False + activation_recompute_enabled: bool = False + activation_recompute_phase: bool = False fp8_graph_capturing: bool = False autocast_depth: int = 0 global_amax_buffer: Dict[str, list] = field(default_factory=dict) @@ -653,7 +656,7 @@ def reduce_tensor_across_group_op_max(tensor: torch.Tensor, group: dist_group_ty ) @classmethod - def reduce_and_update_fp8_tensors( + def reduce_and_update_quantization_state( cls, forward: bool = True, ) -> None: @@ -710,6 +713,29 @@ def reduce_and_update_fp8_tensors( amax_history, scale, get_fp8_max(recipe, forward), recipe ) + # Deprecated alias. + reduce_and_update_fp8_tensors = reduce_and_update_quantization_state + + @classmethod + def schedule_backward_quantization_update(cls) -> None: + """Request a backward quantization state update (e.g. delayed scaling + amax reduction). Called from module backwards; the update itself runs + after the backward pass is complete, at the next top-level autocast + entry. This avoids any assumption about which module's backward + executes last (which does not hold for branched autograd graphs or + activation recompute).""" + cls.quantization_state.pending_backward_quantization_update = True + + @classmethod + def flush_backward_quantization_update(cls) -> None: + """Run the pending backward quantization state update, if any.""" + qstate = cls.quantization_state + if not qstate.pending_backward_quantization_update: + return + qstate.pending_backward_quantization_update = False + with torch.no_grad(): + cls.reduce_and_update_quantization_state(forward=False) + @staticmethod def get_unique_autocast_key( recipe: Optional[Recipe] = None, @@ -740,6 +766,22 @@ def autocast_enter( fp8_recipe = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe autocast_key = cls.get_unique_autocast_key(fp8_recipe, fp8_group) qstate = cls.quantization_state + + # Flush the backward update requested by the previous backward pass. By the + # time a new top-level autocast region is entered, that backward is complete, + # so all backward amaxes are populated. Never flush during activation + # recompute or from inside a running backward pass (an autocast entered by + # an externally driven recompute, e.g. Megatron-Core's). + if ( + qstate.autocast_depth == 0 + and qstate.pending_backward_quantization_update + and not _graph + and not qstate.activation_recompute_phase + and not torch.compiler.is_compiling() + and not cls.fp8_graph_capturing() + and torch._C._current_graph_task_id() == -1 + ): + cls.flush_backward_quantization_update() qstate.autocast_arguments[autocast_key] = ( fp8_recipe, fp8_group, @@ -776,10 +818,18 @@ def autocast_exit(cls, enabled: bool, _graph: bool) -> None: # 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(): + # Skip during activation recompute: the recompute re-enters an autocast, and + # updating here a second time would advance delayed-scaling state twice per step. + if ( + enabled + and qstate.autocast_depth == 0 + and not _graph + and not qstate.activation_recompute_phase + 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) + cls.reduce_and_update_quantization_state(forward=True) @classmethod def copy_forward_fp8_meta_tensors_for_recompute(cls, fp8_meta: Dict[str, Any]) -> None: From eefc7619a8073e4c3b74774bbe95bd2601c84a68 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 17 Aug 2026 12:43:37 +0200 Subject: [PATCH 2/7] Adapt existing tests to deferred backward update timing test_recipe and test_fusible_ops asserted backward scales immediately after backward(); with the deferred update they flush explicitly. Drop raw torch.utils.checkpoint mode from new tests (unsupported with FP8 weight caching independently of this change) and relax exact amax history comparison for nested checkpoints (inner checkpoint re-records amaxes during outer recompute; scales unaffected). Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_fusible_ops.py | 2 ++ .../test_quantization_state_updates.py | 15 ++++++--------- tests/pytorch/test_recipe.py | 3 +++ transformer_engine/pytorch/quantization.py | 19 +++++++++---------- 4 files changed, 20 insertions(+), 19 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 66857d8125..8b0e44b049 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -24,6 +24,7 @@ GRAD_INPUT_BUFFER_KEY, ) from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV +from transformer_engine.pytorch.quantization import FP8GlobalStateManager from transformer_engine.pytorch.ops.fused import ( BackwardActivationBias, @@ -495,6 +496,7 @@ def test_fp8_scale_update( with te.autocast(recipe=recipe): y = model(x) y.backward(dy) + FP8GlobalStateManager.flush_backward_quantization_update() with torch.no_grad(): model.weight.fill_(w_vals[step + 1]) diff --git a/tests/pytorch/test_quantization_state_updates.py b/tests/pytorch/test_quantization_state_updates.py index d088b0b7b7..ad4e4e1ca1 100644 --- a/tests/pytorch/test_quantization_state_updates.py +++ b/tests/pytorch/test_quantization_state_updates.py @@ -108,10 +108,6 @@ def outer(x): return te_checkpoint(outer, x, use_reentrant=True) -def _forward_torch_checkpoint(model, x): - return torch.utils.checkpoint.checkpoint(_run_layers, model, x, use_reentrant=False) - - FORWARD_FNS = { "plain": _forward_plain, "reentrant": _forward_reentrant, @@ -119,7 +115,6 @@ def _forward_torch_checkpoint(model, x): "per_layer_reentrant": _forward_per_layer_reentrant, "per_layer_non_reentrant": _forward_per_layer_non_reentrant, "nested": _forward_nested, - "torch_checkpoint": _forward_torch_checkpoint, } @@ -211,9 +206,7 @@ def test_checkpointed_branch_without_backward(): @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.parametrize( - "mode", ["reentrant", "non_reentrant", "per_layer_reentrant", "nested", "torch_checkpoint"] -) +@pytest.mark.parametrize("mode", ["reentrant", "non_reentrant", "per_layer_reentrant", "nested"]) def test_checkpointing_matches_plain_numerics(mode): """Delayed-scaling state must evolve identically with and without activation checkpointing (duplicate updates would advance it faster).""" @@ -242,7 +235,11 @@ def train(forward): torch.testing.assert_close(l_ckpt, l_ref, rtol=0, atol=0, msg=f"loss diverged @ {step}") for (scale_ref, hist_ref), (scale_ckpt, hist_ckpt) in zip(state_ref, state_ckpt): assert torch.equal(scale_ckpt, scale_ref), "scale diverged under checkpointing" - assert torch.equal(hist_ckpt, hist_ref), "amax history diverged under checkpointing" + # With nested checkpoints the inner checkpoint re-runs its forward (as a + # regular forward frame) during the outer recompute and re-records the + # same amaxes into the history, so exact history equality does not hold. + if mode != "nested": + assert torch.equal(hist_ckpt, hist_ref), "amax history diverged under checkpointing" @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index ccef104a33..1a9a188715 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -94,6 +94,7 @@ def test_fp8_scale_update_with_linear_module( is_first_microbatch=True, ) y.backward(torch.zeros_like(y)) + FP8GlobalStateManager.flush_backward_quantization_update() # Get amax history and scaling factors fp8_meta = module.fp8_meta @@ -147,6 +148,7 @@ def test_fp8_scale_update_with_linear_module( x = torch.randn([16, 16], device="cuda") y = module(x, is_first_microbatch=is_first_microbatch) y.backward(torch.randn_like(y)) + FP8GlobalStateManager.flush_backward_quantization_update() # Check that amax history matches expected values torch.testing.assert_close( @@ -245,6 +247,7 @@ def test_fp8_scale_update_with_linear_fuser_op( with te.autocast(recipe=recipe): y = op(x) y.backward(dy) + FP8GlobalStateManager.flush_backward_quantization_update() def check_metas( test_scale: float, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index be61364af8..107da24665 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -772,16 +772,15 @@ def autocast_enter( # so all backward amaxes are populated. Never flush during activation # recompute or from inside a running backward pass (an autocast entered by # an externally driven recompute, e.g. Megatron-Core's). - if ( - qstate.autocast_depth == 0 - and qstate.pending_backward_quantization_update - and not _graph - and not qstate.activation_recompute_phase - and not torch.compiler.is_compiling() - and not cls.fp8_graph_capturing() - and torch._C._current_graph_task_id() == -1 - ): - cls.flush_backward_quantization_update() + if qstate.autocast_depth == 0 and qstate.pending_backward_quantization_update: + if ( + not _graph + and not qstate.activation_recompute_phase + and not torch.compiler.is_compiling() + and not cls.fp8_graph_capturing() + and torch._C._current_graph_task_id() == -1 + ): + cls.flush_backward_quantization_update() qstate.autocast_arguments[autocast_key] = ( fp8_recipe, fp8_group, From 624cdc610989eaec8bfa559b571c21fa9ea207ae Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 17 Aug 2026 13:00:46 +0200 Subject: [PATCH 3/7] Rename schedule_backward_quantization_update to request_backward_quantization_update The method only idempotently marks the update as pending, so 'request' is more precise. Context booleans renamed to should_request_backward_quantization_update accordingly. Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_backward_override.py | 8 ++++---- .../pytorch/module/grouped_linear.py | 14 +++++++------- .../pytorch/module/layernorm_linear.py | 8 ++++---- transformer_engine/pytorch/module/layernorm_mlp.py | 6 +++--- transformer_engine/pytorch/module/linear.py | 14 ++++++++------ transformer_engine/pytorch/ops/fuser.py | 10 ++++++---- transformer_engine/pytorch/quantization.py | 4 ++-- 7 files changed, 34 insertions(+), 30 deletions(-) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index feff7bccdb..44d88bc9f7 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -419,7 +419,7 @@ def _snapshot_backward_ctx_state( "backward_override", "fp8", "grad_output_quantizer", - "schedule_backward_quantization_update", + "should_request_backward_quantization_update", ) missing_attrs = [attr for attr in required_attrs if not hasattr(state_holder, attr)] if missing_attrs: @@ -430,7 +430,7 @@ def _snapshot_backward_ctx_state( getattr(state_holder, "backward_override"), bool(getattr(state_holder, "fp8")), getattr(state_holder, "grad_output_quantizer"), - bool(getattr(state_holder, "schedule_backward_quantization_update")), + bool(getattr(state_holder, "should_request_backward_quantization_update")), ) @@ -816,7 +816,7 @@ def _run_grouped_linear_single_step_with_ctx_state( required_attrs = ( "backward_override", "fp8", - "schedule_backward_quantization_update", + "should_request_backward_quantization_update", ) missing_attrs = [attr for attr in required_attrs if not hasattr(y.grad_fn, attr)] if missing_attrs: @@ -827,7 +827,7 @@ def _run_grouped_linear_single_step_with_ctx_state( ctx_state = ( getattr(y.grad_fn, "backward_override"), bool(getattr(y.grad_fn, "fp8")), - bool(getattr(y.grad_fn, "schedule_backward_quantization_update")), + bool(getattr(y.grad_fn, "should_request_backward_quantization_update")), ) y.backward(dy) assert x_run.grad is not None diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f042125ed3..a9599f039c 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -963,7 +963,7 @@ def _forward_grouped_tensor( ctx.use_bias = use_bias ctx.inp_shape = inp.shape ctx.requires_dgrad = inp.requires_grad - ctx.schedule_backward_quantization_update = ctx.fp8 and requires_grad( + ctx.should_request_backward_quantization_update = ctx.fp8 and requires_grad( inp, weights[0], biases[0] ) ctx.wgrad_store = wgrad_store @@ -1339,7 +1339,7 @@ def forward( ctx.sequence_parallel = sequence_parallel ctx.inp_shape = inp.shape ctx.requires_dgrad = inp.requires_grad - ctx.schedule_backward_quantization_update = ctx.fp8 and requires_grad( + ctx.should_request_backward_quantization_update = ctx.fp8 and requires_grad( inp, weights[0], biases[0] ) ctx.wgrad_store = wgrad_store @@ -1359,7 +1359,7 @@ def forward( ctx.grad_input_quantizers = [None] * num_gemms ctx.grad_weight_quantizers = [None] * num_gemms ctx.grad_output_quantizers = [None] * num_gemms - ctx.schedule_backward_quantization_update = False + ctx.should_request_backward_quantization_update = False # [*, in_features] -> [*, out_features] except first dimension changes for SP return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @@ -1611,8 +1611,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): else: wgrad_list = [None] * num_weight_args - if ctx.schedule_backward_quantization_update and not is_graph_capturing(): - FP8GlobalStateManager.schedule_backward_quantization_update() + if ctx.should_request_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.request_backward_quantization_update() return ( dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits @@ -1875,8 +1875,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): ): grad_biases = [None] * ctx.num_gemms - if ctx.schedule_backward_quantization_update and not is_graph_capturing(): - FP8GlobalStateManager.schedule_backward_quantization_update() + if ctx.should_request_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.request_backward_quantization_update() return ( dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, None, # m_splits diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 5229464121..f327fd7a6b 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -587,7 +587,7 @@ def forward( ctx.ub_name = ub_name ctx.requires_dgrad = inp_requires_grad ctx.normalization = normalization - ctx.schedule_backward_quantization_update = ctx.fp8 and requires_grad( + ctx.should_request_backward_quantization_update = ctx.fp8 and requires_grad( inp, ln_weight, ln_bias, weight, bias ) ctx.wgrad_store = wgrad_store @@ -604,7 +604,7 @@ def forward( ctx.grad_input_quantizer = None ctx.grad_weight_quantizer = None ctx.grad_output_quantizer = None - ctx.schedule_backward_quantization_update = False + ctx.should_request_backward_quantization_update = False # ------------------------------------------------------ # Cached state for backward pass is ready... @@ -1179,8 +1179,8 @@ def wgrad_gemm( else: wgrad = None - if ctx.schedule_backward_quantization_update and not is_graph_capturing(): - FP8GlobalStateManager.schedule_backward_quantization_update() + if ctx.should_request_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.request_backward_quantization_update() # Scatter fp8 weight buffers # if ctx.fp8 and not isinstance(weight, QuantizedTensorStorage): diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 478df76e13..c5ffbe3727 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -907,7 +907,7 @@ def _forward( inp.requires_grad or ln_weight.requires_grad or ln_bias.requires_grad ) ctx.normalization = normalization - ctx.schedule_backward_quantization_update = ctx.fp8 and requires_grad( + ctx.should_request_backward_quantization_update = ctx.fp8 and requires_grad( inp, ln_weight, ln_bias, fc1_weight, fc2_weight, fc1_bias, fc2_bias ) @@ -1799,8 +1799,8 @@ def fc1_wgrad_gemm( else: fc2_wgrad = None - if ctx.schedule_backward_quantization_update and not is_graph_capturing(): - FP8GlobalStateManager.schedule_backward_quantization_update() + if ctx.should_request_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.request_backward_quantization_update() # FIX THIS # Scatter Fp8 tranposed-weight buffers diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index daad50ffbe..6e21c219ce 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -234,7 +234,7 @@ class LinearBwdArgs: main_grad_func: Optional[Callable[[], torch.Tensor]] = None # --- Quantization state update bookkeeping --- - schedule_backward_quantization_update: bool = False + should_request_backward_quantization_update: bool = False # --- Misc --- cpu_offloading: bool = False @@ -1429,9 +1429,9 @@ def forward( or fwd_args.weight_requires_grad or fwd_args.bias_requires_grad ): - bwd_args.schedule_backward_quantization_update = True + bwd_args.should_request_backward_quantization_update = True if fwd_args.backward_override is not None: - bwd_args.schedule_backward_quantization_update = False + bwd_args.should_request_backward_quantization_update = False return out, new_weight_workspace @@ -1449,13 +1449,15 @@ def backward( if bwd_args.ub_name is not None: nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" result = _linear_backward(bwd_args) + (None,) # fwd_args grad slot - schedule_backward_quantization_update = bwd_args.schedule_backward_quantization_update + should_request_backward_quantization_update = ( + bwd_args.should_request_backward_quantization_update + ) # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, # main_grad closure) so they don't outlive backward via ctx under retain_graph. ctx.backward_objects = None del bwd_args - if schedule_backward_quantization_update and not is_graph_capturing(): - FP8GlobalStateManager.schedule_backward_quantization_update() + if should_request_backward_quantization_update and not is_graph_capturing(): + FP8GlobalStateManager.request_backward_quantization_update() return result diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 99523df17a..292c6f0d9f 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -228,7 +228,7 @@ def forward( func_ctx.tensor_objects = tensor_objects # Whether to request a recipe update in backward pass - schedule_backward_quantization_update = ( + should_request_backward_quantization_update = ( fuser.first_op_requiring_backward < fuser._num_basic_ops and FP8GlobalStateManager.is_fp8_enabled() ) @@ -244,7 +244,9 @@ def forward( func_ctx.basic_op_extra_output_channels = fuser._basic_op_extra_output_channels func_ctx.basic_op_extra_output_consumers = fuser._basic_op_extra_output_consumers func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources - func_ctx.schedule_backward_quantization_update = schedule_backward_quantization_update + func_ctx.should_request_backward_quantization_update = ( + should_request_backward_quantization_update + ) # Mark output tensors as not deletable in backward for tensor in itertools.chain( @@ -385,8 +387,8 @@ def backward( ] # Update FP8 scaling factors - if func_ctx.schedule_backward_quantization_update and not _is_graph_capturing(): - FP8GlobalStateManager.schedule_backward_quantization_update() + if func_ctx.should_request_backward_quantization_update and not _is_graph_capturing(): + FP8GlobalStateManager.request_backward_quantization_update() return ( dx, # input_ diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 2deed671e8..59309b32b1 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -713,11 +713,11 @@ def reduce_and_update_quantization_state( amax_history, scale, get_fp8_max(recipe, forward), recipe ) - # Deprecated alias. + # Compatibility alias (used by Megatron-Core). reduce_and_update_fp8_tensors = reduce_and_update_quantization_state @classmethod - def schedule_backward_quantization_update(cls) -> None: + def request_backward_quantization_update(cls) -> None: """Request a backward quantization state update (e.g. delayed scaling amax reduction). Called from module backwards; the update itself runs after the backward pass is complete, at the next top-level autocast From cfbf8c07d416b98fee78759591b7c6347f94c060 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 17 Aug 2026 13:09:02 +0200 Subject: [PATCH 4/7] Move update-scheduling tests into test_recipe.py Keeps them in a file already enumerated by the L0 QA suite instead of adding a new one. Signed-off-by: Pawel Gadzinski --- .../test_quantization_state_updates.py | 267 ------------------ tests/pytorch/test_recipe.py | 254 +++++++++++++++++ 2 files changed, 254 insertions(+), 267 deletions(-) delete mode 100644 tests/pytorch/test_quantization_state_updates.py diff --git a/tests/pytorch/test_quantization_state_updates.py b/tests/pytorch/test_quantization_state_updates.py deleted file mode 100644 index ad4e4e1ca1..0000000000 --- a/tests/pytorch/test_quantization_state_updates.py +++ /dev/null @@ -1,267 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Tests for delayed-scaling quantization state update scheduling. - -The backward amax reduction must run exactly once per training step, -regardless of activation checkpointing (which re-runs forwards) or the -topology of the autograd graph (branches, unused outputs). -""" - -import pytest -import torch - -import transformer_engine.pytorch as te -from transformer_engine.common.recipe import DelayedScaling -from transformer_engine.pytorch.distributed import checkpoint as te_checkpoint -from transformer_engine.pytorch.quantization import FP8GlobalStateManager - -fp8_available, reason_for_no_fp8 = FP8GlobalStateManager.is_fp8_available() - -HIDDEN = 128 -BATCH = 32 -STEPS = 3 - - -class _UpdateCounter: - """Counts reduce_and_update_quantization_state calls by direction.""" - - def __init__(self): - self.forward = 0 - self.backward = 0 - self._original = None - - def __enter__(self): - self._original = FP8GlobalStateManager.reduce_and_update_quantization_state.__func__ - original = self._original - counter = self - - def counted(cls, forward=True): - if forward: - counter.forward += 1 - else: - counter.backward += 1 - return original(cls, forward=forward) - - FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(counted) - return self - - def __exit__(self, *exc): - FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(self._original) - - -def _make_model(num_layers=3, seed=1234): - torch.manual_seed(seed) - return torch.nn.ModuleList( - [te.Linear(HIDDEN, HIDDEN, bias=True).cuda() for _ in range(num_layers)] - ) - - -def _run_layers(layers, x): - for layer in layers: - x = layer(x) - return x - - -def _train_step(model, x, forward_fn, recipe): - with te.autocast(enabled=True, recipe=recipe): - out = forward_fn(model, x) - loss = out.float().sum() - loss.backward() - return loss - - -def _forward_plain(model, x): - return _run_layers(model, x) - - -def _forward_reentrant(model, x): - return te_checkpoint(_run_layers, model, x, use_reentrant=True) - - -def _forward_non_reentrant(model, x): - return te_checkpoint(_run_layers, model, x, use_reentrant=False) - - -def _forward_per_layer_reentrant(model, x): - for layer in model: - x = te_checkpoint(layer, x, use_reentrant=True) - return x - - -def _forward_per_layer_non_reentrant(model, x): - for layer in model: - x = te_checkpoint(layer, x, use_reentrant=False) - return x - - -def _forward_nested(model, x): - def inner(x): - return te_checkpoint(model[1], x, use_reentrant=True) - - def outer(x): - x = model[0](x) - x = inner(x) - return model[2](x) - - return te_checkpoint(outer, x, use_reentrant=True) - - -FORWARD_FNS = { - "plain": _forward_plain, - "reentrant": _forward_reentrant, - "non_reentrant": _forward_non_reentrant, - "per_layer_reentrant": _forward_per_layer_reentrant, - "per_layer_non_reentrant": _forward_per_layer_non_reentrant, - "nested": _forward_nested, -} - - -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.parametrize("mode", FORWARD_FNS.keys()) -def test_single_update_per_step(mode): - """Exactly one forward and one backward state update per training step.""" - forward_fn = FORWARD_FNS[mode] - model = _make_model() - recipe = DelayedScaling() - - with _UpdateCounter() as counter: - for step in range(STEPS): - x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) - _train_step(model, x, forward_fn, recipe) - # The backward update of this step is flushed at the next - # top-level autocast entry; flush explicitly to count it here. - FP8GlobalStateManager.flush_backward_quantization_update() - assert counter.forward == step + 1, f"{mode}: duplicate/missing forward update" - assert counter.backward == step + 1, f"{mode}: duplicate/missing backward update" - - -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -def test_flush_at_next_autocast_entry(): - """Without an explicit flush, the update runs when the next autocast begins.""" - model = _make_model() - recipe = DelayedScaling() - - with _UpdateCounter() as counter: - x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) - _train_step(model, x, _forward_plain, recipe) - assert counter.backward == 0 - x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) - _train_step(model, x, _forward_plain, recipe) - assert counter.backward == 1 - FP8GlobalStateManager.flush_backward_quantization_update() - - -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.parametrize("checkpoint_first_branch", [True, False]) -def test_branched_graph(checkpoint_first_branch): - """Two checkpointed sibling branches merging into one loss.""" - branch_a = _make_model(num_layers=2, seed=1) - branch_b = _make_model(num_layers=2, seed=2) - recipe = DelayedScaling() - - with _UpdateCounter() as counter: - for step in range(STEPS): - x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) - with te.autocast(enabled=True, recipe=recipe): - if checkpoint_first_branch: - ya = te_checkpoint(_run_layers, branch_a, x, use_reentrant=True) - else: - ya = _run_layers(branch_a, x) - yb = te_checkpoint(_run_layers, branch_b, x, use_reentrant=True) - out = ya + yb - loss = out.float().sum() - loss.backward() - FP8GlobalStateManager.flush_backward_quantization_update() - assert counter.forward == step + 1 - assert counter.backward == step + 1 - assert x.grad is not None and torch.isfinite(x.grad).all() - x.grad = None - - -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -def test_checkpointed_branch_without_backward(): - """A checkpointed branch whose output never receives a gradient must not - strand the backward update (regression test for first-module ownership).""" - used = _make_model(num_layers=2, seed=1) - unused = _make_model(num_layers=2, seed=2) - recipe = DelayedScaling() - - with _UpdateCounter() as counter: - for step in range(STEPS): - x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) - with te.autocast(enabled=True, recipe=recipe): - # The unused branch runs first, so any "first module owns the - # backward update" scheme would assign ownership to a frame - # whose backward never executes. - y_unused = te_checkpoint(_run_layers, unused, x, use_reentrant=True) - y = _run_layers(used, x) - loss = y.float().sum() - loss.backward() - del y_unused - FP8GlobalStateManager.flush_backward_quantization_update() - assert counter.backward == step + 1, "backward update was lost" - x.grad = None - - -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.parametrize("mode", ["reentrant", "non_reentrant", "per_layer_reentrant", "nested"]) -def test_checkpointing_matches_plain_numerics(mode): - """Delayed-scaling state must evolve identically with and without - activation checkpointing (duplicate updates would advance it faster).""" - forward_fn = FORWARD_FNS[mode] - recipe = DelayedScaling() - - def train(forward): - model = _make_model(seed=99) - losses = [] - torch.manual_seed(777) - for _ in range(STEPS + 1): - x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) - losses.append(_train_step(model, x, forward, recipe).detach()) - FP8GlobalStateManager.flush_backward_quantization_update() - state = [] - for layer in model: - for key in ("scaling_fwd", "scaling_bwd"): - meta = layer.fp8_meta[key] - state.append((meta.scale.clone(), meta.amax_history.clone())) - return losses, state - - losses_ref, state_ref = train(_forward_plain) - losses_ckpt, state_ckpt = train(forward_fn) - - for step, (l_ref, l_ckpt) in enumerate(zip(losses_ref, losses_ckpt)): - torch.testing.assert_close(l_ckpt, l_ref, rtol=0, atol=0, msg=f"loss diverged @ {step}") - for (scale_ref, hist_ref), (scale_ckpt, hist_ckpt) in zip(state_ref, state_ckpt): - assert torch.equal(scale_ckpt, scale_ref), "scale diverged under checkpointing" - # With nested checkpoints the inner checkpoint re-runs its forward (as a - # regular forward frame) during the outer recompute and re-records the - # same amaxes into the history, so exact history equality does not hold. - if mode != "nested": - assert torch.equal(hist_ckpt, hist_ref), "amax history diverged under checkpointing" - - -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -def test_exception_in_checkpointed_forward(): - """A failing checkpointed forward must not corrupt update scheduling.""" - model = _make_model() - recipe = DelayedScaling() - - def failing(model, x): - def fn(x): - model[0](x) - raise RuntimeError("boom") - - return te_checkpoint(fn, x, use_reentrant=True) - - with _UpdateCounter() as counter: - x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) - with pytest.raises(RuntimeError, match="boom"): - _train_step(model, x, failing, recipe) - - x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) - _train_step(model, x, _forward_reentrant, recipe) - FP8GlobalStateManager.flush_backward_quantization_update() - assert counter.backward == 1 - assert torch.isfinite(x.grad).all() diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 1a9a188715..56724ed12b 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -32,6 +32,7 @@ _amax_and_scale_update, ) import transformer_engine.pytorch.ops as te_ops +from transformer_engine.pytorch.distributed import checkpoint as te_checkpoint from transformer_engine.common.recipe import ( CustomRecipe, DelayedScaling, @@ -783,3 +784,256 @@ def test_stateful_unknown_or_malformed_pickled_extra_state_requires_opt_in(paylo monkeypatch.setenv(UNSAFE_PICKLE_EXTRA_STATE_ENV, "1") assert should_load_extra_state_pickle(payload, "test") + + +############################################################################### +# Backward quantization state update scheduling +############################################################################### + + +HIDDEN = 128 +BATCH = 32 +STEPS = 3 + + +class _UpdateCounter: + """Counts reduce_and_update_quantization_state calls by direction.""" + + def __init__(self): + self.forward = 0 + self.backward = 0 + self._original = None + + def __enter__(self): + self._original = FP8GlobalStateManager.reduce_and_update_quantization_state.__func__ + original = self._original + counter = self + + def counted(cls, forward=True): + if forward: + counter.forward += 1 + else: + counter.backward += 1 + return original(cls, forward=forward) + + FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(counted) + return self + + def __exit__(self, *exc): + FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(self._original) + + +def _make_model(num_layers=3, seed=1234): + torch.manual_seed(seed) + return torch.nn.ModuleList( + [te.Linear(HIDDEN, HIDDEN, bias=True).cuda() for _ in range(num_layers)] + ) + + +def _run_layers(layers, x): + for layer in layers: + x = layer(x) + return x + + +def _train_step(model, x, forward_fn, recipe): + with te.autocast(enabled=True, recipe=recipe): + out = forward_fn(model, x) + loss = out.float().sum() + loss.backward() + return loss + + +def _forward_plain(model, x): + return _run_layers(model, x) + + +def _forward_reentrant(model, x): + return te_checkpoint(_run_layers, model, x, use_reentrant=True) + + +def _forward_non_reentrant(model, x): + return te_checkpoint(_run_layers, model, x, use_reentrant=False) + + +def _forward_per_layer_reentrant(model, x): + for layer in model: + x = te_checkpoint(layer, x, use_reentrant=True) + return x + + +def _forward_per_layer_non_reentrant(model, x): + for layer in model: + x = te_checkpoint(layer, x, use_reentrant=False) + return x + + +def _forward_nested(model, x): + def inner(x): + return te_checkpoint(model[1], x, use_reentrant=True) + + def outer(x): + x = model[0](x) + x = inner(x) + return model[2](x) + + return te_checkpoint(outer, x, use_reentrant=True) + + +FORWARD_FNS = { + "plain": _forward_plain, + "reentrant": _forward_reentrant, + "non_reentrant": _forward_non_reentrant, + "per_layer_reentrant": _forward_per_layer_reentrant, + "per_layer_non_reentrant": _forward_per_layer_non_reentrant, + "nested": _forward_nested, +} + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("mode", FORWARD_FNS.keys()) +def test_single_update_per_step(mode): + """Exactly one forward and one backward state update per training step.""" + forward_fn = FORWARD_FNS[mode] + model = _make_model() + recipe = DelayedScaling() + + with _UpdateCounter() as counter: + for step in range(STEPS): + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + _train_step(model, x, forward_fn, recipe) + # The backward update of this step is flushed at the next + # top-level autocast entry; flush explicitly to count it here. + FP8GlobalStateManager.flush_backward_quantization_update() + assert counter.forward == step + 1, f"{mode}: duplicate/missing forward update" + assert counter.backward == step + 1, f"{mode}: duplicate/missing backward update" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_flush_at_next_autocast_entry(): + """Without an explicit flush, the update runs when the next autocast begins.""" + model = _make_model() + recipe = DelayedScaling() + + with _UpdateCounter() as counter: + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + _train_step(model, x, _forward_plain, recipe) + assert counter.backward == 0 + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + _train_step(model, x, _forward_plain, recipe) + assert counter.backward == 1 + FP8GlobalStateManager.flush_backward_quantization_update() + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("checkpoint_first_branch", [True, False]) +def test_branched_graph(checkpoint_first_branch): + """Two checkpointed sibling branches merging into one loss.""" + branch_a = _make_model(num_layers=2, seed=1) + branch_b = _make_model(num_layers=2, seed=2) + recipe = DelayedScaling() + + with _UpdateCounter() as counter: + for step in range(STEPS): + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + with te.autocast(enabled=True, recipe=recipe): + if checkpoint_first_branch: + ya = te_checkpoint(_run_layers, branch_a, x, use_reentrant=True) + else: + ya = _run_layers(branch_a, x) + yb = te_checkpoint(_run_layers, branch_b, x, use_reentrant=True) + out = ya + yb + loss = out.float().sum() + loss.backward() + FP8GlobalStateManager.flush_backward_quantization_update() + assert counter.forward == step + 1 + assert counter.backward == step + 1 + assert x.grad is not None and torch.isfinite(x.grad).all() + x.grad = None + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_checkpointed_branch_without_backward(): + """A checkpointed branch whose output never receives a gradient must not + strand the backward update (regression test for first-module ownership).""" + used = _make_model(num_layers=2, seed=1) + unused = _make_model(num_layers=2, seed=2) + recipe = DelayedScaling() + + with _UpdateCounter() as counter: + for step in range(STEPS): + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + with te.autocast(enabled=True, recipe=recipe): + # The unused branch runs first, so any "first module owns the + # backward update" scheme would assign ownership to a frame + # whose backward never executes. + y_unused = te_checkpoint(_run_layers, unused, x, use_reentrant=True) + y = _run_layers(used, x) + loss = y.float().sum() + loss.backward() + del y_unused + FP8GlobalStateManager.flush_backward_quantization_update() + assert counter.backward == step + 1, "backward update was lost" + x.grad = None + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("mode", ["reentrant", "non_reentrant", "per_layer_reentrant", "nested"]) +def test_checkpointing_matches_plain_numerics(mode): + """Delayed-scaling state must evolve identically with and without + activation checkpointing (duplicate updates would advance it faster).""" + forward_fn = FORWARD_FNS[mode] + recipe = DelayedScaling() + + def train(forward): + model = _make_model(seed=99) + losses = [] + torch.manual_seed(777) + for _ in range(STEPS + 1): + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + losses.append(_train_step(model, x, forward, recipe).detach()) + FP8GlobalStateManager.flush_backward_quantization_update() + state = [] + for layer in model: + for key in ("scaling_fwd", "scaling_bwd"): + meta = layer.fp8_meta[key] + state.append((meta.scale.clone(), meta.amax_history.clone())) + return losses, state + + losses_ref, state_ref = train(_forward_plain) + losses_ckpt, state_ckpt = train(forward_fn) + + for step, (l_ref, l_ckpt) in enumerate(zip(losses_ref, losses_ckpt)): + torch.testing.assert_close(l_ckpt, l_ref, rtol=0, atol=0, msg=f"loss diverged @ {step}") + for (scale_ref, hist_ref), (scale_ckpt, hist_ckpt) in zip(state_ref, state_ckpt): + assert torch.equal(scale_ckpt, scale_ref), "scale diverged under checkpointing" + # With nested checkpoints the inner checkpoint re-runs its forward (as a + # regular forward frame) during the outer recompute and re-records the + # same amaxes into the history, so exact history equality does not hold. + if mode != "nested": + assert torch.equal(hist_ckpt, hist_ref), "amax history diverged under checkpointing" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_exception_in_checkpointed_forward(): + """A failing checkpointed forward must not corrupt update scheduling.""" + model = _make_model() + recipe = DelayedScaling() + + def failing(model, x): + def fn(x): + model[0](x) + raise RuntimeError("boom") + + return te_checkpoint(fn, x, use_reentrant=True) + + with _UpdateCounter() as counter: + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + with pytest.raises(RuntimeError, match="boom"): + _train_step(model, x, failing, recipe) + + x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True) + _train_step(model, x, _forward_reentrant, recipe) + FP8GlobalStateManager.flush_backward_quantization_update() + assert counter.backward == 1 + assert torch.isfinite(x.grad).all() From 38a6abf3014d8917572b392cf8f7d7011ea7c3a3 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 17 Aug 2026 13:25:06 +0200 Subject: [PATCH 5/7] Credit original investigation of the duplicate FP8 update bug The activation-checkpoint duplicate-update problem and its regression test scenarios were first diagnosed and addressed in #3213; this PR supersedes that approach by removing first-module ownership entirely. Co-authored-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com> Signed-off-by: Pawel Gadzinski From d64dcadba5d06570f172f5a4f3b4b4d1e9d3b68b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 17 Aug 2026 13:39:04 +0200 Subject: [PATCH 6/7] Make quantization state updates per-recipe via RecipeState FP8GlobalStateManager now only orchestrates when updates run; how a recipe updates its process-global state is a RecipeState classmethod (reduce_and_update_global_state), no-op by default and overridden by DelayedScalingRecipeState with the amax reduction + scale recompute. The recipe-to-state-class mapping is extracted from RecipeState.create into class_for_recipe and reused for dispatch. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/quantization.py | 148 +++++++++++++-------- 1 file changed, 94 insertions(+), 54 deletions(-) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 59309b32b1..eb7d9f9199 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -660,59 +660,36 @@ def reduce_and_update_quantization_state( cls, forward: bool = True, ) -> None: - """Delayed scaling only. Concatenate, reduce, and split amaxes in the global buffer.""" - # global_amax_buffer should only be non-empty for fp8 delayed scaling + """Run the per-recipe quantization state updates over the global buffers. + + The manager only orchestrates *when* updates run (forward at autocast + exit, backward after the backward pass); *how* a recipe updates its + state is defined by `RecipeState.reduce_and_update_global_state` + overrides. Currently only delayed scaling has one, so the buffers + are empty for other recipes. + """ qstate = cls.quantization_state for ( buffer_key, amax_buffer, ) in qstate.global_amax_buffer.items(): - # Check for forward or backward reduction. + # Check for forward or backward update. fwd_update, autocast_key = cls.split_key_in_buffer(buffer_key) if fwd_update != forward: continue if len(amax_buffer) == 0: continue - # Retrieve autocast specific args and concat amaxes. recipe, group = qstate.autocast_arguments[autocast_key] - contiguous_amax = torch.cat(amax_buffer) - - # Reduction. - if ( - recipe.reduce_amax - and torch.distributed.is_initialized() - and torch.distributed.get_world_size(group=group) > 1 - ): - cls.reduce_tensor_across_group_op_max(contiguous_amax, group) - - # Amax and scale update. - unfused_update = ( - bool(int(os.getenv("NVTE_UNFUSED_FP8_UPDATE", "0"))) - or callable(recipe.amax_compute_algo) - or callable(recipe.scaling_factor_compute_algo) + RecipeState.class_for_recipe(recipe).reduce_and_update_global_state( + recipe, + group, + forward, + amax_buffer, + qstate.global_amax_history_buffer[buffer_key], + qstate.global_scale_buffer[buffer_key], ) - if not unfused_update: - tex.fused_amax_and_scale_update_after_reduction( - contiguous_amax, - qstate.global_amax_history_buffer[buffer_key], - qstate.global_scale_buffer[buffer_key], - recipe.amax_compute_algo, - get_fp8_te_dtype(recipe, forward), - recipe.margin, - ) - else: - split_and_copy(contiguous_amax, amax_buffer, [x.numel() for x in amax_buffer]) - - for amax_history, scale in zip( - qstate.global_amax_history_buffer[buffer_key], - qstate.global_scale_buffer[buffer_key], - ): - _amax_and_scale_update( - amax_history, scale, get_fp8_max(recipe, forward), recipe - ) - # Compatibility alias (used by Megatron-Core). reduce_and_update_fp8_tensors = reduce_and_update_quantization_state @@ -1372,21 +1349,7 @@ def create( """ - cls = None - if recipe.delayed(): - cls = DelayedScalingRecipeState - elif recipe.mxfp8(): - cls = MXFP8BlockScalingRecipeState - elif recipe.float8_current_scaling(): - cls = Float8CurrentScalingRecipeState - elif recipe.float8_block_scaling(): - cls = Float8BlockScalingRecipeState - elif recipe.nvfp4(): - cls = NVFP4BlockScalingRecipeState - elif recipe.custom(): - cls = CustomRecipeState - else: - raise ValueError(f"{recipe.__class__.__name__} is not supported") + cls = RecipeState.class_for_recipe(recipe) return cls( recipe, mode=mode, @@ -1395,6 +1358,42 @@ def create( roles=roles, ) + @staticmethod + def class_for_recipe(recipe: Recipe) -> type: + """Resolve the RecipeState subclass handling a recipe.""" + if recipe.delayed(): + return DelayedScalingRecipeState + if recipe.mxfp8(): + return MXFP8BlockScalingRecipeState + if recipe.float8_current_scaling(): + return Float8CurrentScalingRecipeState + if recipe.float8_block_scaling(): + return Float8BlockScalingRecipeState + if recipe.nvfp4(): + return NVFP4BlockScalingRecipeState + if recipe.custom(): + return CustomRecipeState + raise ValueError(f"{recipe.__class__.__name__} is not supported") + + @classmethod + def reduce_and_update_global_state( + cls, + recipe: Recipe, + group: Optional[dist_group_type], + forward: bool, + amax_buffer: List[torch.Tensor], + amax_history_buffer: List[torch.Tensor], + scale_buffer: List[torch.Tensor], + ) -> None: + """Run this recipe's post-step update of process-global state. + + Called by `FP8GlobalStateManager` once per autocast key, with the + state buffers of all participating modules — after the forward pass + (autocast exit) and after the backward pass. The manager decides + *when* updates run; recipes define *what* they update. The default + is a no-op: most recipes keep no process-global state. + """ + @abc.abstractmethod def make_quantizers(self) -> list: """Convert recipe state to quantizers. @@ -1463,6 +1462,47 @@ class DelayedScalingRecipeState(RecipeState): # See ``RecipeState.inherit_state_from``. _persistent_state_buffers = ("scale", "amax_history") + @classmethod + def reduce_and_update_global_state( + cls, + recipe: DelayedScaling, + group: Optional[dist_group_type], + forward: bool, + amax_buffer: List[torch.Tensor], + amax_history_buffer: List[torch.Tensor], + scale_buffer: List[torch.Tensor], + ) -> None: + """Concatenate, reduce, and split amaxes, then recompute scales.""" + contiguous_amax = torch.cat(amax_buffer) + + if ( + recipe.reduce_amax + and torch.distributed.is_initialized() + and torch.distributed.get_world_size(group=group) > 1 + ): + FP8GlobalStateManager.reduce_tensor_across_group_op_max(contiguous_amax, group) + + unfused_update = ( + bool(int(os.getenv("NVTE_UNFUSED_FP8_UPDATE", "0"))) + or callable(recipe.amax_compute_algo) + or callable(recipe.scaling_factor_compute_algo) + ) + + if not unfused_update: + tex.fused_amax_and_scale_update_after_reduction( + contiguous_amax, + amax_history_buffer, + scale_buffer, + recipe.amax_compute_algo, + get_fp8_te_dtype(recipe, forward), + recipe.margin, + ) + else: + split_and_copy(contiguous_amax, amax_buffer, [x.numel() for x in amax_buffer]) + + for amax_history, scale in zip(amax_history_buffer, scale_buffer): + _amax_and_scale_update(amax_history, scale, get_fp8_max(recipe, forward), recipe) + def __init__( self, recipe: DelayedScaling, From ef588896292f9ce0e992345b912ed803ae264460 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Mon, 17 Aug 2026 13:46:04 +0200 Subject: [PATCH 7/7] Restore NVTX range around the backward quantization state update The reduction no longer runs in module backwards, so the old per-module NVTX markers had nothing left to wrap; mark it at its new single execution site in flush_backward_quantization_update instead. Signed-off-by: Pawel Gadzinski --- transformer_engine/pytorch/quantization.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index eb7d9f9199..b52aa93939 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -28,7 +28,7 @@ ) from .constants import dist_group_type, DType -from .utils import get_device_compute_capability +from .utils import get_device_compute_capability, nvtx_range_push, nvtx_range_pop from .jit import jit_fuser @@ -710,8 +710,10 @@ def flush_backward_quantization_update(cls) -> None: if not qstate.pending_backward_quantization_update: return qstate.pending_backward_quantization_update = False + nvtx_range_push("transformer_engine.reduce_and_update_quantization_state.backward") with torch.no_grad(): cls.reduce_and_update_quantization_state(forward=False) + nvtx_range_pop("transformer_engine.reduce_and_update_quantization_state.backward") @staticmethod def get_unique_autocast_key(