From b065243b3d208e7e5494b85f77b750ac31346a2b Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 20 Aug 2026 17:54:58 +0200
Subject: [PATCH 01/13] [PyTorch] Schedule delayed-scaling updates after
 backward

Queue one quantization state update at the autograd boundary instead of assigning it to the first FP8 module seen in forward. Add an optional logical-backward scope for multi-backward schedules and delayed weight-gradient computation.

Co-authored-by: AlbertYang514 <201034045+AlbertYang514@users.noreply.github.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 docs/api/pytorch.rst                          |   2 +
 tests/pytorch/test_backward_override.py       |  34 +--
 tests/pytorch/test_recipe.py                  | 205 ++++++++++++++++++
 transformer_engine/pytorch/__init__.py        |   1 +
 transformer_engine/pytorch/distributed.py     |  28 ++-
 .../pytorch/module/grouped_linear.py          |  33 ++-
 .../pytorch/module/layernorm_linear.py        |  21 +-
 .../pytorch/module/layernorm_mlp.py           |  21 +-
 transformer_engine/pytorch/module/linear.py   |  31 +--
 transformer_engine/pytorch/ops/fuser.py       |  19 +-
 transformer_engine/pytorch/quantization.py    |  95 +++++++-
 11 files changed, 390 insertions(+), 100 deletions(-)

diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst
index 5fac0a89a6..3115fb027f 100644
--- a/docs/api/pytorch.rst
+++ b/docs/api/pytorch.rst
@@ -40,6 +40,8 @@ PyTorch
 
 .. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None)
 
+.. autoapifunction:: transformer_engine.pytorch.backward_quantization_update_scope
+
 .. autoapifunction:: transformer_engine.pytorch.quantized_model_init
 
 .. autoapifunction:: transformer_engine.pytorch.checkpoint
diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py
index c0acf2e6b3..00bf5a8f2e 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",
+        "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, "reduce_and_update_bwd_fp8_tensors")),
+        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",
-            "reduce_and_update_bwd_fp8_tensors",
+            "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, "reduce_and_update_bwd_fp8_tensors")),
+            bool(getattr(y.grad_fn, "should_request_backward_quantization_update")),
         )
     y.backward(dy)
     assert x_run.grad is not None
@@ -1453,33 +1453,34 @@ def test_linear_like_runtime_backward_override_switch_updates_ctx(
         default_mode,
         default_fp8,
         default_grad_output_quantizer,
-        default_reduce_and_update,
+        default_should_request_update,
     ) = default_ctx
+    expected_request = default_recipe.delayed() or default_recipe.custom()
     assert default_mode is None
     assert default_fp8
     assert default_grad_output_quantizer is not None
-    assert default_reduce_and_update
+    assert default_should_request_update == expected_request
 
     *_, switched_ctx = _run_single_step_with_ctx_state(module, x, dy, mode_recipe)
-    switched_mode, switched_fp8, switched_grad_output_quantizer, switched_reduce_and_update = (
+    switched_mode, switched_fp8, switched_grad_output_quantizer, switched_should_request_update = (
         switched_ctx
     )
     assert switched_mode == backward_override
     assert not switched_fp8
     assert switched_grad_output_quantizer is None
-    assert not switched_reduce_and_update
+    assert not switched_should_request_update
 
     *_, default_ctx_after = _run_single_step_with_ctx_state(module, x, dy, default_recipe)
     (
         default_mode_after,
         default_fp8_after,
         default_grad_output_quantizer_after,
-        default_reduce_and_update_after,
+        default_should_request_update_after,
     ) = default_ctx_after
     assert default_mode_after is None
     assert default_fp8_after
     assert default_grad_output_quantizer_after is not None
-    assert default_reduce_and_update_after
+    assert default_should_request_update_after == expected_request
 
 
 @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
@@ -1526,10 +1527,11 @@ def test_grouped_linear_runtime_backward_override_switch_updates_ctx(
         dy,
         default_recipe,
     )
-    default_mode, default_fp8, default_reduce_and_update = default_ctx
+    default_mode, default_fp8, default_should_request_update = default_ctx
+    expected_request = default_recipe.delayed() or default_recipe.custom()
     assert default_mode is None
     assert default_fp8
-    assert default_reduce_and_update
+    assert default_should_request_update == expected_request
 
     *_, switched_ctx = _run_grouped_linear_single_step_with_ctx_state(
         module,
@@ -1538,10 +1540,10 @@ def test_grouped_linear_runtime_backward_override_switch_updates_ctx(
         dy,
         mode_recipe,
     )
-    switched_mode, switched_fp8, switched_reduce_and_update = switched_ctx
+    switched_mode, switched_fp8, switched_should_request_update = switched_ctx
     assert switched_mode == backward_override
     assert not switched_fp8
-    assert not switched_reduce_and_update
+    assert not switched_should_request_update
 
     *_, default_ctx_after = _run_grouped_linear_single_step_with_ctx_state(
         module,
@@ -1550,10 +1552,10 @@ def test_grouped_linear_runtime_backward_override_switch_updates_ctx(
         dy,
         default_recipe,
     )
-    default_mode_after, default_fp8_after, default_reduce_and_update_after = default_ctx_after
+    default_mode_after, default_fp8_after, default_should_request_update_after = default_ctx_after
     assert default_mode_after is None
     assert default_fp8_after
-    assert default_reduce_and_update_after
+    assert default_should_request_update_after == expected_request
 
 
 @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py
index ccef104a33..b077d70383 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,
@@ -780,3 +781,207 @@ 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")
+
+
+_UPDATE_TEST_HIDDEN = 128
+_UPDATE_TEST_BATCH = 32
+_UPDATE_TEST_STEPS = 3
+
+
+class _UpdateCounter:
+    def __init__(self):
+        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 not forward:
+                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_update_test_model(num_layers=3, seed=1234):
+    torch.manual_seed(seed)
+    return torch.nn.ModuleList(
+        [
+            te.Linear(_UPDATE_TEST_HIDDEN, _UPDATE_TEST_HIDDEN, bias=True).cuda()
+            for _ in range(num_layers)
+        ]
+    )
+
+
+def _run_update_test_layers(layers, x):
+    for layer in layers:
+        x = layer(x)
+    return x
+
+
+def _run_update_test_step(model, x, forward_fn, recipe):
+    with te.autocast(enabled=True, recipe=recipe):
+        out = forward_fn(model, x)
+    loss = out.float().sum()
+    loss.backward()
+
+
+def _update_forward_plain(model, x):
+    return _run_update_test_layers(model, x)
+
+
+def _update_forward_reentrant(model, x):
+    return te_checkpoint(_run_update_test_layers, model, x, use_reentrant=True)
+
+
+def _update_forward_non_reentrant(model, x):
+    return te_checkpoint(_run_update_test_layers, model, x, use_reentrant=False)
+
+
+def _update_forward_per_layer_reentrant(model, x):
+    for layer in model:
+        x = te_checkpoint(layer, x, use_reentrant=True)
+    return x
+
+
+def _update_forward_nested(model, x):
+    def inner(value):
+        return te_checkpoint(model[1], value, use_reentrant=True)
+
+    def outer(value):
+        return model[2](inner(model[0](value)))
+
+    return te_checkpoint(outer, x, use_reentrant=True)
+
+
+_UPDATE_FORWARD_FNS = {
+    "plain": _update_forward_plain,
+    "reentrant": _update_forward_reentrant,
+    "non_reentrant": _update_forward_non_reentrant,
+    "per_layer_reentrant": _update_forward_per_layer_reentrant,
+    "nested": _update_forward_nested,
+}
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("mode", _UPDATE_FORWARD_FNS.keys())
+def test_delayed_scaling_updates_once_per_backward(mode):
+    FP8GlobalStateManager.reset()
+    model = _make_update_test_model()
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter:
+        for step in range(_UPDATE_TEST_STEPS):
+            x = torch.randn(
+                _UPDATE_TEST_BATCH,
+                _UPDATE_TEST_HIDDEN,
+                device="cuda",
+                requires_grad=True,
+            )
+            _run_update_test_step(model, x, _UPDATE_FORWARD_FNS[mode], recipe)
+            assert counter.backward == step + 1
+            qstate = FP8GlobalStateManager.quantization_state
+            assert not qstate.pending_backward_quantization_update
+            assert qstate.backward_quantization_update_callback_task_id is None
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_backward_quantization_update_scope_groups_autograd_calls():
+    FP8GlobalStateManager.reset()
+    model = _make_update_test_model()
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter, te.backward_quantization_update_scope():
+        for _ in range(_UPDATE_TEST_STEPS):
+            x = torch.randn(
+                _UPDATE_TEST_BATCH,
+                _UPDATE_TEST_HIDDEN,
+                device="cuda",
+                requires_grad=True,
+            )
+            _run_update_test_step(model, x, _update_forward_plain, recipe)
+            assert counter.backward == 0
+    assert counter.backward == 1
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_backward_quantization_update_scope_covers_delayed_wgrad():
+    FP8GlobalStateManager.reset()
+    model = te.Linear(
+        _UPDATE_TEST_HIDDEN,
+        _UPDATE_TEST_HIDDEN,
+        bias=True,
+        delay_wgrad_compute=True,
+    ).cuda()
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter, te.backward_quantization_update_scope():
+        x = torch.randn(
+            _UPDATE_TEST_BATCH,
+            _UPDATE_TEST_HIDDEN,
+            device="cuda",
+            requires_grad=True,
+        )
+        with te.autocast(enabled=True, recipe=recipe):
+            out = model(x)
+        out.float().sum().backward()
+        assert counter.backward == 0
+        model.backward_dw()
+        assert model.weight.grad is not None
+        assert counter.backward == 0
+    assert counter.backward == 1
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("checkpoint_first_branch", [True, False])
+def test_delayed_scaling_update_on_branched_graph(checkpoint_first_branch):
+    FP8GlobalStateManager.reset()
+    branch_a = _make_update_test_model(num_layers=2, seed=1)
+    branch_b = _make_update_test_model(num_layers=2, seed=2)
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter:
+        x = torch.randn(
+            _UPDATE_TEST_BATCH,
+            _UPDATE_TEST_HIDDEN,
+            device="cuda",
+            requires_grad=True,
+        )
+        with te.autocast(enabled=True, recipe=recipe):
+            if checkpoint_first_branch:
+                out_a = te_checkpoint(_run_update_test_layers, branch_a, x, use_reentrant=True)
+            else:
+                out_a = _run_update_test_layers(branch_a, x)
+            out_b = te_checkpoint(_run_update_test_layers, branch_b, x, use_reentrant=True)
+        (out_a + out_b).float().sum().backward()
+        assert counter.backward == 1
+        assert x.grad is not None and torch.isfinite(x.grad).all()
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_unused_checkpoint_branch_does_not_own_backward_update():
+    FP8GlobalStateManager.reset()
+    used = _make_update_test_model(num_layers=2, seed=1)
+    unused = _make_update_test_model(num_layers=2, seed=2)
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter:
+        x = torch.randn(
+            _UPDATE_TEST_BATCH,
+            _UPDATE_TEST_HIDDEN,
+            device="cuda",
+            requires_grad=True,
+        )
+        with te.autocast(enabled=True, recipe=recipe):
+            unused_out = te_checkpoint(_run_update_test_layers, unused, x, use_reentrant=True)
+            out = _run_update_test_layers(used, x)
+        out.float().sum().backward()
+        del unused_out
+        assert counter.backward == 1
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index 2b1803bfb2..26f10bef9f 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/transformer_engine/pytorch/__init__.py
@@ -45,6 +45,7 @@
 from transformer_engine.pytorch.quantization import fp8_autocast
 from transformer_engine.pytorch.quantization import fp8_model_init
 from transformer_engine.pytorch.quantization import autocast
+from transformer_engine.pytorch.quantization import backward_quantization_update_scope
 from transformer_engine.pytorch.quantization import quantized_model_init
 from transformer_engine.pytorch.quantization import is_fp8_available
 from transformer_engine.pytorch.quantization import is_mxfp8_available
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 8605a4746b..2fc6b85340 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -39,7 +39,7 @@
 )
 
 from .constants import dist_group_type
-from .quantization import FP8GlobalStateManager, autocast
+from .quantization import FP8GlobalStateManager, autocast, backward_quantization_update_scope
 from .tensor.float8_tensor import Float8Quantizer, Float8Tensor, Float8CurrentScalingQuantizer
 from .tensor.mxfp8_tensor import MXFP8Quantizer
 from .tensor.nvfp4_tensor import NVFP4Quantizer
@@ -247,8 +247,6 @@ 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
@@ -264,12 +262,6 @@ def __enter__(self):
         _IN_ACTIVATION_RECOMPUTE_REGION = self.activation_recompute
         _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)
-
     def __exit__(self, *exc_details):
         global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE
         _IN_ACTIVATION_RECOMPUTE_REGION = False
@@ -407,6 +399,20 @@ def backward(
         ctx, *args: Tuple[Union[torch.Tensor, None], ...]
     ) -> Tuple[Union[torch.Tensor, None], ...]:
         """Call backward function with activation recomputation."""
+        recipe = ctx.fp8_recipe
+        update_scope = (
+            backward_quantization_update_scope()
+            if ctx.fp8 and (recipe.delayed() or recipe.custom())
+            else nullcontext()
+        )
+        with update_scope:
+            return _CheckpointFunction._backward(ctx, *args)
+
+    @staticmethod
+    def _backward(
+        ctx, *args: Tuple[Union[torch.Tensor, None], ...]
+    ) -> Tuple[Union[torch.Tensor, None], ...]:
+        """Recompute the forward and run its nested backward."""
         if not torch.autograd._is_checkpoint_valid():
             raise RuntimeError(
                 "Checkpointing is not compatible with .grad(), please use .backward() if possible"
@@ -440,9 +446,7 @@ def backward(
         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
-        ), autocast(
-            enabled=ctx.fp8, recipe=ctx.fp8_recipe
-        ):
+        ), autocast(enabled=ctx.fp8, recipe=ctx.fp8_recipe):
             outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
 
         # Set the states back to what it was at the start of this function.
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 612a430966..5cae0bb265 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -50,6 +50,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,
@@ -599,12 +600,11 @@ 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.should_request_backward_quantization_update = (
+                ctx.fp8
+                and (ctx.fp8_recipe.delayed() or ctx.fp8_recipe.custom())
+                and requires_grad(inp, weights[0], biases[0])
+            )
             ctx.wgrad_store = wgrad_store
             ctx.debug = False
             ctx.save_original_input = save_original_input
@@ -975,12 +975,11 @@ 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.should_request_backward_quantization_update = (
+                ctx.fp8
+                and (ctx.fp8_recipe.delayed() or ctx.fp8_recipe.custom())
+                and requires_grad(inp, weights[0], biases[0])
+            )
             ctx.wgrad_store = wgrad_store
             ctx.debug = debug
             ctx.save_original_input = save_original_input
@@ -998,7 +997,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.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
@@ -1250,8 +1249,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
         else:
             wgrad_list = [None] * num_weight_args
 
-        if ctx.reduce_and_update_bwd_fp8_tensors:
-            FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
+        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
@@ -1517,8 +1516,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.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 561e813348..1ed8326572 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,11 @@ 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.should_request_backward_quantization_update = (
+                ctx.fp8
+                and (ctx.fp8_recipe.delayed() or ctx.fp8_recipe.custom())
+                and requires_grad(inp, ln_weight, ln_bias, weight, bias)
+            )
             ctx.wgrad_store = wgrad_store
             ctx.debug = debug
 
@@ -609,7 +606,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.should_request_backward_quantization_update = False
 
         # ------------------------------------------------------
         # Cached state for backward pass is ready...
@@ -1184,10 +1181,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.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 3ee0cda50c..3939f7e21d 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,13 @@ 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(
-                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.should_request_backward_quantization_update = (
+                ctx.fp8
+                and (ctx.fp8_recipe.delayed() or ctx.fp8_recipe.custom())
+                and requires_grad(
+                    inp, ln_weight, ln_bias, fc1_weight, fc2_weight, fc1_bias, fc2_bias
+                )
+            )
 
             ctx.wgrad_store = wgrad_store
             if is_recomputation:  # return the recomputed tensors
@@ -1806,8 +1803,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.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 56622db5e6..b6e50f039f 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 ---
+    should_request_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,12 @@ 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()
+                recipe = FP8GlobalStateManager.get_fp8_recipe()
+                bwd_args.should_request_backward_quantization_update = (
+                    recipe.delayed() or recipe.custom()
+                )
             if fwd_args.backward_override is not None:
-                bwd_args.reduce_and_update_bwd_fp8_tensors = False
+                bwd_args.should_request_backward_quantization_update = False
 
         return out, new_weight_workspace
 
@@ -1459,15 +1452,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
-        reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors
+        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 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 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 fd66529ba8..3adc2ed2cc 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -227,10 +227,12 @@ 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()
+            recipe = FP8GlobalStateManager.get_fp8_recipe()
+            should_request_backward_quantization_update = (
+                fuser.first_op_requiring_backward < fuser._num_basic_ops
+                and FP8GlobalStateManager.is_fp8_enabled()
+                and (recipe.delayed() or recipe.custom())
+            )
 
             # Other context
             func_ctx.backward_ops = fuser._backward_ops
@@ -243,7 +245,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.is_first_module = is_first_module
+            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(
@@ -383,9 +387,8 @@ def backward(
             for op_idx, input_idx in func_ctx.external_extra_input_slots
         ]
 
-        # 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.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 8c7ea22263..752261fbe3 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -28,12 +28,13 @@
 )
 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
 
 
 __all__ = [
     "autocast",
+    "backward_quantization_update_scope",
     "quantized_model_init",
     "is_fp8_available",
     "is_mxfp8_available",
@@ -400,6 +401,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
+    backward_quantization_update_callback_task_id: Optional[int] = None
+    backward_quantization_update_scope_depth: int = 0
     fp8_graph_capturing: bool = False
     autocast_depth: int = 0
     global_amax_buffer: Dict[str, list] = field(default_factory=dict)
@@ -653,7 +657,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:
@@ -680,6 +684,68 @@ def reduce_and_update_fp8_tensors(
                 qstate.global_scale_buffer[buffer_key],
             )
 
+    # Compatibility alias used by Megatron-Core.
+    reduce_and_update_fp8_tensors = reduce_and_update_quantization_state
+
+    @classmethod
+    def request_backward_quantization_update(cls) -> None:
+        """Request an update after the enclosing logical backward."""
+        qstate = cls.quantization_state
+        qstate.pending_backward_quantization_update = True
+        if qstate.backward_quantization_update_scope_depth == 0:
+            cls._queue_backward_quantization_update_callback()
+
+    @classmethod
+    def _queue_backward_quantization_update_callback(cls, task_id: Optional[int] = None) -> None:
+        """Queue an update after an autograd task."""
+        qstate = cls.quantization_state
+        if task_id is None:
+            task_id = torch._C._current_graph_task_id()
+        if task_id == -1:
+            raise RuntimeError("Backward quantization update must be requested during backward")
+        if qstate.backward_quantization_update_callback_task_id == task_id:
+            return
+
+        qstate.backward_quantization_update_callback_task_id = task_id
+
+        def callback() -> None:
+            cls._run_backward_quantization_update_callback(task_id)
+
+        try:
+            torch.autograd.Variable._execution_engine.queue_callback(callback)
+        except RuntimeError:
+            if qstate.backward_quantization_update_callback_task_id == task_id:
+                qstate.backward_quantization_update_callback_task_id = None
+            raise
+
+    @classmethod
+    def _run_backward_quantization_update_callback(cls, task_id: int) -> None:
+        """Run the update callback for an autograd task."""
+        qstate = cls.quantization_state
+        if qstate.backward_quantization_update_callback_task_id != task_id:
+            return
+        qstate.backward_quantization_update_callback_task_id = None
+        if qstate.backward_quantization_update_scope_depth == 0:
+            cls._run_pending_backward_quantization_update()
+
+    @classmethod
+    def _run_pending_backward_quantization_update(cls) -> None:
+        """Run the pending backward update, if any."""
+        qstate = cls.quantization_state
+        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")
+        update_succeeded = False
+        try:
+            with torch.no_grad():
+                cls.reduce_and_update_quantization_state(forward=False)
+            update_succeeded = True
+        finally:
+            if not update_succeeded:
+                qstate.pending_backward_quantization_update = True
+            nvtx_range_pop("transformer_engine.reduce_and_update_quantization_state.backward")
+
     @staticmethod
     def get_unique_autocast_key(
         recipe: Optional[Recipe] = None,
@@ -749,7 +815,7 @@ def autocast_exit(cls, enabled: bool, _graph: bool) -> None:
         if enabled and qstate.autocast_depth == 0 and not _graph 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:
@@ -813,6 +879,29 @@ def restore_fp8_meta_tensors(fp8_meta: Dict[str, Any]) -> None:
         fp8_meta["scaling_fwd"].scale.copy_(fp8_meta["updated_scale_fwd"])
 
 
+@contextmanager
+def backward_quantization_update_scope() -> None:
+    """Delay the quantization state update until the end of a logical backward.
+
+    Ordinary backward calls update automatically and do not require this scope.
+    Use it when a logical backward spans multiple autograd calls or includes
+    delayed work such as ``module.backward_dw()``. Nested scopes update once
+    when the outermost scope exits.
+    """
+    qstate = FP8GlobalStateManager.quantization_state
+    outermost = qstate.backward_quantization_update_scope_depth == 0
+    task_id = torch._C._current_graph_task_id() if outermost else -1
+    if task_id != -1:
+        FP8GlobalStateManager._queue_backward_quantization_update_callback(task_id)
+    qstate.backward_quantization_update_scope_depth += 1
+    try:
+        yield
+    finally:
+        qstate.backward_quantization_update_scope_depth -= 1
+        if outermost and task_id == -1:
+            FP8GlobalStateManager._run_pending_backward_quantization_update()
+
+
 @contextmanager
 def fp8_model_init(
     enabled: bool = True,

From 238c591c94b3f257242807fb3ee0404da7258f31 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 20 Aug 2026 18:03:22 +0200
Subject: [PATCH 02/13] test: cover independent backward graphs

Document that graphs produced under one autocast need an explicit logical-backward scope when their backward calls should share one delayed-scaling update.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_recipe.py               | 34 +++++++++++++---------
 transformer_engine/pytorch/quantization.py |  7 +++--
 2 files changed, 25 insertions(+), 16 deletions(-)

diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py
index b077d70383..bb0294e889 100644
--- a/tests/pytorch/test_recipe.py
+++ b/tests/pytorch/test_recipe.py
@@ -893,22 +893,30 @@ def test_delayed_scaling_updates_once_per_backward(mode):
 
 
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-def test_backward_quantization_update_scope_groups_autograd_calls():
+def test_backward_quantization_update_scope_groups_independent_graphs():
     FP8GlobalStateManager.reset()
-    model = _make_update_test_model()
+    models = [_make_update_test_model(num_layers=2, seed=seed) for seed in (1, 2)]
+    inputs = [
+        torch.randn(
+            _UPDATE_TEST_BATCH,
+            _UPDATE_TEST_HIDDEN,
+            device="cuda",
+            requires_grad=True,
+        )
+        for _ in models
+    ]
     recipe = DelayedScaling()
 
-    with _UpdateCounter() as counter, te.backward_quantization_update_scope():
-        for _ in range(_UPDATE_TEST_STEPS):
-            x = torch.randn(
-                _UPDATE_TEST_BATCH,
-                _UPDATE_TEST_HIDDEN,
-                device="cuda",
-                requires_grad=True,
-            )
-            _run_update_test_step(model, x, _update_forward_plain, recipe)
-            assert counter.backward == 0
-    assert counter.backward == 1
+    with _UpdateCounter() as counter:
+        with te.autocast(enabled=True, recipe=recipe):
+            outputs = [_run_update_test_layers(model, x) for model, x in zip(models, inputs)]
+        with te.backward_quantization_update_scope():
+            for output in outputs:
+                output.float().sum().backward()
+                assert counter.backward == 0
+        assert counter.backward == 1
+        for x in inputs:
+            assert x.grad is not None
 
 
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 752261fbe3..3da01f1f91 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -884,9 +884,10 @@ def backward_quantization_update_scope() -> None:
     """Delay the quantization state update until the end of a logical backward.
 
     Ordinary backward calls update automatically and do not require this scope.
-    Use it when a logical backward spans multiple autograd calls or includes
-    delayed work such as ``module.backward_dw()``. Nested scopes update once
-    when the outermost scope exits.
+    Use it when a logical backward spans multiple autograd calls, including
+    independent graphs produced under one autocast, or includes delayed work
+    such as ``module.backward_dw()``. Nested scopes update once when the
+    outermost scope exits.
     """
     qstate = FP8GlobalStateManager.quantization_state
     outermost = qstate.backward_quantization_update_scope_depth == 0

From d3e476ed11128facdb789aa5e1d4689648ba23ed Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 1 Sep 2026 11:58:23 +0000
Subject: [PATCH 03/13] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/pytorch/distributed.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 2fc6b85340..bb0e387246 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -446,7 +446,9 @@ def _backward(
         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
-        ), autocast(enabled=ctx.fp8, recipe=ctx.fp8_recipe):
+        ), autocast(
+            enabled=ctx.fp8, recipe=ctx.fp8_recipe
+        ):
             outputs = ctx.run_function(*detached_inputs, **ctx.kwargs)
 
         # Set the states back to what it was at the start of this function.

From 3d0988dc5618fa9fbbf8f6f9021e8fc72dba91c5 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 1 Sep 2026 14:15:10 +0200
Subject: [PATCH 04/13] Rename backward_quantization_update_scope to
 quantization_backward_scope

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 docs/api/pytorch.rst                       |  2 +-
 tests/pytorch/test_recipe.py               |  8 ++++----
 transformer_engine/pytorch/__init__.py     |  2 +-
 transformer_engine/pytorch/distributed.py  |  4 ++--
 transformer_engine/pytorch/quantization.py | 16 ++++++++--------
 5 files changed, 16 insertions(+), 16 deletions(-)

diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst
index 3115fb027f..84d850eebc 100644
--- a/docs/api/pytorch.rst
+++ b/docs/api/pytorch.rst
@@ -40,7 +40,7 @@ PyTorch
 
 .. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None)
 
-.. autoapifunction:: transformer_engine.pytorch.backward_quantization_update_scope
+.. autoapifunction:: transformer_engine.pytorch.quantization_backward_scope
 
 .. autoapifunction:: transformer_engine.pytorch.quantized_model_init
 
diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py
index bb0294e889..a5d5fd6df8 100644
--- a/tests/pytorch/test_recipe.py
+++ b/tests/pytorch/test_recipe.py
@@ -893,7 +893,7 @@ def test_delayed_scaling_updates_once_per_backward(mode):
 
 
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-def test_backward_quantization_update_scope_groups_independent_graphs():
+def test_quantization_backward_scope_groups_independent_graphs():
     FP8GlobalStateManager.reset()
     models = [_make_update_test_model(num_layers=2, seed=seed) for seed in (1, 2)]
     inputs = [
@@ -910,7 +910,7 @@ def test_backward_quantization_update_scope_groups_independent_graphs():
     with _UpdateCounter() as counter:
         with te.autocast(enabled=True, recipe=recipe):
             outputs = [_run_update_test_layers(model, x) for model, x in zip(models, inputs)]
-        with te.backward_quantization_update_scope():
+        with te.quantization_backward_scope():
             for output in outputs:
                 output.float().sum().backward()
                 assert counter.backward == 0
@@ -920,7 +920,7 @@ def test_backward_quantization_update_scope_groups_independent_graphs():
 
 
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-def test_backward_quantization_update_scope_covers_delayed_wgrad():
+def test_quantization_backward_scope_covers_delayed_wgrad():
     FP8GlobalStateManager.reset()
     model = te.Linear(
         _UPDATE_TEST_HIDDEN,
@@ -930,7 +930,7 @@ def test_backward_quantization_update_scope_covers_delayed_wgrad():
     ).cuda()
     recipe = DelayedScaling()
 
-    with _UpdateCounter() as counter, te.backward_quantization_update_scope():
+    with _UpdateCounter() as counter, te.quantization_backward_scope():
         x = torch.randn(
             _UPDATE_TEST_BATCH,
             _UPDATE_TEST_HIDDEN,
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index 26f10bef9f..286eef2561 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/transformer_engine/pytorch/__init__.py
@@ -45,7 +45,7 @@
 from transformer_engine.pytorch.quantization import fp8_autocast
 from transformer_engine.pytorch.quantization import fp8_model_init
 from transformer_engine.pytorch.quantization import autocast
-from transformer_engine.pytorch.quantization import backward_quantization_update_scope
+from transformer_engine.pytorch.quantization import quantization_backward_scope
 from transformer_engine.pytorch.quantization import quantized_model_init
 from transformer_engine.pytorch.quantization import is_fp8_available
 from transformer_engine.pytorch.quantization import is_mxfp8_available
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index bb0e387246..a2d5a621b6 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -39,7 +39,7 @@
 )
 
 from .constants import dist_group_type
-from .quantization import FP8GlobalStateManager, autocast, backward_quantization_update_scope
+from .quantization import FP8GlobalStateManager, autocast, quantization_backward_scope
 from .tensor.float8_tensor import Float8Quantizer, Float8Tensor, Float8CurrentScalingQuantizer
 from .tensor.mxfp8_tensor import MXFP8Quantizer
 from .tensor.nvfp4_tensor import NVFP4Quantizer
@@ -401,7 +401,7 @@ def backward(
         """Call backward function with activation recomputation."""
         recipe = ctx.fp8_recipe
         update_scope = (
-            backward_quantization_update_scope()
+            quantization_backward_scope()
             if ctx.fp8 and (recipe.delayed() or recipe.custom())
             else nullcontext()
         )
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 3da01f1f91..d0078d7612 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -34,7 +34,7 @@
 
 __all__ = [
     "autocast",
-    "backward_quantization_update_scope",
+    "quantization_backward_scope",
     "quantized_model_init",
     "is_fp8_available",
     "is_mxfp8_available",
@@ -403,7 +403,7 @@ class FP8GlobalState:
     is_first_fp8_module: bool = False
     pending_backward_quantization_update: bool = False
     backward_quantization_update_callback_task_id: Optional[int] = None
-    backward_quantization_update_scope_depth: int = 0
+    quantization_backward_scope_depth: int = 0
     fp8_graph_capturing: bool = False
     autocast_depth: int = 0
     global_amax_buffer: Dict[str, list] = field(default_factory=dict)
@@ -692,7 +692,7 @@ def request_backward_quantization_update(cls) -> None:
         """Request an update after the enclosing logical backward."""
         qstate = cls.quantization_state
         qstate.pending_backward_quantization_update = True
-        if qstate.backward_quantization_update_scope_depth == 0:
+        if qstate.quantization_backward_scope_depth == 0:
             cls._queue_backward_quantization_update_callback()
 
     @classmethod
@@ -725,7 +725,7 @@ def _run_backward_quantization_update_callback(cls, task_id: int) -> None:
         if qstate.backward_quantization_update_callback_task_id != task_id:
             return
         qstate.backward_quantization_update_callback_task_id = None
-        if qstate.backward_quantization_update_scope_depth == 0:
+        if qstate.quantization_backward_scope_depth == 0:
             cls._run_pending_backward_quantization_update()
 
     @classmethod
@@ -880,7 +880,7 @@ def restore_fp8_meta_tensors(fp8_meta: Dict[str, Any]) -> None:
 
 
 @contextmanager
-def backward_quantization_update_scope() -> None:
+def quantization_backward_scope() -> None:
     """Delay the quantization state update until the end of a logical backward.
 
     Ordinary backward calls update automatically and do not require this scope.
@@ -890,15 +890,15 @@ def backward_quantization_update_scope() -> None:
     outermost scope exits.
     """
     qstate = FP8GlobalStateManager.quantization_state
-    outermost = qstate.backward_quantization_update_scope_depth == 0
+    outermost = qstate.quantization_backward_scope_depth == 0
     task_id = torch._C._current_graph_task_id() if outermost else -1
     if task_id != -1:
         FP8GlobalStateManager._queue_backward_quantization_update_callback(task_id)
-    qstate.backward_quantization_update_scope_depth += 1
+    qstate.quantization_backward_scope_depth += 1
     try:
         yield
     finally:
-        qstate.backward_quantization_update_scope_depth -= 1
+        qstate.quantization_backward_scope_depth -= 1
         if outermost and task_id == -1:
             FP8GlobalStateManager._run_pending_backward_quantization_update()
 

From e65294e84fda8f4da6c0ef01379145e529f298bf Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 11:17:46 +0200
Subject: [PATCH 05/13] Centralize backward update request and make
 quantization_backward_scope always update

Modules call FP8GlobalStateManager.request_backward_quantization_update(recipe)
from backward whenever they ran in FP8; the recipe and graph-capture checks
live in that helper instead of being repeated per module.

quantization_backward_scope now marks the update pending on entry, so ranks
that ran no quantized backward inside the scope still join the amax
reduction. Add a distributed test covering a module skipped on some ranks and
a rank with no backward inside the scope.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 qa/L1_pytorch_distributed_unittest/test.sh    |   1 +
 .../distributed/run_backward_update_ranks.py  | 110 ++++++++++++++++++
 .../distributed/test_backward_update_ranks.py |  28 +++++
 tests/pytorch/test_backward_override.py       |  38 ++----
 .../pytorch/module/grouped_linear.py          |  21 +---
 .../pytorch/module/layernorm_linear.py        |  12 +-
 .../pytorch/module/layernorm_mlp.py           |  14 +--
 transformer_engine/pytorch/module/linear.py   |  25 +---
 transformer_engine/pytorch/ops/fuser.py       |  37 ++----
 transformer_engine/pytorch/quantization.py    |  20 +++-
 10 files changed, 184 insertions(+), 122 deletions(-)
 create mode 100644 tests/pytorch/distributed/run_backward_update_ranks.py
 create mode 100644 tests/pytorch/distributed/test_backward_update_ranks.py

diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh
index ec19492ee7..c127d559b3 100644
--- a/qa/L1_pytorch_distributed_unittest/test.sh
+++ b/qa/L1_pytorch_distributed_unittest/test.sh
@@ -45,6 +45,7 @@ fi
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py"
+python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_backward_update_ranks.xml $TE_PATH/tests/pytorch/distributed/test_backward_update_ranks.py || test_fail "test_backward_update_ranks.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py -k "not hybrid" || test_fail "test_torch_fsdp2.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py"
diff --git a/tests/pytorch/distributed/run_backward_update_ranks.py b/tests/pytorch/distributed/run_backward_update_ranks.py
new file mode 100644
index 0000000000..4f7da74446
--- /dev/null
+++ b/tests/pytorch/distributed/run_backward_update_ranks.py
@@ -0,0 +1,110 @@
+#!/usr/bin/python3
+
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Delayed-scaling backward update when some ranks skip backward.
+
+Every rank runs forward for every module so the amax buffers match across
+ranks, but only some ranks run backward. The update must still happen once
+per step on every rank or the amax all-reduce hangs.
+"""
+
+import argparse
+import datetime
+import os
+
+import torch
+import torch.distributed as dist
+
+import transformer_engine.pytorch as te
+from transformer_engine.common.recipe import DelayedScaling
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
+
+BATCH, HIDDEN, STEPS = 32, 128, 3
+
+
+def _make_model(seed):
+    torch.manual_seed(seed)
+    return torch.nn.ModuleList([te.Linear(HIDDEN, HIDDEN, bias=True) for _ in range(2)]).cuda()
+
+
+def _step_skipped_module_backward(model, recipe, rank):
+    """Odd ranks feed the first module an empty batch and drop its output."""
+    rows = BATCH if rank % 2 == 0 else 0
+    x_a = torch.randn(rows, HIDDEN, device="cuda", requires_grad=True)
+    x_b = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+    with te.autocast(enabled=True, recipe=recipe):
+        y_a = model[0](x_a)
+        y_b = model[1](x_b)
+    loss = y_b.float().sum()
+    if y_a.numel() > 0:
+        loss = loss + y_a.float().sum()
+    loss.backward()
+
+
+def _step_no_backward_in_scope(model, recipe, rank):
+    """Odd ranks run no backward at all; the scope still triggers the update."""
+    x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+    with te.quantization_backward_scope():
+        with te.autocast(enabled=True, recipe=recipe):
+            y = model[1](model[0](x))
+        if rank % 2 == 0:
+            y.float().sum().backward()
+
+
+_CASES = {
+    "skipped_module_backward": _step_skipped_module_backward,
+    "no_backward_in_scope": _step_no_backward_in_scope,
+}
+
+
+def _bwd_state(model):
+    tensors = []
+    for module in model:
+        state = module.fp8_meta["scaling_bwd"]
+        tensors += [state.amax_history.clone(), state.scale.clone()]
+    return tensors
+
+
+def _assert_same_on_all_ranks(tensors, world_size):
+    for t in tensors:
+        gathered = [torch.empty_like(t) for _ in range(world_size)]
+        dist.all_gather(gathered, t)
+        for other in gathered[1:]:
+            assert torch.equal(other, gathered[0]), f"{gathered[0]} vs {other}"
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--case", choices=_CASES.keys(), required=True)
+    parser.add_argument("--backend", default="nccl")
+    args = parser.parse_args()
+
+    local_rank = int(os.getenv("LOCAL_RANK", "0"))
+    torch.cuda.set_device(local_rank % torch.cuda.device_count())
+    dist.init_process_group(backend=args.backend, timeout=datetime.timedelta(seconds=120))
+    rank, world_size = dist.get_rank(), dist.get_world_size()
+
+    model = _make_model(seed=1234)
+    recipe = DelayedScaling(reduce_amax=True)
+    step = _CASES[args.case]
+    qstate = FP8GlobalStateManager.quantization_state
+
+    for _ in range(STEPS):
+        model.zero_grad(set_to_none=True)
+        step(model, recipe, rank)
+        assert not qstate.pending_backward_quantization_update
+        assert qstate.backward_quantization_update_callback_task_id is None
+        _assert_same_on_all_ranks(_bwd_state(model), world_size)
+
+    # Ranks that skipped backward must have received the other ranks' amaxes.
+    for module in model:
+        assert module.fp8_meta["scaling_bwd"].amax_history.abs().sum() > 0
+
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/pytorch/distributed/test_backward_update_ranks.py b/tests/pytorch/distributed/test_backward_update_ranks.py
new file mode 100644
index 0000000000..b2987ee7d9
--- /dev/null
+++ b/tests/pytorch/distributed/test_backward_update_ranks.py
@@ -0,0 +1,28 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import os
+import subprocess
+from pathlib import Path
+
+import pytest
+import torch
+import transformer_engine.pytorch as te
+
+if torch.cuda.device_count() < 2:
+    pytest.skip("Distributed training needs at least 2 GPUs.", allow_module_level=True)
+
+fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)
+
+TEST_ROOT = Path(__file__).parent.resolve()
+NUM_PROCS: int = min(4, torch.cuda.device_count())
+LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"]
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("case", ["skipped_module_backward", "no_backward_in_scope"])
+def test_backward_update_ranks(case):
+    test_cmd = LAUNCH_CMD + [str(TEST_ROOT / "run_backward_update_ranks.py"), "--case", case]
+    result = subprocess.run(test_cmd, env=os.environ, check=False, timeout=600)
+    assert result.returncode == 0
diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py
index 00bf5a8f2e..3dc308ed72 100644
--- a/tests/pytorch/test_backward_override.py
+++ b/tests/pytorch/test_backward_override.py
@@ -408,7 +408,7 @@ def _snapshot_layout_invariants(
 
 def _snapshot_backward_ctx_state(
     output: torch.Tensor,
-) -> tuple[str, bool, object, bool]:
+) -> tuple[str, bool, object]:
     if output.grad_fn is None:
         raise RuntimeError("Output tensor has no grad_fn; cannot inspect backward context state.")
     # ``Linear`` packs backward state into ``grad_fn.backward_objects``
@@ -419,7 +419,6 @@ def _snapshot_backward_ctx_state(
         "backward_override",
         "fp8",
         "grad_output_quantizer",
-        "should_request_backward_quantization_update",
     )
     missing_attrs = [attr for attr in required_attrs if not hasattr(state_holder, attr)]
     if missing_attrs:
@@ -430,7 +429,6 @@ 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, "should_request_backward_quantization_update")),
     )
 
 
@@ -816,7 +814,6 @@ def _run_grouped_linear_single_step_with_ctx_state(
         required_attrs = (
             "backward_override",
             "fp8",
-            "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 +824,6 @@ 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, "should_request_backward_quantization_update")),
         )
     y.backward(dy)
     assert x_run.grad is not None
@@ -1449,38 +1445,22 @@ def test_linear_like_runtime_backward_override_switch_updates_ctx(
     skip_unsupported_backward_override(module_type, mode_recipe, backward_override)
 
     *_, default_ctx = _run_single_step_with_ctx_state(module, x, dy, default_recipe)
-    (
-        default_mode,
-        default_fp8,
-        default_grad_output_quantizer,
-        default_should_request_update,
-    ) = default_ctx
-    expected_request = default_recipe.delayed() or default_recipe.custom()
+    default_mode, default_fp8, default_grad_output_quantizer = default_ctx
     assert default_mode is None
     assert default_fp8
     assert default_grad_output_quantizer is not None
-    assert default_should_request_update == expected_request
 
     *_, switched_ctx = _run_single_step_with_ctx_state(module, x, dy, mode_recipe)
-    switched_mode, switched_fp8, switched_grad_output_quantizer, switched_should_request_update = (
-        switched_ctx
-    )
+    switched_mode, switched_fp8, switched_grad_output_quantizer = switched_ctx
     assert switched_mode == backward_override
     assert not switched_fp8
     assert switched_grad_output_quantizer is None
-    assert not switched_should_request_update
 
     *_, default_ctx_after = _run_single_step_with_ctx_state(module, x, dy, default_recipe)
-    (
-        default_mode_after,
-        default_fp8_after,
-        default_grad_output_quantizer_after,
-        default_should_request_update_after,
-    ) = default_ctx_after
+    default_mode_after, default_fp8_after, default_grad_output_quantizer_after = default_ctx_after
     assert default_mode_after is None
     assert default_fp8_after
     assert default_grad_output_quantizer_after is not None
-    assert default_should_request_update_after == expected_request
 
 
 @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
@@ -1527,11 +1507,9 @@ def test_grouped_linear_runtime_backward_override_switch_updates_ctx(
         dy,
         default_recipe,
     )
-    default_mode, default_fp8, default_should_request_update = default_ctx
-    expected_request = default_recipe.delayed() or default_recipe.custom()
+    default_mode, default_fp8 = default_ctx
     assert default_mode is None
     assert default_fp8
-    assert default_should_request_update == expected_request
 
     *_, switched_ctx = _run_grouped_linear_single_step_with_ctx_state(
         module,
@@ -1540,10 +1518,9 @@ def test_grouped_linear_runtime_backward_override_switch_updates_ctx(
         dy,
         mode_recipe,
     )
-    switched_mode, switched_fp8, switched_should_request_update = switched_ctx
+    switched_mode, switched_fp8 = switched_ctx
     assert switched_mode == backward_override
     assert not switched_fp8
-    assert not switched_should_request_update
 
     *_, default_ctx_after = _run_grouped_linear_single_step_with_ctx_state(
         module,
@@ -1552,10 +1529,9 @@ def test_grouped_linear_runtime_backward_override_switch_updates_ctx(
         dy,
         default_recipe,
     )
-    default_mode_after, default_fp8_after, default_should_request_update_after = default_ctx_after
+    default_mode_after, default_fp8_after = default_ctx_after
     assert default_mode_after is None
     assert default_fp8_after
-    assert default_should_request_update_after == expected_request
 
 
 @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 5cae0bb265..010d6c661b 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -40,7 +40,6 @@
     clear_tensor_data,
     get_device_compute_capability,
     init_method_constant,
-    requires_grad,
     resolve_grouped_linear_single_param_flags,
     get_nvtx_range_context,
 )
@@ -50,7 +49,6 @@
     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,
@@ -600,11 +598,6 @@ def _forward_grouped_tensor(
             ctx.use_bias = use_bias
             ctx.inp_shape = inp.shape
             ctx.requires_dgrad = inp.requires_grad
-            ctx.should_request_backward_quantization_update = (
-                ctx.fp8
-                and (ctx.fp8_recipe.delayed() or ctx.fp8_recipe.custom())
-                and requires_grad(inp, weights[0], biases[0])
-            )
             ctx.wgrad_store = wgrad_store
             ctx.debug = False
             ctx.save_original_input = save_original_input
@@ -975,11 +968,6 @@ def forward(
             ctx.sequence_parallel = sequence_parallel
             ctx.inp_shape = inp.shape
             ctx.requires_dgrad = inp.requires_grad
-            ctx.should_request_backward_quantization_update = (
-                ctx.fp8
-                and (ctx.fp8_recipe.delayed() or ctx.fp8_recipe.custom())
-                and requires_grad(inp, weights[0], biases[0])
-            )
             ctx.wgrad_store = wgrad_store
             ctx.debug = debug
             ctx.save_original_input = save_original_input
@@ -997,7 +985,6 @@ def forward(
                 ctx.grad_input_quantizers = [None] * num_gemms
                 ctx.grad_weight_quantizers = [None] * num_gemms
                 ctx.grad_output_quantizers = [None] * num_gemms
-                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
@@ -1249,8 +1236,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
         else:
             wgrad_list = [None] * num_weight_args
 
-        if ctx.should_request_backward_quantization_update and not is_graph_capturing():
-            FP8GlobalStateManager.request_backward_quantization_update()
+        if ctx.fp8:
+            FP8GlobalStateManager.request_backward_quantization_update(ctx.fp8_recipe)
         return (
             dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
             None,  # m_splits
@@ -1516,8 +1503,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
             ):
                 grad_biases = [None] * ctx.num_gemms
 
-        if ctx.should_request_backward_quantization_update and not is_graph_capturing():
-            FP8GlobalStateManager.request_backward_quantization_update()
+        if ctx.fp8:
+            FP8GlobalStateManager.request_backward_quantization_update(ctx.fp8_recipe)
         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 1ed8326572..145fe4cbf9 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -41,7 +41,6 @@
     init_method_constant,
     nvtx_range_pop,
     nvtx_range_push,
-    requires_grad,
     needs_quantized_gemm,
     get_nvtx_range_context,
 )
@@ -63,7 +62,6 @@
 )
 from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type
 from ..jit import no_torch_dynamo
-from ..graph import is_graph_capturing
 from ._common import (
     apply_normalization,
     noop_cat,
@@ -587,11 +585,6 @@ def forward(
             ctx.ub_name = ub_name
             ctx.requires_dgrad = inp_requires_grad
             ctx.normalization = normalization
-            ctx.should_request_backward_quantization_update = (
-                ctx.fp8
-                and (ctx.fp8_recipe.delayed() or ctx.fp8_recipe.custom())
-                and requires_grad(inp, ln_weight, ln_bias, weight, bias)
-            )
             ctx.wgrad_store = wgrad_store
             ctx.debug = debug
 
@@ -606,7 +599,6 @@ def forward(
                 ctx.grad_input_quantizer = None
                 ctx.grad_weight_quantizer = None
                 ctx.grad_output_quantizer = None
-                ctx.should_request_backward_quantization_update = False
 
         # ------------------------------------------------------
         # Cached state for backward pass is ready...
@@ -1181,8 +1173,8 @@ def wgrad_gemm(
         else:
             wgrad = None
 
-        if ctx.should_request_backward_quantization_update and not is_graph_capturing():
-            FP8GlobalStateManager.request_backward_quantization_update()
+        if ctx.fp8:
+            FP8GlobalStateManager.request_backward_quantization_update(ctx.fp8_recipe)
 
         # 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 3939f7e21d..0b3f7558a5 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -46,7 +46,6 @@
     cast_if_needed,
     assert_dim_for_fp8_exec,
     clear_tensor_data,
-    requires_grad,
     needs_quantized_gemm,
     get_nvtx_range_context,
 )
@@ -64,7 +63,6 @@
 )
 from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, dist_group_type
 from ..jit import no_torch_dynamo
-from ..graph import is_graph_capturing
 from ..tensor.float8_tensor import Float8Tensor
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
 from ..tensor.nvfp4_tensor import NVFP4Quantizer
@@ -907,14 +905,6 @@ def _forward(
                 inp.requires_grad or ln_weight.requires_grad or ln_bias.requires_grad
             )
             ctx.normalization = normalization
-            ctx.should_request_backward_quantization_update = (
-                ctx.fp8
-                and (ctx.fp8_recipe.delayed() or ctx.fp8_recipe.custom())
-                and requires_grad(
-                    inp, ln_weight, ln_bias, fc1_weight, fc2_weight, fc1_bias, fc2_bias
-                )
-            )
-
             ctx.wgrad_store = wgrad_store
             if is_recomputation:  # return the recomputed tensors
                 return (
@@ -1803,8 +1793,8 @@ def fc1_wgrad_gemm(
         else:
             fc2_wgrad = None
 
-        if ctx.should_request_backward_quantization_update and not is_graph_capturing():
-            FP8GlobalStateManager.request_backward_quantization_update()
+        if ctx.fp8:
+            FP8GlobalStateManager.request_backward_quantization_update(ctx.fp8_recipe)
 
         # FIX THIS
         # Scatter Fp8 tranposed-weight buffers
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index b6e50f039f..9631f5e7a2 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -72,7 +72,6 @@
 )
 from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type
 from ..jit import no_torch_dynamo
-from ..graph import is_graph_capturing
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
@@ -198,6 +197,7 @@ class LinearBwdArgs:
     # --- Numerical / dtype config ---
     activation_dtype: Optional[torch.dtype] = None
     fp8: bool = False
+    fp8_recipe: Optional[Recipe] = None
     dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD
     wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD
     backward_override: Optional[str] = None
@@ -233,9 +233,6 @@ class LinearBwdArgs:
     origin_weight_overwrites_main_grad: bool = False
     main_grad_func: Optional[Callable[[], torch.Tensor]] = None
 
-    # --- Quantization state update bookkeeping ---
-    should_request_backward_quantization_update: bool = False
-
     # --- Misc ---
     cpu_offloading: bool = False
     owns_input: bool = False
@@ -705,6 +702,7 @@ def _linear_setup_ctx(
     # Numerical / dtype config
     bwd_args.activation_dtype = fwd_args.activation_dtype
     bwd_args.fp8 = fp8
+    bwd_args.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
     bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator
     bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator
     bwd_args.backward_override = backward_override
@@ -1424,17 +1422,6 @@ def forward(
             ctx.save_for_backward(*tensors_to_save)
             ctx.tensor_objects = tensor_objects
             ctx.backward_objects = bwd_args
-            if fwd_args.fp8 and (
-                fwd_args.input_requires_grad
-                or fwd_args.weight_requires_grad
-                or fwd_args.bias_requires_grad
-            ):
-                recipe = FP8GlobalStateManager.get_fp8_recipe()
-                bwd_args.should_request_backward_quantization_update = (
-                    recipe.delayed() or recipe.custom()
-                )
-            if fwd_args.backward_override is not None:
-                bwd_args.should_request_backward_quantization_update = False
 
         return out, new_weight_workspace
 
@@ -1452,15 +1439,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
-        should_request_backward_quantization_update = (
-            bwd_args.should_request_backward_quantization_update
-        )
+        fp8_recipe = bwd_args.fp8_recipe if bwd_args.fp8 else None
         # 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 should_request_backward_quantization_update and not is_graph_capturing():
-            FP8GlobalStateManager.request_backward_quantization_update()
+        if fp8_recipe is not None:
+            FP8GlobalStateManager.request_backward_quantization_update(fp8_recipe)
         return result
 
 
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 3adc2ed2cc..9d8ab42662 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -5,7 +5,7 @@
 """Manager class for a pipeline of fusible operations."""
 
 from __future__ import annotations
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Iterable, Sequence
 import itertools
 from typing import Any, Optional, TypeAlias
 
@@ -26,24 +26,6 @@ def _split_tuple(t: tuple, idx: int) -> tuple[tuple, tuple]:
     return t[:idx], t[idx:]
 
 
-# Lazily imported function used in _is_graph_capturing
-_is_graph_capturing_function: Optional[Callable[[], bool]] = None
-
-
-def _is_graph_capturing() -> bool:
-    """Whether function is called within ``make_graphed_callables``
-
-    Avoid circular import with lazy import.
-
-    """
-    global _is_graph_capturing_function
-    if _is_graph_capturing_function is None:
-        from ..graph import is_graph_capturing
-
-        _is_graph_capturing_function = is_graph_capturing
-    return _is_graph_capturing_function()
-
-
 # Type alias for a function that may perform operation fusion
 OperationFusionFunction: TypeAlias = (
     "Callable[tuple[list[FusibleOperation], ...], list[FusibleOperation]]"
@@ -227,13 +209,6 @@ def forward(
             func_ctx.save_for_backward(*tensors_to_save)
             func_ctx.tensor_objects = tensor_objects
 
-            recipe = FP8GlobalStateManager.get_fp8_recipe()
-            should_request_backward_quantization_update = (
-                fuser.first_op_requiring_backward < fuser._num_basic_ops
-                and FP8GlobalStateManager.is_fp8_enabled()
-                and (recipe.delayed() or recipe.custom())
-            )
-
             # Other context
             func_ctx.backward_ops = fuser._backward_ops
             func_ctx.basic_ops = fuser._basic_ops
@@ -245,8 +220,10 @@ 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.should_request_backward_quantization_update = (
-                should_request_backward_quantization_update
+            func_ctx.fp8_recipe = (
+                FP8GlobalStateManager.get_fp8_recipe()
+                if FP8GlobalStateManager.is_fp8_enabled()
+                else None
             )
 
         # Mark output tensors as not deletable in backward
@@ -387,8 +364,8 @@ def backward(
             for op_idx, input_idx in func_ctx.external_extra_input_slots
         ]
 
-        if func_ctx.should_request_backward_quantization_update and not _is_graph_capturing():
-            FP8GlobalStateManager.request_backward_quantization_update()
+        if func_ctx.fp8_recipe is not None:
+            FP8GlobalStateManager.request_backward_quantization_update(func_ctx.fp8_recipe)
 
         return (
             dx,  # input_
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index d0078d7612..310a5eed46 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -688,8 +688,18 @@ def reduce_and_update_quantization_state(
     reduce_and_update_fp8_tensors = reduce_and_update_quantization_state
 
     @classmethod
-    def request_backward_quantization_update(cls) -> None:
-        """Request an update after the enclosing logical backward."""
+    def request_backward_quantization_update(cls, recipe: Recipe) -> None:
+        """Request an update after the enclosing logical backward.
+
+        No-op for recipes without delayed-scaling state and inside CUDA graph
+        capture, where the graphed wrapper performs the update.
+        """
+        if not (recipe.delayed() or recipe.custom()):
+            return
+        from .graph import is_graph_capturing  # pylint: disable=import-outside-toplevel
+
+        if is_graph_capturing():
+            return
         qstate = cls.quantization_state
         qstate.pending_backward_quantization_update = True
         if qstate.quantization_backward_scope_depth == 0:
@@ -888,9 +898,15 @@ def quantization_backward_scope() -> None:
     independent graphs produced under one autocast, or includes delayed work
     such as ``module.backward_dw()``. Nested scopes update once when the
     outermost scope exits.
+
+    The update runs on scope exit even if no quantized module ran backward
+    inside it, so ranks that skipped every backward (e.g. an expert with no
+    tokens) still participate in the amax reduction.
     """
     qstate = FP8GlobalStateManager.quantization_state
     outermost = qstate.quantization_backward_scope_depth == 0
+    if outermost:
+        qstate.pending_backward_quantization_update = True
     task_id = torch._C._current_graph_task_id() if outermost else -1
     if task_id != -1:
         FP8GlobalStateManager._queue_backward_quantization_update_callback(task_id)

From 2edba6af02d014ec9cbec50d7ca9d04add2d4dc4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 11:24:07 +0200
Subject: [PATCH 06/13] Move skipped-rank backward update test into
 run_numerics.py

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 qa/L1_pytorch_distributed_unittest/test.sh    |   1 -
 .../distributed/run_backward_update_ranks.py  | 110 ------------------
 tests/pytorch/distributed/run_numerics.py     |  68 +++++++++++
 .../distributed/test_backward_update_ranks.py |  28 -----
 4 files changed, 68 insertions(+), 139 deletions(-)
 delete mode 100644 tests/pytorch/distributed/run_backward_update_ranks.py
 delete mode 100644 tests/pytorch/distributed/test_backward_update_ranks.py

diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh
index c127d559b3..ec19492ee7 100644
--- a/qa/L1_pytorch_distributed_unittest/test.sh
+++ b/qa/L1_pytorch_distributed_unittest/test.sh
@@ -45,7 +45,6 @@ fi
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py"
-python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_backward_update_ranks.xml $TE_PATH/tests/pytorch/distributed/test_backward_update_ranks.py || test_fail "test_backward_update_ranks.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops.py || test_fail "test_fusible_ops.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py -k "not hybrid" || test_fail "test_torch_fsdp2.py"
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py"
diff --git a/tests/pytorch/distributed/run_backward_update_ranks.py b/tests/pytorch/distributed/run_backward_update_ranks.py
deleted file mode 100644
index 4f7da74446..0000000000
--- a/tests/pytorch/distributed/run_backward_update_ranks.py
+++ /dev/null
@@ -1,110 +0,0 @@
-#!/usr/bin/python3
-
-# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-#
-# See LICENSE for license information.
-
-"""Delayed-scaling backward update when some ranks skip backward.
-
-Every rank runs forward for every module so the amax buffers match across
-ranks, but only some ranks run backward. The update must still happen once
-per step on every rank or the amax all-reduce hangs.
-"""
-
-import argparse
-import datetime
-import os
-
-import torch
-import torch.distributed as dist
-
-import transformer_engine.pytorch as te
-from transformer_engine.common.recipe import DelayedScaling
-from transformer_engine.pytorch.quantization import FP8GlobalStateManager
-
-BATCH, HIDDEN, STEPS = 32, 128, 3
-
-
-def _make_model(seed):
-    torch.manual_seed(seed)
-    return torch.nn.ModuleList([te.Linear(HIDDEN, HIDDEN, bias=True) for _ in range(2)]).cuda()
-
-
-def _step_skipped_module_backward(model, recipe, rank):
-    """Odd ranks feed the first module an empty batch and drop its output."""
-    rows = BATCH if rank % 2 == 0 else 0
-    x_a = torch.randn(rows, HIDDEN, device="cuda", requires_grad=True)
-    x_b = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
-    with te.autocast(enabled=True, recipe=recipe):
-        y_a = model[0](x_a)
-        y_b = model[1](x_b)
-    loss = y_b.float().sum()
-    if y_a.numel() > 0:
-        loss = loss + y_a.float().sum()
-    loss.backward()
-
-
-def _step_no_backward_in_scope(model, recipe, rank):
-    """Odd ranks run no backward at all; the scope still triggers the update."""
-    x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
-    with te.quantization_backward_scope():
-        with te.autocast(enabled=True, recipe=recipe):
-            y = model[1](model[0](x))
-        if rank % 2 == 0:
-            y.float().sum().backward()
-
-
-_CASES = {
-    "skipped_module_backward": _step_skipped_module_backward,
-    "no_backward_in_scope": _step_no_backward_in_scope,
-}
-
-
-def _bwd_state(model):
-    tensors = []
-    for module in model:
-        state = module.fp8_meta["scaling_bwd"]
-        tensors += [state.amax_history.clone(), state.scale.clone()]
-    return tensors
-
-
-def _assert_same_on_all_ranks(tensors, world_size):
-    for t in tensors:
-        gathered = [torch.empty_like(t) for _ in range(world_size)]
-        dist.all_gather(gathered, t)
-        for other in gathered[1:]:
-            assert torch.equal(other, gathered[0]), f"{gathered[0]} vs {other}"
-
-
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--case", choices=_CASES.keys(), required=True)
-    parser.add_argument("--backend", default="nccl")
-    args = parser.parse_args()
-
-    local_rank = int(os.getenv("LOCAL_RANK", "0"))
-    torch.cuda.set_device(local_rank % torch.cuda.device_count())
-    dist.init_process_group(backend=args.backend, timeout=datetime.timedelta(seconds=120))
-    rank, world_size = dist.get_rank(), dist.get_world_size()
-
-    model = _make_model(seed=1234)
-    recipe = DelayedScaling(reduce_amax=True)
-    step = _CASES[args.case]
-    qstate = FP8GlobalStateManager.quantization_state
-
-    for _ in range(STEPS):
-        model.zero_grad(set_to_none=True)
-        step(model, recipe, rank)
-        assert not qstate.pending_backward_quantization_update
-        assert qstate.backward_quantization_update_callback_task_id is None
-        _assert_same_on_all_ranks(_bwd_state(model), world_size)
-
-    # Ranks that skipped backward must have received the other ranks' amaxes.
-    for module in model:
-        assert module.fp8_meta["scaling_bwd"].amax_history.abs().sum() > 0
-
-    dist.destroy_process_group()
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py
index fe02f990b4..51041b7b5d 100644
--- a/tests/pytorch/distributed/run_numerics.py
+++ b/tests/pytorch/distributed/run_numerics.py
@@ -28,6 +28,7 @@
 from transformer_engine.pytorch import Float8CurrentScalingQuantizer, NVFP4Quantizer
 from transformer_engine.pytorch.constants import NVFP4_BLOCK_SCALING_SIZE
 from transformer_engine.pytorch.distributed import gather_along_first_dim
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
 from run_layer_with_overlap import _compare_tensors
 
 SEQ_LEN, BATCH_SIZE = 16, 16
@@ -132,6 +133,7 @@ def main(argv=None, namespace=None):
         test_layernorm_linear,
         test_layernorm_mlp,
         test_transformer_layer,
+        test_backward_update_with_skipped_ranks,
     ]
 
     for test in test_dict:
@@ -1139,5 +1141,71 @@ def test_transformer_layer():
             _test_transformer_layer_parallel(sequence_parallel, **kwargs)
 
 
+############################################
+#      Delayed-scaling backward update     #
+############################################
+
+
+def _assert_bwd_state_matches_across_ranks(model):
+    for module in model:
+        state = module.fp8_meta["scaling_bwd"]
+        for t in (state.amax_history, state.scale):
+            gathered = [torch.empty_like(t) for _ in range(WORLD_SIZE)]
+            dist.all_gather(gathered, t)
+            for other in gathered[1:]:
+                assert torch.equal(other, gathered[0]), f"{gathered[0]} vs {other}"
+
+
+def _backward_update_step_skipped_module(model, recipe):
+    """Odd ranks feed the first module an empty batch and drop its output."""
+    rows = BATCH_SIZE if WORLD_RANK % 2 == 0 else 0
+    x_a = torch.randn(rows, HIDDEN_SIZE, device="cuda", requires_grad=True)
+    x_b = torch.randn(BATCH_SIZE, HIDDEN_SIZE, device="cuda", requires_grad=True)
+    with te.autocast(enabled=True, recipe=recipe):
+        y_a = model[0](x_a)
+        y_b = model[1](x_b)
+    loss = y_b.float().sum()
+    if y_a.numel() > 0:
+        loss = loss + y_a.float().sum()
+    loss.backward()
+
+
+def _backward_update_step_no_backward_in_scope(model, recipe):
+    """Odd ranks run no backward at all; the scope still triggers the update."""
+    x = torch.randn(BATCH_SIZE, HIDDEN_SIZE, device="cuda", requires_grad=True)
+    with te.quantization_backward_scope():
+        with te.autocast(enabled=True, recipe=recipe):
+            y = model[1](model[0](x))
+        if WORLD_RANK % 2 == 0:
+            y.float().sum().backward()
+
+
+@run_distributed_test()
+def _test_backward_update_with_skipped_ranks(step_fn):
+    model = nn.ModuleList([te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=True) for _ in range(2)]).cuda()
+    recipe = DelayedScaling(reduce_amax=True)
+    qstate = FP8GlobalStateManager.quantization_state
+    for _ in range(3):
+        model.zero_grad(set_to_none=True)
+        step_fn(model, recipe)
+        assert not qstate.pending_backward_quantization_update
+        assert qstate.backward_quantization_update_callback_task_id is None
+        _assert_bwd_state_matches_across_ranks(model)
+    # Ranks that skipped backward must have received the other ranks' amaxes.
+    for module in model:
+        assert module.fp8_meta["scaling_bwd"].amax_history.abs().sum() > 0
+
+
+def test_backward_update_with_skipped_ranks():
+    """Every rank must join the amax reduction even if it skipped backward."""
+    if QUANTIZATION != "fp8":
+        return
+    for step_fn in (
+        _backward_update_step_skipped_module,
+        _backward_update_step_no_backward_in_scope,
+    ):
+        _test_backward_update_with_skipped_ranks(step_fn)
+
+
 if __name__ == "__main__":
     sys.exit(main())
diff --git a/tests/pytorch/distributed/test_backward_update_ranks.py b/tests/pytorch/distributed/test_backward_update_ranks.py
deleted file mode 100644
index b2987ee7d9..0000000000
--- a/tests/pytorch/distributed/test_backward_update_ranks.py
+++ /dev/null
@@ -1,28 +0,0 @@
-# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-#
-# See LICENSE for license information.
-
-import os
-import subprocess
-from pathlib import Path
-
-import pytest
-import torch
-import transformer_engine.pytorch as te
-
-if torch.cuda.device_count() < 2:
-    pytest.skip("Distributed training needs at least 2 GPUs.", allow_module_level=True)
-
-fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)
-
-TEST_ROOT = Path(__file__).parent.resolve()
-NUM_PROCS: int = min(4, torch.cuda.device_count())
-LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"]
-
-
-@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-@pytest.mark.parametrize("case", ["skipped_module_backward", "no_backward_in_scope"])
-def test_backward_update_ranks(case):
-    test_cmd = LAUNCH_CMD + [str(TEST_ROOT / "run_backward_update_ranks.py"), "--case", case]
-    result = subprocess.run(test_cmd, env=os.environ, check=False, timeout=600)
-    assert result.returncode == 0

From f678207f4aef530f33f3c88e0da9d318420a21ab Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 11:28:40 +0200
Subject: [PATCH 07/13] Keep the recipe out of autograd ctx for the backward
 update request

Decide in forward via FP8GlobalStateManager.backward_quantization_update_needed()
and carry only a bool through ctx; request_backward_quantization_update()
takes no arguments.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/module/grouped_linear.py              | 14 ++++++++++----
 .../pytorch/module/layernorm_linear.py            |  7 +++++--
 .../pytorch/module/layernorm_mlp.py               |  7 +++++--
 transformer_engine/pytorch/module/linear.py       | 12 +++++++-----
 transformer_engine/pytorch/ops/fuser.py           | 10 ++++------
 transformer_engine/pytorch/quantization.py        | 15 ++++++++++-----
 6 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 010d6c661b..f62d500d0f 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -598,6 +598,9 @@ def _forward_grouped_tensor(
             ctx.use_bias = use_bias
             ctx.inp_shape = inp.shape
             ctx.requires_dgrad = inp.requires_grad
+            ctx.request_backward_quantization_update = (
+                ctx.fp8 and FP8GlobalStateManager.backward_quantization_update_needed()
+            )
             ctx.wgrad_store = wgrad_store
             ctx.debug = False
             ctx.save_original_input = save_original_input
@@ -985,6 +988,9 @@ def forward(
                 ctx.grad_input_quantizers = [None] * num_gemms
                 ctx.grad_weight_quantizers = [None] * num_gemms
                 ctx.grad_output_quantizers = [None] * num_gemms
+            ctx.request_backward_quantization_update = (
+                ctx.fp8 and FP8GlobalStateManager.backward_quantization_update_needed()
+            )
 
         # [*, in_features] -> [*, out_features] except first dimension changes for SP
         return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces
@@ -1236,8 +1242,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
         else:
             wgrad_list = [None] * num_weight_args
 
-        if ctx.fp8:
-            FP8GlobalStateManager.request_backward_quantization_update(ctx.fp8_recipe)
+        if ctx.request_backward_quantization_update:
+            FP8GlobalStateManager.request_backward_quantization_update()
         return (
             dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None,
             None,  # m_splits
@@ -1503,8 +1509,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
             ):
                 grad_biases = [None] * ctx.num_gemms
 
-        if ctx.fp8:
-            FP8GlobalStateManager.request_backward_quantization_update(ctx.fp8_recipe)
+        if ctx.request_backward_quantization_update:
+            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 145fe4cbf9..13dc428161 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -599,6 +599,9 @@ def forward(
                 ctx.grad_input_quantizer = None
                 ctx.grad_weight_quantizer = None
                 ctx.grad_output_quantizer = None
+            ctx.request_backward_quantization_update = (
+                ctx.fp8 and FP8GlobalStateManager.backward_quantization_update_needed()
+            )
 
         # ------------------------------------------------------
         # Cached state for backward pass is ready...
@@ -1173,8 +1176,8 @@ def wgrad_gemm(
         else:
             wgrad = None
 
-        if ctx.fp8:
-            FP8GlobalStateManager.request_backward_quantization_update(ctx.fp8_recipe)
+        if ctx.request_backward_quantization_update:
+            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 0b3f7558a5..da9c1d3cc5 100644
--- a/transformer_engine/pytorch/module/layernorm_mlp.py
+++ b/transformer_engine/pytorch/module/layernorm_mlp.py
@@ -906,6 +906,9 @@ def _forward(
             )
             ctx.normalization = normalization
             ctx.wgrad_store = wgrad_store
+            ctx.request_backward_quantization_update = (
+                ctx.fp8 and FP8GlobalStateManager.backward_quantization_update_needed()
+            )
             if is_recomputation:  # return the recomputed tensors
                 return (
                     ctx,
@@ -1793,8 +1796,8 @@ def fc1_wgrad_gemm(
         else:
             fc2_wgrad = None
 
-        if ctx.fp8:
-            FP8GlobalStateManager.request_backward_quantization_update(ctx.fp8_recipe)
+        if ctx.request_backward_quantization_update:
+            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 9631f5e7a2..16d8bb2d3d 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -197,7 +197,7 @@ class LinearBwdArgs:
     # --- Numerical / dtype config ---
     activation_dtype: Optional[torch.dtype] = None
     fp8: bool = False
-    fp8_recipe: Optional[Recipe] = None
+    request_backward_quantization_update: bool = False
     dgrad_use_split_accumulator: bool = _2X_ACC_DGRAD
     wgrad_use_split_accumulator: bool = _2X_ACC_WGRAD
     backward_override: Optional[str] = None
@@ -702,7 +702,6 @@ def _linear_setup_ctx(
     # Numerical / dtype config
     bwd_args.activation_dtype = fwd_args.activation_dtype
     bwd_args.fp8 = fp8
-    bwd_args.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
     bwd_args.dgrad_use_split_accumulator = fwd_args.dgrad_use_split_accumulator
     bwd_args.wgrad_use_split_accumulator = fwd_args.wgrad_use_split_accumulator
     bwd_args.backward_override = backward_override
@@ -757,6 +756,9 @@ def _linear_setup_ctx(
         bwd_args.grad_input_quantizer = None
         bwd_args.grad_weight_quantizer = None
         bwd_args.grad_output_quantizer = None
+    bwd_args.request_backward_quantization_update = (
+        bwd_args.fp8 and FP8GlobalStateManager.backward_quantization_update_needed()
+    )
 
     saved_inputmat, wt_save, saved_weight, saved_bias = tensors_to_save_from_forward
     inputmat_alias, wt_save_alias, saved_weight_alias, bias_alias = ctx_attrs[
@@ -1439,13 +1441,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
-        fp8_recipe = bwd_args.fp8_recipe if bwd_args.fp8 else None
+        request_update = bwd_args.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 fp8_recipe is not None:
-            FP8GlobalStateManager.request_backward_quantization_update(fp8_recipe)
+        if request_update:
+            FP8GlobalStateManager.request_backward_quantization_update()
         return result
 
 
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 9d8ab42662..7dac79810c 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -220,10 +220,8 @@ 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.fp8_recipe = (
-                FP8GlobalStateManager.get_fp8_recipe()
-                if FP8GlobalStateManager.is_fp8_enabled()
-                else None
+            func_ctx.request_backward_quantization_update = (
+                FP8GlobalStateManager.backward_quantization_update_needed()
             )
 
         # Mark output tensors as not deletable in backward
@@ -364,8 +362,8 @@ def backward(
             for op_idx, input_idx in func_ctx.external_extra_input_slots
         ]
 
-        if func_ctx.fp8_recipe is not None:
-            FP8GlobalStateManager.request_backward_quantization_update(func_ctx.fp8_recipe)
+        if func_ctx.request_backward_quantization_update:
+            FP8GlobalStateManager.request_backward_quantization_update()
 
         return (
             dx,  # input_
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 310a5eed46..b5554a4412 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -688,14 +688,19 @@ def reduce_and_update_quantization_state(
     reduce_and_update_fp8_tensors = reduce_and_update_quantization_state
 
     @classmethod
-    def request_backward_quantization_update(cls, recipe: Recipe) -> None:
+    def backward_quantization_update_needed(cls) -> bool:
+        """Whether modules in the active autocast must request an update after backward."""
+        if not cls.is_fp8_enabled():
+            return False
+        recipe = cls.get_fp8_recipe()
+        return recipe.delayed() or recipe.custom()
+
+    @classmethod
+    def request_backward_quantization_update(cls) -> None:
         """Request an update after the enclosing logical backward.
 
-        No-op for recipes without delayed-scaling state and inside CUDA graph
-        capture, where the graphed wrapper performs the update.
+        No-op inside CUDA graph capture, where the graphed wrapper performs the update.
         """
-        if not (recipe.delayed() or recipe.custom()):
-            return
         from .graph import is_graph_capturing  # pylint: disable=import-outside-toplevel
 
         if is_graph_capturing():

From 576545b4e03ab90c8d48f84c2f40e5799d3318a7 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 11:49:42 +0200
Subject: [PATCH 08/13] test: assert backward update runs after all modules
 finished backward

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_recipe.py | 30 +++++++++++++++++++++++-------
 1 file changed, 23 insertions(+), 7 deletions(-)

diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py
index a5d5fd6df8..b4aed54ccd 100644
--- a/tests/pytorch/test_recipe.py
+++ b/tests/pytorch/test_recipe.py
@@ -789,9 +789,19 @@ def test_stateful_unknown_or_malformed_pickled_extra_state_requires_opt_in(paylo
 
 
 class _UpdateCounter:
-    def __init__(self):
+    """Count backward updates and check they run after every given module's backward.
+
+    A module that has run backward holds a nonzero grad amax in the current
+    history slot; the update rolls the history and clears that slot.
+    """
+
+    def __init__(self, *modules):
         self.backward = 0
         self._original = None
+        self._modules = [m for root in modules for m in root.modules() if hasattr(m, "fp8_meta")]
+
+    def _current_bwd_amax(self):
+        return [m.fp8_meta["scaling_bwd"].amax_history[0] for m in self._modules]
 
     def __enter__(self):
         self._original = FP8GlobalStateManager.reduce_and_update_quantization_state.__func__
@@ -801,7 +811,13 @@ def __enter__(self):
         def counted(cls, forward=True):
             if not forward:
                 counter.backward += 1
-            return original(cls, forward=forward)
+                for amax in counter._current_bwd_amax():
+                    assert amax.any(), "backward update ran before all modules finished backward"
+            result = original(cls, forward=forward)
+            if not forward:
+                for amax in counter._current_bwd_amax():
+                    assert not amax.any(), "backward update did not roll the amax history"
+            return result
 
         FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(counted)
         return self
@@ -877,7 +893,7 @@ def test_delayed_scaling_updates_once_per_backward(mode):
     model = _make_update_test_model()
     recipe = DelayedScaling()
 
-    with _UpdateCounter() as counter:
+    with _UpdateCounter(model) as counter:
         for step in range(_UPDATE_TEST_STEPS):
             x = torch.randn(
                 _UPDATE_TEST_BATCH,
@@ -907,7 +923,7 @@ def test_quantization_backward_scope_groups_independent_graphs():
     ]
     recipe = DelayedScaling()
 
-    with _UpdateCounter() as counter:
+    with _UpdateCounter(*models) as counter:
         with te.autocast(enabled=True, recipe=recipe):
             outputs = [_run_update_test_layers(model, x) for model, x in zip(models, inputs)]
         with te.quantization_backward_scope():
@@ -930,7 +946,7 @@ def test_quantization_backward_scope_covers_delayed_wgrad():
     ).cuda()
     recipe = DelayedScaling()
 
-    with _UpdateCounter() as counter, te.quantization_backward_scope():
+    with _UpdateCounter(model) as counter, te.quantization_backward_scope():
         x = torch.randn(
             _UPDATE_TEST_BATCH,
             _UPDATE_TEST_HIDDEN,
@@ -955,7 +971,7 @@ def test_delayed_scaling_update_on_branched_graph(checkpoint_first_branch):
     branch_b = _make_update_test_model(num_layers=2, seed=2)
     recipe = DelayedScaling()
 
-    with _UpdateCounter() as counter:
+    with _UpdateCounter(branch_a, branch_b) as counter:
         x = torch.randn(
             _UPDATE_TEST_BATCH,
             _UPDATE_TEST_HIDDEN,
@@ -980,7 +996,7 @@ def test_unused_checkpoint_branch_does_not_own_backward_update():
     unused = _make_update_test_model(num_layers=2, seed=2)
     recipe = DelayedScaling()
 
-    with _UpdateCounter() as counter:
+    with _UpdateCounter(used) as counter:
         x = torch.randn(
             _UPDATE_TEST_BATCH,
             _UPDATE_TEST_HIDDEN,

From 524070567ffbd4dfa21dd34f073de42ae16dccf8 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 13:53:14 +0200
Subject: [PATCH 09/13] test: merge backward update tests into one parametrized
 case; reset global state in distributed test

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/distributed/run_numerics.py |   3 +
 tests/pytorch/test_recipe.py              | 211 ++++++++++------------
 2 files changed, 101 insertions(+), 113 deletions(-)

diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py
index 51041b7b5d..5d8529e909 100644
--- a/tests/pytorch/distributed/run_numerics.py
+++ b/tests/pytorch/distributed/run_numerics.py
@@ -1182,6 +1182,8 @@ def _backward_update_step_no_backward_in_scope(model, recipe):
 
 @run_distributed_test()
 def _test_backward_update_with_skipped_ranks(step_fn):
+    # Drop amax buffers registered by earlier tests so only this model is reduced.
+    FP8GlobalStateManager.reset()
     model = nn.ModuleList([te.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=True) for _ in range(2)]).cuda()
     recipe = DelayedScaling(reduce_amax=True)
     qstate = FP8GlobalStateManager.quantization_state
@@ -1205,6 +1207,7 @@ def test_backward_update_with_skipped_ranks():
         _backward_update_step_no_backward_in_scope,
     ):
         _test_backward_update_with_skipped_ranks(step_fn)
+    FP8GlobalStateManager.reset()
 
 
 if __name__ == "__main__":
diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py
index b4aed54ccd..76edbe6e73 100644
--- a/tests/pytorch/test_recipe.py
+++ b/tests/pytorch/test_recipe.py
@@ -842,13 +842,6 @@ def _run_update_test_layers(layers, x):
     return x
 
 
-def _run_update_test_step(model, x, forward_fn, recipe):
-    with te.autocast(enabled=True, recipe=recipe):
-        out = forward_fn(model, x)
-    loss = out.float().sum()
-    loss.backward()
-
-
 def _update_forward_plain(model, x):
     return _run_update_test_layers(model, x)
 
@@ -877,135 +870,127 @@ def outer(value):
     return te_checkpoint(outer, x, use_reentrant=True)
 
 
-_UPDATE_FORWARD_FNS = {
-    "plain": _update_forward_plain,
-    "reentrant": _update_forward_reentrant,
-    "non_reentrant": _update_forward_non_reentrant,
-    "per_layer_reentrant": _update_forward_per_layer_reentrant,
-    "nested": _update_forward_nested,
-}
+# Each case builds its modules and returns (step, tracked_modules). ``step``
+# runs one forward + backward; ``tracked_modules`` must all finish backward
+# before the single update of that step.
 
 
-@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-@pytest.mark.parametrize("mode", _UPDATE_FORWARD_FNS.keys())
-def test_delayed_scaling_updates_once_per_backward(mode):
-    FP8GlobalStateManager.reset()
-    model = _make_update_test_model()
-    recipe = DelayedScaling()
+def _update_case_single_graph(forward_fn):
+    def build():
+        model = _make_update_test_model()
 
-    with _UpdateCounter(model) as counter:
-        for step in range(_UPDATE_TEST_STEPS):
-            x = torch.randn(
-                _UPDATE_TEST_BATCH,
-                _UPDATE_TEST_HIDDEN,
-                device="cuda",
-                requires_grad=True,
-            )
-            _run_update_test_step(model, x, _UPDATE_FORWARD_FNS[mode], recipe)
-            assert counter.backward == step + 1
-            qstate = FP8GlobalStateManager.quantization_state
-            assert not qstate.pending_backward_quantization_update
-            assert qstate.backward_quantization_update_callback_task_id is None
+        def step(x, recipe, counter):
+            with te.autocast(enabled=True, recipe=recipe):
+                out = forward_fn(model, x)
+            out.float().sum().backward()
 
+        return step, [model]
 
-@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-def test_quantization_backward_scope_groups_independent_graphs():
-    FP8GlobalStateManager.reset()
+    return build
+
+
+def _update_case_branched(checkpoint_first_branch):
+    def build():
+        branch_a = _make_update_test_model(num_layers=2, seed=1)
+        branch_b = _make_update_test_model(num_layers=2, seed=2)
+
+        def step(x, recipe, counter):
+            with te.autocast(enabled=True, recipe=recipe):
+                if checkpoint_first_branch:
+                    out_a = te_checkpoint(_run_update_test_layers, branch_a, x, use_reentrant=True)
+                else:
+                    out_a = _run_update_test_layers(branch_a, x)
+                out_b = te_checkpoint(_run_update_test_layers, branch_b, x, use_reentrant=True)
+            (out_a + out_b).float().sum().backward()
+
+        return step, [branch_a, branch_b]
+
+    return build
+
+
+def _update_case_unused_checkpoint_branch():
+    used = _make_update_test_model(num_layers=2, seed=1)
+    unused = _make_update_test_model(num_layers=2, seed=2)
+
+    def step(x, recipe, counter):
+        with te.autocast(enabled=True, recipe=recipe):
+            unused_out = te_checkpoint(_run_update_test_layers, unused, x, use_reentrant=True)
+            out = _run_update_test_layers(used, x)
+        out.float().sum().backward()
+        del unused_out
+
+    return step, [used]
+
+
+def _update_case_scope_independent_graphs():
     models = [_make_update_test_model(num_layers=2, seed=seed) for seed in (1, 2)]
-    inputs = [
-        torch.randn(
-            _UPDATE_TEST_BATCH,
-            _UPDATE_TEST_HIDDEN,
-            device="cuda",
-            requires_grad=True,
-        )
-        for _ in models
-    ]
-    recipe = DelayedScaling()
 
-    with _UpdateCounter(*models) as counter:
+    def step(x, recipe, counter):
         with te.autocast(enabled=True, recipe=recipe):
-            outputs = [_run_update_test_layers(model, x) for model, x in zip(models, inputs)]
+            outputs = [_run_update_test_layers(model, x) for model in models]
+        updates_before = counter.backward
         with te.quantization_backward_scope():
             for output in outputs:
                 output.float().sum().backward()
-                assert counter.backward == 0
-        assert counter.backward == 1
-        for x in inputs:
-            assert x.grad is not None
+                assert counter.backward == updates_before
 
+    return step, models
 
-@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-def test_quantization_backward_scope_covers_delayed_wgrad():
-    FP8GlobalStateManager.reset()
+
+def _update_case_scope_delayed_wgrad():
     model = te.Linear(
         _UPDATE_TEST_HIDDEN,
         _UPDATE_TEST_HIDDEN,
         bias=True,
         delay_wgrad_compute=True,
     ).cuda()
-    recipe = DelayedScaling()
-
-    with _UpdateCounter(model) as counter, te.quantization_backward_scope():
-        x = torch.randn(
-            _UPDATE_TEST_BATCH,
-            _UPDATE_TEST_HIDDEN,
-            device="cuda",
-            requires_grad=True,
-        )
-        with te.autocast(enabled=True, recipe=recipe):
-            out = model(x)
-        out.float().sum().backward()
-        assert counter.backward == 0
-        model.backward_dw()
-        assert model.weight.grad is not None
-        assert counter.backward == 0
-    assert counter.backward == 1
 
-
-@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-@pytest.mark.parametrize("checkpoint_first_branch", [True, False])
-def test_delayed_scaling_update_on_branched_graph(checkpoint_first_branch):
-    FP8GlobalStateManager.reset()
-    branch_a = _make_update_test_model(num_layers=2, seed=1)
-    branch_b = _make_update_test_model(num_layers=2, seed=2)
-    recipe = DelayedScaling()
-
-    with _UpdateCounter(branch_a, branch_b) as counter:
-        x = torch.randn(
-            _UPDATE_TEST_BATCH,
-            _UPDATE_TEST_HIDDEN,
-            device="cuda",
-            requires_grad=True,
-        )
-        with te.autocast(enabled=True, recipe=recipe):
-            if checkpoint_first_branch:
-                out_a = te_checkpoint(_run_update_test_layers, branch_a, x, use_reentrant=True)
-            else:
-                out_a = _run_update_test_layers(branch_a, x)
-            out_b = te_checkpoint(_run_update_test_layers, branch_b, x, use_reentrant=True)
-        (out_a + out_b).float().sum().backward()
-        assert counter.backward == 1
-        assert x.grad is not None and torch.isfinite(x.grad).all()
+    def step(x, recipe, counter):
+        updates_before = counter.backward
+        with te.quantization_backward_scope():
+            with te.autocast(enabled=True, recipe=recipe):
+                out = model(x)
+            out.float().sum().backward()
+            assert counter.backward == updates_before
+            model.backward_dw()
+            assert model.weight.grad is not None
+            assert counter.backward == updates_before
+
+    return step, [model]
+
+
+_UPDATE_TEST_CASES = {
+    "plain": _update_case_single_graph(_update_forward_plain),
+    "reentrant": _update_case_single_graph(_update_forward_reentrant),
+    "non_reentrant": _update_case_single_graph(_update_forward_non_reentrant),
+    "per_layer_reentrant": _update_case_single_graph(_update_forward_per_layer_reentrant),
+    "nested": _update_case_single_graph(_update_forward_nested),
+    "branched": _update_case_branched(checkpoint_first_branch=False),
+    "branched_checkpoint_first": _update_case_branched(checkpoint_first_branch=True),
+    "unused_checkpoint_branch": _update_case_unused_checkpoint_branch,
+    "scope_independent_graphs": _update_case_scope_independent_graphs,
+    "scope_delayed_wgrad": _update_case_scope_delayed_wgrad,
+}
 
 
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-def test_unused_checkpoint_branch_does_not_own_backward_update():
+@pytest.mark.parametrize("case", _UPDATE_TEST_CASES.keys())
+def test_delayed_scaling_updates_once_per_backward(case):
     FP8GlobalStateManager.reset()
-    used = _make_update_test_model(num_layers=2, seed=1)
-    unused = _make_update_test_model(num_layers=2, seed=2)
+    step, tracked = _UPDATE_TEST_CASES[case]()
     recipe = DelayedScaling()
 
-    with _UpdateCounter(used) as counter:
-        x = torch.randn(
-            _UPDATE_TEST_BATCH,
-            _UPDATE_TEST_HIDDEN,
-            device="cuda",
-            requires_grad=True,
-        )
-        with te.autocast(enabled=True, recipe=recipe):
-            unused_out = te_checkpoint(_run_update_test_layers, unused, x, use_reentrant=True)
-            out = _run_update_test_layers(used, x)
-        out.float().sum().backward()
-        del unused_out
-        assert counter.backward == 1
+    with _UpdateCounter(*tracked) as counter:
+        for i in range(_UPDATE_TEST_STEPS):
+            x = torch.randn(
+                _UPDATE_TEST_BATCH,
+                _UPDATE_TEST_HIDDEN,
+                device="cuda",
+                requires_grad=True,
+            )
+            step(x, recipe, counter)
+            assert counter.backward == i + 1
+            assert x.grad is not None and torch.isfinite(x.grad).all()
+            qstate = FP8GlobalStateManager.quantization_state
+            assert not qstate.pending_backward_quantization_update
+            assert qstate.backward_quantization_update_callback_task_id is None

From b161f624263525c3e845df1f2815ff49e81b1571 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 14:01:53 +0200
Subject: [PATCH 10/13] test: move backward update tests next to the FP8
 checkpoint tests in test_numerics.py

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_numerics.py | 214 +++++++++++++++++++++++++++++++++
 tests/pytorch/test_recipe.py   | 214 ---------------------------------
 2 files changed, 214 insertions(+), 214 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index e6a83d92bc..1a0bc6f7d9 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -13,6 +13,7 @@
 
 from transformer_engine.pytorch.quantization import (
     FP8GlobalStateManager,
+    quantization_backward_scope,
 )
 from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV
 from transformer_engine.pytorch.utils import (
@@ -955,6 +956,219 @@ def body(value):
     assert _FP8_RECOMPUTE_KEY in fp8_layer.fp8_meta
 
 
+_UPDATE_TEST_HIDDEN = 128
+_UPDATE_TEST_BATCH = 32
+_UPDATE_TEST_STEPS = 3
+
+
+class _UpdateCounter:
+    """Count backward updates and check they run after every given module's backward.
+
+    A module that has run backward holds a nonzero grad amax in the current
+    history slot; the update rolls the history and clears that slot.
+    """
+
+    def __init__(self, *modules):
+        self.backward = 0
+        self._original = None
+        self._modules = [m for root in modules for m in root.modules() if hasattr(m, "fp8_meta")]
+
+    def _current_bwd_amax(self):
+        return [m.fp8_meta["scaling_bwd"].amax_history[0] for m in self._modules]
+
+    def __enter__(self):
+        self._original = FP8GlobalStateManager.reduce_and_update_quantization_state.__func__
+        original = self._original
+        counter = self
+
+        def counted(cls, forward=True):
+            if not forward:
+                counter.backward += 1
+                for amax in counter._current_bwd_amax():
+                    assert amax.any(), "backward update ran before all modules finished backward"
+            result = original(cls, forward=forward)
+            if not forward:
+                for amax in counter._current_bwd_amax():
+                    assert not amax.any(), "backward update did not roll the amax history"
+            return result
+
+        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_update_test_model(num_layers=3, seed=1234):
+    torch.manual_seed(seed)
+    return torch.nn.ModuleList(
+        [
+            Linear(_UPDATE_TEST_HIDDEN, _UPDATE_TEST_HIDDEN, bias=True).cuda()
+            for _ in range(num_layers)
+        ]
+    )
+
+
+def _run_update_test_layers(layers, x):
+    for layer in layers:
+        x = layer(x)
+    return x
+
+
+def _update_forward_plain(model, x):
+    return _run_update_test_layers(model, x)
+
+
+def _update_forward_reentrant(model, x):
+    return te_checkpoint(_run_update_test_layers, model, x, use_reentrant=True)
+
+
+def _update_forward_non_reentrant(model, x):
+    return te_checkpoint(_run_update_test_layers, model, x, use_reentrant=False)
+
+
+def _update_forward_per_layer_reentrant(model, x):
+    for layer in model:
+        x = te_checkpoint(layer, x, use_reentrant=True)
+    return x
+
+
+def _update_forward_nested(model, x):
+    def inner(value):
+        return te_checkpoint(model[1], value, use_reentrant=True)
+
+    def outer(value):
+        return model[2](inner(model[0](value)))
+
+    return te_checkpoint(outer, x, use_reentrant=True)
+
+
+# Each case builds its modules and returns (step, tracked_modules). ``step``
+# runs one forward + backward; ``tracked_modules`` must all finish backward
+# before the single update of that step.
+
+
+def _update_case_single_graph(forward_fn):
+    def build():
+        model = _make_update_test_model()
+
+        def step(x, recipe, counter):
+            with autocast(enabled=True, recipe=recipe):
+                out = forward_fn(model, x)
+            out.float().sum().backward()
+
+        return step, [model]
+
+    return build
+
+
+def _update_case_branched(checkpoint_first_branch):
+    def build():
+        branch_a = _make_update_test_model(num_layers=2, seed=1)
+        branch_b = _make_update_test_model(num_layers=2, seed=2)
+
+        def step(x, recipe, counter):
+            with autocast(enabled=True, recipe=recipe):
+                if checkpoint_first_branch:
+                    out_a = te_checkpoint(_run_update_test_layers, branch_a, x, use_reentrant=True)
+                else:
+                    out_a = _run_update_test_layers(branch_a, x)
+                out_b = te_checkpoint(_run_update_test_layers, branch_b, x, use_reentrant=True)
+            (out_a + out_b).float().sum().backward()
+
+        return step, [branch_a, branch_b]
+
+    return build
+
+
+def _update_case_unused_checkpoint_branch():
+    used = _make_update_test_model(num_layers=2, seed=1)
+    unused = _make_update_test_model(num_layers=2, seed=2)
+
+    def step(x, recipe, counter):
+        with autocast(enabled=True, recipe=recipe):
+            unused_out = te_checkpoint(_run_update_test_layers, unused, x, use_reentrant=True)
+            out = _run_update_test_layers(used, x)
+        out.float().sum().backward()
+        del unused_out
+
+    return step, [used]
+
+
+def _update_case_scope_independent_graphs():
+    models = [_make_update_test_model(num_layers=2, seed=seed) for seed in (1, 2)]
+
+    def step(x, recipe, counter):
+        with autocast(enabled=True, recipe=recipe):
+            outputs = [_run_update_test_layers(model, x) for model in models]
+        updates_before = counter.backward
+        with quantization_backward_scope():
+            for output in outputs:
+                output.float().sum().backward()
+                assert counter.backward == updates_before
+
+    return step, models
+
+
+def _update_case_scope_delayed_wgrad():
+    model = Linear(
+        _UPDATE_TEST_HIDDEN,
+        _UPDATE_TEST_HIDDEN,
+        bias=True,
+        delay_wgrad_compute=True,
+    ).cuda()
+
+    def step(x, recipe, counter):
+        updates_before = counter.backward
+        with quantization_backward_scope():
+            with autocast(enabled=True, recipe=recipe):
+                out = model(x)
+            out.float().sum().backward()
+            assert counter.backward == updates_before
+            model.backward_dw()
+            assert model.weight.grad is not None
+            assert counter.backward == updates_before
+
+    return step, [model]
+
+
+_UPDATE_TEST_CASES = {
+    "plain": _update_case_single_graph(_update_forward_plain),
+    "reentrant": _update_case_single_graph(_update_forward_reentrant),
+    "non_reentrant": _update_case_single_graph(_update_forward_non_reentrant),
+    "per_layer_reentrant": _update_case_single_graph(_update_forward_per_layer_reentrant),
+    "nested": _update_case_single_graph(_update_forward_nested),
+    "branched": _update_case_branched(checkpoint_first_branch=False),
+    "branched_checkpoint_first": _update_case_branched(checkpoint_first_branch=True),
+    "unused_checkpoint_branch": _update_case_unused_checkpoint_branch,
+    "scope_independent_graphs": _update_case_scope_independent_graphs,
+    "scope_delayed_wgrad": _update_case_scope_delayed_wgrad,
+}
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("case", _UPDATE_TEST_CASES.keys())
+def test_delayed_scaling_updates_once_per_backward(case):
+    FP8GlobalStateManager.reset()
+    step, tracked = _UPDATE_TEST_CASES[case]()
+    fp8_recipe = recipe.DelayedScaling()
+
+    with _UpdateCounter(*tracked) as counter:
+        for i in range(_UPDATE_TEST_STEPS):
+            x = torch.randn(
+                _UPDATE_TEST_BATCH,
+                _UPDATE_TEST_HIDDEN,
+                device="cuda",
+                requires_grad=True,
+            )
+            step(x, fp8_recipe, counter)
+            assert counter.backward == i + 1
+            assert x.grad is not None and torch.isfinite(x.grad).all()
+            qstate = FP8GlobalStateManager.quantization_state
+            assert not qstate.pending_backward_quantization_update
+            assert qstate.backward_quantization_update_callback_task_id is None
+
+
 def _test_e2e_checkpointing_get_model(config, dtype):
     sigma = 0.023
     init_method = init_method_normal(sigma)
diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py
index 76edbe6e73..ccef104a33 100644
--- a/tests/pytorch/test_recipe.py
+++ b/tests/pytorch/test_recipe.py
@@ -32,7 +32,6 @@
     _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,
@@ -781,216 +780,3 @@ 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")
-
-
-_UPDATE_TEST_HIDDEN = 128
-_UPDATE_TEST_BATCH = 32
-_UPDATE_TEST_STEPS = 3
-
-
-class _UpdateCounter:
-    """Count backward updates and check they run after every given module's backward.
-
-    A module that has run backward holds a nonzero grad amax in the current
-    history slot; the update rolls the history and clears that slot.
-    """
-
-    def __init__(self, *modules):
-        self.backward = 0
-        self._original = None
-        self._modules = [m for root in modules for m in root.modules() if hasattr(m, "fp8_meta")]
-
-    def _current_bwd_amax(self):
-        return [m.fp8_meta["scaling_bwd"].amax_history[0] for m in self._modules]
-
-    def __enter__(self):
-        self._original = FP8GlobalStateManager.reduce_and_update_quantization_state.__func__
-        original = self._original
-        counter = self
-
-        def counted(cls, forward=True):
-            if not forward:
-                counter.backward += 1
-                for amax in counter._current_bwd_amax():
-                    assert amax.any(), "backward update ran before all modules finished backward"
-            result = original(cls, forward=forward)
-            if not forward:
-                for amax in counter._current_bwd_amax():
-                    assert not amax.any(), "backward update did not roll the amax history"
-            return result
-
-        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_update_test_model(num_layers=3, seed=1234):
-    torch.manual_seed(seed)
-    return torch.nn.ModuleList(
-        [
-            te.Linear(_UPDATE_TEST_HIDDEN, _UPDATE_TEST_HIDDEN, bias=True).cuda()
-            for _ in range(num_layers)
-        ]
-    )
-
-
-def _run_update_test_layers(layers, x):
-    for layer in layers:
-        x = layer(x)
-    return x
-
-
-def _update_forward_plain(model, x):
-    return _run_update_test_layers(model, x)
-
-
-def _update_forward_reentrant(model, x):
-    return te_checkpoint(_run_update_test_layers, model, x, use_reentrant=True)
-
-
-def _update_forward_non_reentrant(model, x):
-    return te_checkpoint(_run_update_test_layers, model, x, use_reentrant=False)
-
-
-def _update_forward_per_layer_reentrant(model, x):
-    for layer in model:
-        x = te_checkpoint(layer, x, use_reentrant=True)
-    return x
-
-
-def _update_forward_nested(model, x):
-    def inner(value):
-        return te_checkpoint(model[1], value, use_reentrant=True)
-
-    def outer(value):
-        return model[2](inner(model[0](value)))
-
-    return te_checkpoint(outer, x, use_reentrant=True)
-
-
-# Each case builds its modules and returns (step, tracked_modules). ``step``
-# runs one forward + backward; ``tracked_modules`` must all finish backward
-# before the single update of that step.
-
-
-def _update_case_single_graph(forward_fn):
-    def build():
-        model = _make_update_test_model()
-
-        def step(x, recipe, counter):
-            with te.autocast(enabled=True, recipe=recipe):
-                out = forward_fn(model, x)
-            out.float().sum().backward()
-
-        return step, [model]
-
-    return build
-
-
-def _update_case_branched(checkpoint_first_branch):
-    def build():
-        branch_a = _make_update_test_model(num_layers=2, seed=1)
-        branch_b = _make_update_test_model(num_layers=2, seed=2)
-
-        def step(x, recipe, counter):
-            with te.autocast(enabled=True, recipe=recipe):
-                if checkpoint_first_branch:
-                    out_a = te_checkpoint(_run_update_test_layers, branch_a, x, use_reentrant=True)
-                else:
-                    out_a = _run_update_test_layers(branch_a, x)
-                out_b = te_checkpoint(_run_update_test_layers, branch_b, x, use_reentrant=True)
-            (out_a + out_b).float().sum().backward()
-
-        return step, [branch_a, branch_b]
-
-    return build
-
-
-def _update_case_unused_checkpoint_branch():
-    used = _make_update_test_model(num_layers=2, seed=1)
-    unused = _make_update_test_model(num_layers=2, seed=2)
-
-    def step(x, recipe, counter):
-        with te.autocast(enabled=True, recipe=recipe):
-            unused_out = te_checkpoint(_run_update_test_layers, unused, x, use_reentrant=True)
-            out = _run_update_test_layers(used, x)
-        out.float().sum().backward()
-        del unused_out
-
-    return step, [used]
-
-
-def _update_case_scope_independent_graphs():
-    models = [_make_update_test_model(num_layers=2, seed=seed) for seed in (1, 2)]
-
-    def step(x, recipe, counter):
-        with te.autocast(enabled=True, recipe=recipe):
-            outputs = [_run_update_test_layers(model, x) for model in models]
-        updates_before = counter.backward
-        with te.quantization_backward_scope():
-            for output in outputs:
-                output.float().sum().backward()
-                assert counter.backward == updates_before
-
-    return step, models
-
-
-def _update_case_scope_delayed_wgrad():
-    model = te.Linear(
-        _UPDATE_TEST_HIDDEN,
-        _UPDATE_TEST_HIDDEN,
-        bias=True,
-        delay_wgrad_compute=True,
-    ).cuda()
-
-    def step(x, recipe, counter):
-        updates_before = counter.backward
-        with te.quantization_backward_scope():
-            with te.autocast(enabled=True, recipe=recipe):
-                out = model(x)
-            out.float().sum().backward()
-            assert counter.backward == updates_before
-            model.backward_dw()
-            assert model.weight.grad is not None
-            assert counter.backward == updates_before
-
-    return step, [model]
-
-
-_UPDATE_TEST_CASES = {
-    "plain": _update_case_single_graph(_update_forward_plain),
-    "reentrant": _update_case_single_graph(_update_forward_reentrant),
-    "non_reentrant": _update_case_single_graph(_update_forward_non_reentrant),
-    "per_layer_reentrant": _update_case_single_graph(_update_forward_per_layer_reentrant),
-    "nested": _update_case_single_graph(_update_forward_nested),
-    "branched": _update_case_branched(checkpoint_first_branch=False),
-    "branched_checkpoint_first": _update_case_branched(checkpoint_first_branch=True),
-    "unused_checkpoint_branch": _update_case_unused_checkpoint_branch,
-    "scope_independent_graphs": _update_case_scope_independent_graphs,
-    "scope_delayed_wgrad": _update_case_scope_delayed_wgrad,
-}
-
-
-@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-@pytest.mark.parametrize("case", _UPDATE_TEST_CASES.keys())
-def test_delayed_scaling_updates_once_per_backward(case):
-    FP8GlobalStateManager.reset()
-    step, tracked = _UPDATE_TEST_CASES[case]()
-    recipe = DelayedScaling()
-
-    with _UpdateCounter(*tracked) as counter:
-        for i in range(_UPDATE_TEST_STEPS):
-            x = torch.randn(
-                _UPDATE_TEST_BATCH,
-                _UPDATE_TEST_HIDDEN,
-                device="cuda",
-                requires_grad=True,
-            )
-            step(x, recipe, counter)
-            assert counter.backward == i + 1
-            assert x.grad is not None and torch.isfinite(x.grad).all()
-            qstate = FP8GlobalStateManager.quantization_state
-            assert not qstate.pending_backward_quantization_update
-            assert qstate.backward_quantization_update_callback_task_id is None

From 1988d058badd47332232e0473af0162cacd69934 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 14:20:33 +0200
Subject: [PATCH 11/13] Rewrite quantization_backward_scope docstring

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/quantization.py | 33 ++++++++++++++++------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index b5554a4412..023ff4b48f 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -896,17 +896,32 @@ def restore_fp8_meta_tensors(fp8_meta: Dict[str, Any]) -> None:
 
 @contextmanager
 def quantization_backward_scope() -> None:
-    """Delay the quantization state update until the end of a logical backward.
+    """Run the delayed-scaling update once, at the end of a logical backward.
 
-    Ordinary backward calls update automatically and do not require this scope.
-    Use it when a logical backward spans multiple autograd calls, including
-    independent graphs produced under one autocast, or includes delayed work
-    such as ``module.backward_dw()``. Nested scopes update once when the
-    outermost scope exits.
+    With delayed scaling, each backward pass ends by reducing the gradient amaxes
+    across the amax reduction group and recomputing the scales. Ordinary
+    ``loss.backward()`` calls do this automatically and do not need this scope.
 
-    The update runs on scope exit even if no quantized module ran backward
-    inside it, so ranks that skipped every backward (e.g. an expert with no
-    tokens) still participate in the amax reduction.
+    Use it when one training step spans several autograd calls, so that the update
+    runs once when the outermost scope exits instead of after every call:
+
+    .. code-block:: python
+
+        with te.quantization_backward_scope():
+            for loss in microbatch_losses:  # e.g. a 1F1B pipeline schedule
+                loss.backward()
+            model.backward_dw()  # deferred weight gradients
+
+    Nested scopes are no-ops; only the outermost one triggers the update. A scope
+    entered inside a running backward (e.g. in a hook) defers the update to the
+    end of that backward.
+
+    The update also runs if no quantized module ran backward inside the scope, so
+    a rank that skipped all of them (e.g. an expert that received no tokens) still
+    joins the amax reduction and does not stall the other ranks. Every rank must
+    therefore enter and exit the scope the same number of times per step.
+
+    For recipes without delayed-scaling state this scope has no effect.
     """
     qstate = FP8GlobalStateManager.quantization_state
     outermost = qstate.quantization_backward_scope_depth == 0

From 7aed56fcbab9601b9a49db0a43cbe0ef23f7972e Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 14:24:56 +0200
Subject: [PATCH 12/13] docs: describe quantization_backward_scope in the
 delayed scaling feature page

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../fp8_delayed_scaling/fp8_delayed_scaling.rst     | 13 +++++++++++++
 .../pytorch_delayed_scaling_distributed_example.py  |  7 +++++++
 2 files changed, 20 insertions(+)

diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
index 9d05305eda..74a1e90dcd 100644
--- a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
+++ b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
@@ -146,6 +146,19 @@ However, amax reduction works slightly differently in different frameworks.
         skipped - if no rank executes a module, its history is not rotated and scale 
         remains unchanged.
 
+      The gradient amaxes are reduced separately, once at the end of every ``backward()``
+      call, on each rank where at least one quantized module ran backward. When a
+      training step consists of several ``backward()`` calls (e.g. a 1F1B pipeline
+      schedule) or ends with deferred weight gradients (``backward_dw()``), wrap them in
+      ``quantization_backward_scope`` so the update runs once per step. The scope also
+      runs the update on ranks where no quantized module ran backward, so, like
+      ``autocast``, it must be entered and exited on all ranks:
+
+      .. literalinclude:: pytorch_delayed_scaling_distributed_example.py
+         :language: python
+         :start-after: # START_BACKWARD_SCOPE_EXAMPLE
+         :end-before: # END_BACKWARD_SCOPE_EXAMPLE
+
 
    .. tab:: JAX
 
diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py b/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py
index 863b71e8c6..afb21c200d 100644
--- a/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py
+++ b/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py
@@ -16,3 +16,10 @@
     output = model(inp)
 
 # END_AMAX_REDUCTION_EXAMPLE
+
+# START_BACKWARD_SCOPE_EXAMPLE
+with te.quantization_backward_scope():
+    for loss in microbatch_losses:
+        loss.backward()
+    model.backward_dw()  # only with delay_wgrad_compute=True
+# END_BACKWARD_SCOPE_EXAMPLE

From 5f795b1dad77623dc445f207c43fee0868450488 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 2 Sep 2026 14:28:35 +0200
Subject: [PATCH 13/13] docs: drop backward_dw from the
 quantization_backward_scope examples

Gradient quantization happens in backward(); backward_dw() only runs the
stored GEMM, so it does not motivate the scope.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../fp8_delayed_scaling/fp8_delayed_scaling.rst              | 4 ++--
 .../pytorch_delayed_scaling_distributed_example.py           | 1 -
 transformer_engine/pytorch/quantization.py                   | 5 ++---
 3 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
index 74a1e90dcd..30665018ff 100644
--- a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
+++ b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
@@ -149,8 +149,8 @@ However, amax reduction works slightly differently in different frameworks.
       The gradient amaxes are reduced separately, once at the end of every ``backward()``
       call, on each rank where at least one quantized module ran backward. When a
       training step consists of several ``backward()`` calls (e.g. a 1F1B pipeline
-      schedule) or ends with deferred weight gradients (``backward_dw()``), wrap them in
-      ``quantization_backward_scope`` so the update runs once per step. The scope also
+      schedule), wrap them in ``quantization_backward_scope`` so the update runs once
+      per step. The scope also
       runs the update on ranks where no quantized module ran backward, so, like
       ``autocast``, it must be entered and exited on all ranks:
 
diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py b/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py
index afb21c200d..230e7fb7f0 100644
--- a/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py
+++ b/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py
@@ -21,5 +21,4 @@
 with te.quantization_backward_scope():
     for loss in microbatch_losses:
         loss.backward()
-    model.backward_dw()  # only with delay_wgrad_compute=True
 # END_BACKWARD_SCOPE_EXAMPLE
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 023ff4b48f..02fd2cae39 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -902,15 +902,14 @@ def quantization_backward_scope() -> None:
     across the amax reduction group and recomputing the scales. Ordinary
     ``loss.backward()`` calls do this automatically and do not need this scope.
 
-    Use it when one training step spans several autograd calls, so that the update
-    runs once when the outermost scope exits instead of after every call:
+    Use it when one training step consists of several ``backward()`` calls, so that
+    the update runs once when the outermost scope exits instead of after every call:
 
     .. code-block:: python
 
         with te.quantization_backward_scope():
             for loss in microbatch_losses:  # e.g. a 1F1B pipeline schedule
                 loss.backward()
-            model.backward_dw()  # deferred weight gradients
 
     Nested scopes are no-ops; only the outermost one triggers the update. A scope
     entered inside a running backward (e.g. in a hook) defers the update to the