[PyTorch] Deferred backward quantization state update (removes first-FP8-module ownership)

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

tests/pytorch/test_fusible_ops.py

@@ -25,6 +25,7 @@
 )
 from transformer_engine.pytorch.ops.fuser import OperationFuser
 from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
 
 from transformer_engine.pytorch.ops.fused import (
     BackwardActivationBias,
@@ -1030,6 +1031,7 @@ def test_fp8_scale_update(
             with te.autocast(recipe=recipe):
                 y = model(x)
             y.backward(dy)
+            FP8GlobalStateManager.flush_backward_quantization_update()
             with torch.no_grad():
                 model.weight.fill_(w_vals[step + 1])
 

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,
@@ -94,6 +95,7 @@ def test_fp8_scale_update_with_linear_module(
                 is_first_microbatch=True,
             )
         y.backward(torch.zeros_like(y))
+        FP8GlobalStateManager.flush_backward_quantization_update()
 
         # Get amax history and scaling factors
         fp8_meta = module.fp8_meta
@@ -147,6 +149,7 @@ def test_fp8_scale_update_with_linear_module(
             x = torch.randn([16, 16], device="cuda")
             y = module(x, is_first_microbatch=is_first_microbatch)
         y.backward(torch.randn_like(y))
+        FP8GlobalStateManager.flush_backward_quantization_update()
 
         # Check that amax history matches expected values
         torch.testing.assert_close(
@@ -245,6 +248,7 @@ def test_fp8_scale_update_with_linear_fuser_op(
             with te.autocast(recipe=recipe):
                 y = op(x)
             y.backward(dy)
+            FP8GlobalStateManager.flush_backward_quantization_update()
 
             def check_metas(
                 test_scale: float,
@@ -780,3 +784,256 @@ def test_stateful_unknown_or_malformed_pickled_extra_state_requires_opt_in(paylo
 
     monkeypatch.setenv(UNSAFE_PICKLE_EXTRA_STATE_ENV, "1")
     assert should_load_extra_state_pickle(payload, "test")
+
+
+###############################################################################
+# Backward quantization state update scheduling
+###############################################################################
+
+
+HIDDEN = 128
+BATCH = 32
+STEPS = 3
+
+
+class _UpdateCounter:
+    """Counts reduce_and_update_quantization_state calls by direction."""
+
+    def __init__(self):
+        self.forward = 0
+        self.backward = 0
+        self._original = None
+
+    def __enter__(self):
+        self._original = FP8GlobalStateManager.reduce_and_update_quantization_state.__func__
+        original = self._original
+        counter = self
+
+        def counted(cls, forward=True):
+            if forward:
+                counter.forward += 1
+            else:
+                counter.backward += 1
+            return original(cls, forward=forward)
+
+        FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(counted)
+        return self
+
+    def __exit__(self, *exc):
+        FP8GlobalStateManager.reduce_and_update_quantization_state = classmethod(self._original)
+
+
+def _make_model(num_layers=3, seed=1234):
+    torch.manual_seed(seed)
+    return torch.nn.ModuleList(
+        [te.Linear(HIDDEN, HIDDEN, bias=True).cuda() for _ in range(num_layers)]
+    )
+
+
+def _run_layers(layers, x):
+    for layer in layers:
+        x = layer(x)
+    return x
+
+
+def _train_step(model, x, forward_fn, recipe):
+    with te.autocast(enabled=True, recipe=recipe):
+        out = forward_fn(model, x)
+    loss = out.float().sum()
+    loss.backward()
+    return loss
+
+
+def _forward_plain(model, x):
+    return _run_layers(model, x)
+
+
+def _forward_reentrant(model, x):
+    return te_checkpoint(_run_layers, model, x, use_reentrant=True)
+
+
+def _forward_non_reentrant(model, x):
+    return te_checkpoint(_run_layers, model, x, use_reentrant=False)
+
+
+def _forward_per_layer_reentrant(model, x):
+    for layer in model:
+        x = te_checkpoint(layer, x, use_reentrant=True)
+    return x
+
+
+def _forward_per_layer_non_reentrant(model, x):
+    for layer in model:
+        x = te_checkpoint(layer, x, use_reentrant=False)
+    return x
+
+
+def _forward_nested(model, x):
+    def inner(x):
+        return te_checkpoint(model[1], x, use_reentrant=True)
+
+    def outer(x):
+        x = model[0](x)
+        x = inner(x)
+        return model[2](x)
+
+    return te_checkpoint(outer, x, use_reentrant=True)
+
+
+FORWARD_FNS = {
+    "plain": _forward_plain,
+    "reentrant": _forward_reentrant,
+    "non_reentrant": _forward_non_reentrant,
+    "per_layer_reentrant": _forward_per_layer_reentrant,
+    "per_layer_non_reentrant": _forward_per_layer_non_reentrant,
+    "nested": _forward_nested,
+}
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("mode", FORWARD_FNS.keys())
+def test_single_update_per_step(mode):
+    """Exactly one forward and one backward state update per training step."""
+    forward_fn = FORWARD_FNS[mode]
+    model = _make_model()
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter:
+        for step in range(STEPS):
+            x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+            _train_step(model, x, forward_fn, recipe)
+            # The backward update of this step is flushed at the next
+            # top-level autocast entry; flush explicitly to count it here.
+            FP8GlobalStateManager.flush_backward_quantization_update()
+            assert counter.forward == step + 1, f"{mode}: duplicate/missing forward update"
+            assert counter.backward == step + 1, f"{mode}: duplicate/missing backward update"
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_flush_at_next_autocast_entry():
+    """Without an explicit flush, the update runs when the next autocast begins."""
+    model = _make_model()
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter:
+        x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+        _train_step(model, x, _forward_plain, recipe)
+        assert counter.backward == 0
+        x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+        _train_step(model, x, _forward_plain, recipe)
+        assert counter.backward == 1
+    FP8GlobalStateManager.flush_backward_quantization_update()
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("checkpoint_first_branch", [True, False])
+def test_branched_graph(checkpoint_first_branch):
+    """Two checkpointed sibling branches merging into one loss."""
+    branch_a = _make_model(num_layers=2, seed=1)
+    branch_b = _make_model(num_layers=2, seed=2)
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter:
+        for step in range(STEPS):
+            x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+            with te.autocast(enabled=True, recipe=recipe):
+                if checkpoint_first_branch:
+                    ya = te_checkpoint(_run_layers, branch_a, x, use_reentrant=True)
+                else:
+                    ya = _run_layers(branch_a, x)
+                yb = te_checkpoint(_run_layers, branch_b, x, use_reentrant=True)
+                out = ya + yb
+            loss = out.float().sum()
+            loss.backward()
+            FP8GlobalStateManager.flush_backward_quantization_update()
+            assert counter.forward == step + 1
+            assert counter.backward == step + 1
+            assert x.grad is not None and torch.isfinite(x.grad).all()
+            x.grad = None
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_checkpointed_branch_without_backward():
+    """A checkpointed branch whose output never receives a gradient must not
+    strand the backward update (regression test for first-module ownership)."""
+    used = _make_model(num_layers=2, seed=1)
+    unused = _make_model(num_layers=2, seed=2)
+    recipe = DelayedScaling()
+
+    with _UpdateCounter() as counter:
+        for step in range(STEPS):
+            x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+            with te.autocast(enabled=True, recipe=recipe):
+                # The unused branch runs first, so any "first module owns the
+                # backward update" scheme would assign ownership to a frame
+                # whose backward never executes.
+                y_unused = te_checkpoint(_run_layers, unused, x, use_reentrant=True)
+                y = _run_layers(used, x)
+            loss = y.float().sum()
+            loss.backward()
+            del y_unused
+            FP8GlobalStateManager.flush_backward_quantization_update()
+            assert counter.backward == step + 1, "backward update was lost"
+            x.grad = None
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("mode", ["reentrant", "non_reentrant", "per_layer_reentrant", "nested"])
+def test_checkpointing_matches_plain_numerics(mode):
+    """Delayed-scaling state must evolve identically with and without
+    activation checkpointing (duplicate updates would advance it faster)."""
+    forward_fn = FORWARD_FNS[mode]
+    recipe = DelayedScaling()
+
+    def train(forward):
+        model = _make_model(seed=99)
+        losses = []
+        torch.manual_seed(777)
+        for _ in range(STEPS + 1):
+            x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+            losses.append(_train_step(model, x, forward, recipe).detach())
+        FP8GlobalStateManager.flush_backward_quantization_update()
+        state = []
+        for layer in model:
+            for key in ("scaling_fwd", "scaling_bwd"):
+                meta = layer.fp8_meta[key]
+                state.append((meta.scale.clone(), meta.amax_history.clone()))
+        return losses, state
+
+    losses_ref, state_ref = train(_forward_plain)
+    losses_ckpt, state_ckpt = train(forward_fn)
+
+    for step, (l_ref, l_ckpt) in enumerate(zip(losses_ref, losses_ckpt)):
+        torch.testing.assert_close(l_ckpt, l_ref, rtol=0, atol=0, msg=f"loss diverged @ {step}")
+    for (scale_ref, hist_ref), (scale_ckpt, hist_ckpt) in zip(state_ref, state_ckpt):
+        assert torch.equal(scale_ckpt, scale_ref), "scale diverged under checkpointing"
+        # With nested checkpoints the inner checkpoint re-runs its forward (as a
+        # regular forward frame) during the outer recompute and re-records the
+        # same amaxes into the history, so exact history equality does not hold.
+        if mode != "nested":
+            assert torch.equal(hist_ckpt, hist_ref), "amax history diverged under checkpointing"
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_exception_in_checkpointed_forward():
+    """A failing checkpointed forward must not corrupt update scheduling."""
+    model = _make_model()
+    recipe = DelayedScaling()
+
+    def failing(model, x):
+        def fn(x):
+            model[0](x)
+            raise RuntimeError("boom")
+
+        return te_checkpoint(fn, x, use_reentrant=True)
+
+    with _UpdateCounter() as counter:
+        x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+        with pytest.raises(RuntimeError, match="boom"):
+            _train_step(model, x, failing, recipe)
+
+        x = torch.randn(BATCH, HIDDEN, device="cuda", requires_grad=True)
+        _train_step(model, x, _forward_reentrant, recipe)
+        FP8GlobalStateManager.flush_backward_quantization_update()
+        assert counter.backward == 1
+        assert torch.isfinite(x.grad).all()