[PyTorch] Schedule delayed-scaling updates after backward

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.quantization_backward_scope
+
 .. autoapifunction:: transformer_engine.pytorch.quantized_model_init
 
 .. autoapifunction:: transformer_engine.pytorch.checkpoint

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), 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
 

docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py

@@ -16,3 +16,9 @@
     output = model(inp)
 
 # END_AMAX_REDUCTION_EXAMPLE
+
+# START_BACKWARD_SCOPE_EXAMPLE
+with te.quantization_backward_scope():
+    for loss in microbatch_losses:
+        loss.backward()
+# END_BACKWARD_SCOPE_EXAMPLE

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,74 @@ 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):
+    # 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
+    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)
+    FP8GlobalStateManager.reset()
+
+
 if __name__ == "__main__":
     sys.exit(main())

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",
-        "reduce_and_update_bwd_fp8_tensors",
     )
     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, "reduce_and_update_bwd_fp8_tensors")),
     )
 
 
@@ -816,7 +814,6 @@ def _run_grouped_linear_single_step_with_ctx_state(
         required_attrs = (
             "backward_override",
             "fp8",
-            "reduce_and_update_bwd_fp8_tensors",
         )
         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, "reduce_and_update_bwd_fp8_tensors")),
         )
     y.backward(dy)
     assert x_run.grad is not None
@@ -1449,37 +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_reduce_and_update,
-    ) = default_ctx
+    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_reduce_and_update
 
     *_, 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_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_reduce_and_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_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_reduce_and_update_after
 
 
 @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)
@@ -1526,10 +1507,9 @@ 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_ctx
     assert default_mode is None
     assert default_fp8
-    assert default_reduce_and_update
 
     *_, switched_ctx = _run_grouped_linear_single_step_with_ctx_state(
         module,
@@ -1538,10 +1518,9 @@ 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_ctx
     assert switched_mode == backward_override
     assert not switched_fp8
-    assert not switched_reduce_and_update
 
     *_, default_ctx_after = _run_grouped_linear_single_step_with_ctx_state(
         module,
@@ -1550,10 +1529,9 @@ 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_ctx_after
     assert default_mode_after is None
     assert default_fp8_after
-    assert default_reduce_and_update_after
 
 
 @pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list)