[PyTorch] Finalize quantization updates at logical backward boundaries

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

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_numerics.py

@@ -43,8 +43,10 @@
 )
 from transformer_engine.pytorch import checkpoint as te_checkpoint
 from transformer_engine.pytorch.distributed import (
-    is_fp8_activation_recompute_enabled,
+    activation_recompute_forward,
+    in_fp8_activation_recompute_forward_phase,
     in_fp8_activation_recompute_phase,
+    is_fp8_activation_recompute_enabled,
 )
 from transformer_engine.pytorch.cpp_extensions import general_gemm
 from transformer_engine.common import recipe
@@ -906,6 +908,23 @@ def _checkpointed_linear_backward(body, use_reentrant, *layers):
         assert torch.isfinite(layer.weight.grad).all()
 
 
+def test_nested_activation_recompute_phases():
+    """Nested checkpoint forwards preserve an active outer recompute phase."""
+    FP8GlobalStateManager.reset()
+
+    with activation_recompute_forward(True, True):
+        assert in_fp8_activation_recompute_phase()
+        assert not in_fp8_activation_recompute_forward_phase()
+        with activation_recompute_forward(True, False):
+            assert in_fp8_activation_recompute_phase()
+            assert in_fp8_activation_recompute_forward_phase()
+        assert in_fp8_activation_recompute_phase()
+        assert not in_fp8_activation_recompute_forward_phase()
+
+    assert not in_fp8_activation_recompute_phase()
+    assert not in_fp8_activation_recompute_forward_phase()
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 @pytest.mark.parametrize("use_reentrant", all_boolean)
 def test_checkpoint_inner_autocast_is_an_fp8_recompute_region(use_reentrant):