[PyTorch] Scope the quantized-param caching flag to its own graph capture

[xiuhu17]

Description

make_graphed_callables(cache_quantized_params=True) leaks its weight-update flag into every later capture in the same process, which can leave the later module silently training on a stale quantized weight.

The mechanism

Quantized-parameter caching lets a graphed module reuse its quantized weight across microbatches. Because a CUDA graph must replay an identical kernel sequence, the quantize kernel is always launched and a noop_flag — a GPU scalar — decides at runtime whether it writes. That flag has to be process-global: its address is baked into the captured graph, and replay fills it in place.

Nothing scopes it to the capture that asked for it:

• _make_graphed_callables creates the flag when cache_quantized_params=True (graph.py) and never clears it.
• Linear, GroupedLinear, LayerNormLinear and LayerNormMLP read it whenever FP8GlobalStateManager.fp8_graph_capturing() is true — they never check whether this capture requested caching. Seeing a non-None flag forces is_first_microbatch = False, so the module takes the cached-workspace path and bakes the flag into its graph.
• Only a callable graphed with caching writes the flag at replay (graph.py, Graphed.forward).

So:

capture A (cache_quantized_params=True)        -> flag tensor created
run A with is_first_microbatch=False           -> flag = "skip the weight update"
capture B (cache_quantized_params=False)       -> B reads A's leftover flag,
                                                  bakes it in as its noop flag
replay B                                       -> B never writes the flag; it stays
                                                  "skip" forever

B's optimizer keeps updating the master weight, but its GEMM keeps using the quantized weight from capture time. No error, no warning — just output that drifts further from eager every step. In a two-Linear repro the graphed output diverges from eager by ~22 in bf16.

This is not specific to any recipe (reproduced on MXFP8, NVFP4, Float8BlockScaling, Float8CurrentScaling and DelayedScaling). It has stayed hidden because the CUDA-graph tests run each case in a fresh process, so no test ever performs two captures with different caching settings.

The fix

The flag tensor cannot be cleared or reallocated between captures: already-captured graphs bake in its address, and freeing it would make an earlier graph read recycled memory — worse than the original bug. So the read is gated instead:

• FP8GlobalState gains caching_quantized_params, tracking whether the capture in progress asked for caching.
• get_skip_fp8_weight_update_tensor() returns None unless that is set.
• make_graphed_callables sets the scope for the duration of the capture and clears it in a finally.
• The four modules read through the accessor instead of touching quantization_state directly.

The tensor's address, lifetime and write path are unchanged; only its visibility is now limited to the capture that declared it.

Test

test_make_graphed_callables_weight_caching_does_not_leak in tests/pytorch/test_cuda_graphs.py captures with caching, runs one cached microbatch to leave the flag set to "skip", then captures a second module without caching and asserts its graphed outputs match eager across two optimizer steps. Parametrized over every available recipe.

The test was mutation-checked: with the scope check removed at runtime it fails for all 7 recipes, and passes for all 7 with the fix in place. (Note the flag is only written for the first module of an autocast context, so each poisoning microbatch needs its own context — without that the test passes either way and proves nothing.)

Full tests/pytorch/test_cuda_graphs.py on B200: 676 passed, 1121 skipped, 0 failed.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings
• I have added tests that prove my fix is effective

[greptile-apps]

Greptile Summary

The PR confines the process-global quantized-weight update tensor to CUDA graph captures that explicitly request parameter caching.

• Adds capture-scoped caching state with guaranteed cleanup after graph construction.
• Routes four fused PyTorch modules through the scoped tensor accessor.
• Adds a multi-recipe regression test covering consecutive captures with different caching settings.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/graph.py	Establishes the caching scope before graph construction and clears it reliably in a finally block.
transformer_engine/pytorch/quantization.py	Adds the capture-scoping state and gates access to the persistent quantized-weight update tensor.
transformer_engine/pytorch/module/linear.py	Uses the scoped accessor so unrelated captures cannot inherit the persistent update tensor.
transformer_engine/pytorch/module/grouped_linear.py	Applies the scoped accessor consistently to both existing grouped-linear lookup sites.
transformer_engine/pytorch/module/layernorm_linear.py	Uses the scoped accessor for fused normalization and linear graph capture.
transformer_engine/pytorch/module/layernorm_mlp.py	Uses the scoped accessor for both cached MLP weight workspaces.
tests/pytorch/test_cuda_graphs.py	Adds regression coverage for sequential cached and uncached captures across supported quantization recipes.

Sequence Diagram

sequenceDiagram
  participant User
  participant Graph as make_graphed_callables
  participant State as FP8GlobalStateManager
  participant Module as Quantized module
  User->>Graph: "capture(cache_quantized_params=True)"
  Graph->>State: enable caching scope
  Graph->>State: initialize persistent update tensor
  Graph->>Module: capture forward
  Module->>State: request scoped update tensor
  State-->>Module: persistent tensor
  Graph->>State: disable caching scope in finally
  User->>Graph: "later capture(cache_quantized_params=False)"
  Graph->>State: keep caching scope disabled
  Graph->>Module: capture forward
  Module->>State: request scoped update tensor
  State-->>Module: None

Loading

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into fix_cache_quant..." | Re-trigger Greptile}