[PyTorch] Deferred backward quantization state update (removes first-FP8-module ownership)

[pggPL]

Description

Refactor of how the backward delayed-scaling amax reduction is triggered.

Problem. reduce_and_update_fp8_tensors(forward=False) was fired from the backward of the first FP8 module in an autocast, on the assumption that its backward runs last. This breaks in two ways:

1. Activation recompute (Fix duplicate FP8 state updates during activation checkpoint recomputation NVIDIA/TransformerEngine#3213): the recompute re-runs forwards and re-enters autocasts, advancing delayed-scaling state multiple times per step (duplicate forward update at recompute-autocast exit + per-module ownership re-acquisition).
2. Non-linear autograd graphs: with branches (or a checkpointed branch whose output never receives a gradient), the owning module's backward is not guaranteed to run last — or at all — so the reduction can fire early or be lost entirely.

Fix. Module backwards no longer trigger the update. They only request it (FP8GlobalStateManager.request_backward_quantization_update(), a bool flag). The update is flushed exactly once at the next top-level autocast entry — at that point the previous backward is provably complete, whatever the graph topology or checkpointing scheme. Guards ensure the flush never runs during the recompute phase, inside a running backward (torch._C._current_graph_task_id()), under CUDA graph capture, or under torch.compile tracing. The forward update at autocast exit is skipped during the recompute phase (fixes the duplicate forward update; the recompute-phase flag now lives on the quantization state and also covers Megatron-Core's checkpoint, which wraps its recompute in TE's activation_recompute_forward).

This removes the is_first_fp8_module consumption from all eager modules and the save/restore bookkeeping in checkpoint contexts (_is_first_fp8_module FIFO, per-module restore hacks). The CUDA-graphs path (graph.py) keeps its own mechanism and is unchanged. Timing note: the backward update now runs at the start of the next step instead of the end of the current backward — in between, nothing reads the affected state, so scale trajectories are unchanged (asserted exactly in the new tests).

Naming: reduce_and_update_fp8_tensors → reduce_and_update_quantization_state (old name kept as an alias since Megatron-Core calls it), since delayed update methods will be used beyond FP8 (e.g. RHT in NVFP4).

Supersedes the approach in NVIDIA#3213: instead of scoping ownership to checkpoint frames, ownership is removed.

New tests (in tests/pytorch/test_recipe.py, already enumerated by the L0 QA suite)

• exactly 1 fwd + 1 bwd update per step across: plain, reentrant/non-reentrant TE checkpoint, per-layer, nested, torch.utils.checkpoint
• flush timing semantics
• branched graphs (two checkpointed siblings; checkpointed branch with no backward — fails under first-module ownership)
• exact scale/amax-history equivalence checkpointed vs plain over multiple steps (the Fix duplicate FP8 state updates during activation checkpoint recomputation NVIDIA/TransformerEngine#3213 regression)
• exception in checkpointed forward

Type of change

• Bug fix + refactor

🤖 Generated with Claude Code