[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.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
@@ -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)

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,215 @@ 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_independent_graphs():
+    FP8GlobalStateManager.reset()
+    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:
+        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)
+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

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