[PyTorch] Schedule delayed-scaling updates after backward

[pggPL]

Description

Delayed-scaling (amax/scale) updates are currently owned by the first FP8 module seen in forward and executed from that module's backward. This depends on backward traversal order and breaks with unused branches, reentrant checkpointing, or schedules that delay weight-gradient computation.

This PR moves the update to the end of the autograd task instead. Each participating module requests an update and TE queues a single callback on the enclosing GraphTask, so the update runs once after the whole backward has finished. Reentrant TE checkpointing registers the callback on the outer task before entering its nested backward.

Plain .backward() needs no changes. A new quantization_backward_scope() lets applications define a larger logical backward, so the update runs once when the outermost scope exits:

with te.quantization_backward_scope():
    for loss in microbatch_losses:  # e.g. a 1F1B pipeline schedule
        loss.backward()

The update always runs on scope exit, even on ranks that ran no quantized backward inside it, so every rank joins the amax reduction. Like autocast, the scope must be entered and exited on all ranks.

New API

transformer_engine.pytorch.quantization_backward_scope() (context manager, documented in the PyTorch API reference and on the delayed scaling feature page):

• Delays the delayed-scaling update (amax reduction across the reduction group and scale recomputation for the gradient quantizers) until the outermost scope exits. Nested scopes are no-ops. A scope entered inside a running backward defers the update to the end of that backward.
• Always runs the update on exit, even if no quantized module ran backward inside the scope. Like autocast, it must be entered and exited on all ranks.
• Has no effect for recipes without delayed-scaling state.

Internal hooks used by the modules: FP8GlobalStateManager.backward_quantization_update_needed() (decided in forward, stored as a bool on the autograd ctx) and FP8GlobalStateManager.request_backward_quantization_update() (called from backward, queues one callback per autograd task).

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Delegate delayed-scaling updates to RecipeState.
• Queue one recipe-state update at the end of each autograd task instead of in the first module's backward.
• Decide in forward whether a module must request the update (backward_quantization_update_needed()), carry only a bool through the autograd ctx, and request from backward with request_backward_quantization_update().
• Handle reentrant TE checkpoints by registering on the outer GraphTask.
• Add quantization_backward_scope() for custom logical-backward boundaries; it always updates on exit.
• Keep the CUDA graph path and the reduce_and_update_fp8_tensors alias working.
• Add tests for plain, reentrant, non-reentrant, nested, branched, unused-branch, multi-backward, and delayed-wgrad cases, plus a distributed test where some ranks skip backward.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR moves delayed-scaling gradient amax and scale updates from backward traversal order to an end-of-autograd-task callback, and introduces a public scope for grouping multiple backward calls.

• Delegates global delayed-scaling updates through recipe-state implementations.
• Records update intent during forward and requests one callback from participating module and fused-operation backward paths.
• Handles reentrant Transformer Engine checkpointing by associating updates with the enclosing autograd task.
• Adds distributed, checkpoint, branching, multi-backward, and delayed-weight-gradient coverage.
• Documents and exports quantization_backward_scope().

Confidence Score: 5/5

The PR appears safe to merge with no concrete blocking or non-blocking defects identified.

The callback scheduler, recipe-state delegation, module integrations, checkpoint handling, public scope, and distributed participation behavior are internally consistent and covered across the major changed execution paths.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/quantization.py	Introduces global pending-update scheduling, GraphTask callback deduplication, recipe-state update delegation, and the logical backward scope.
transformer_engine/pytorch/distributed.py	Wraps reentrant checkpoint backward so nested recomputation associates delayed-scaling updates with the enclosing backward task.
transformer_engine/pytorch/module/linear.py	Replaces traversal-order-owned FP8 updates with a forward-captured request flag consumed after linear backward.
transformer_engine/pytorch/module/grouped_linear.py	Applies callback-based delayed-scaling update requests to both grouped-tensor and legacy grouped-linear backward paths.
transformer_engine/pytorch/module/layernorm_linear.py	Moves delayed-scaling update ownership from the first FP8 module to the global end-of-task scheduler.
transformer_engine/pytorch/module/layernorm_mlp.py	Integrates the new update request mechanism with fused MLP backward and recomputation paths.
transformer_engine/pytorch/ops/fuser.py	Schedules delayed-scaling updates after the fused operation pipeline completes backward.
tests/pytorch/test_numerics.py	Adds coverage for plain, checkpointed, nested, branched, scoped, multi-graph, and delayed-weight-gradient schedules.
tests/pytorch/distributed/run_numerics.py	Verifies ranks that skip modules or all backward work still participate in scoped delayed-scaling reduction.

Sequence Diagram

sequenceDiagram
    participant App
    participant Module
    participant Autograd
    participant Manager as FP8GlobalStateManager
    participant Recipe as RecipeState

    App->>Module: Quantized forward
    Module->>Module: Save update-needed flag
    App->>Autograd: backward()
    Autograd->>Module: Module backward
    Module->>Manager: Request backward quantization update
    Manager->>Autograd: Queue one GraphTask callback
    Autograd-->>Manager: Run callback after task completes
    Manager->>Recipe: Reduce gradient amaxes and update scales

    opt Multiple backward calls
        App->>Manager: Enter quantization_backward_scope
        App->>Autograd: backward() per microbatch
        Manager->>Manager: Defer pending update
        App->>Manager: Exit outermost scope
        Manager->>Recipe: Run one logical-backward update
    end

Loading

_{Reviews (1): Last reviewed commit: "docs: drop backward_dw from the quantiza..." | Re-trigger Greptile}