From b065243b3d208e7e5494b85f77b750ac31346a2b Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 20 Aug 2026 17:54:58 +0200 Subject: [PATCH 1/2] [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 --- 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 Date: Thu, 20 Aug 2026 18:03:22 +0200 Subject: [PATCH 2/2] 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 --- 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