[PyTorch] Preserve FP8 recompute state for inner autocast

[nvegesna-netizen]

Description

The FP8 activation-recompute stash and its restore are gated on conditions
sampled at different times, so they can disagree.

The stash in prepare_forward is gated on is_fp8_activation_recompute_enabled(),
which was latched from FP8GlobalStateManager.is_fp8_enabled() at the moment
te.checkpoint was entered. The matching restore is gated on self.fp8, which
is read when the module actually runs. If the checkpointed callable opens its own
autocast — selecting precision per layer, for example — FP8 is off at checkpoint
entry and on at module execution. Phase 1 then skips the stash while the recompute
phase still restores, and get_old_fp8_meta_tensors_for_recompute raises
KeyError: 'global_fp8_buffer_pos_fwd_recompute'.

This is delayed-scaling only: the stash helpers early-return for every other recipe.

Megatron already works around this rather than hitting it. transformer_block.py
routes delayed scaling to an outer context and everything else to an inner one, and
multi_token_prediction.py carries an eleven-line comment naming this exact failure
(2b77d32b1, May 2026) explaining why the outer context is required there. The trap
is real enough that someone paid for it; this change removes it rather than requiring
callers to route around it.

Changes

• Track the checkpoint region independently of the FP8 state at entry, and evaluate
  the FP8 state in is_fp8_activation_recompute_enabled() instead, so both gates
  derive from the same live autocast state.
• Rename the two module-private globals to match what they now hold. The public
  getters keep their names.
• Add reentrant and non-reentrant coverage for inner-FP8, non-FP8 and mixed
  checkpoint regions in test_numerics.py, which is already run by L0.

A behaviour change worth calling out

is_fp8_activation_recompute_enabled() also drives columnwise_usage in
linear.py and grouped_linear.py, and that path is not delayed-scaling only.
In the inner-autocast pattern the predicate now returns True for every recipe, so
the weight quantizer is asked for a columnwise copy in phase 1 where it previously
was not. That matches what the outer-autocast pattern already does, and phase 1 runs
under no_grad, so this is the value the recompute needs — but it is a memory and
work change on a recipe-independent path, and it is not covered by the tests here.
Flagging it explicitly for review.

Validation

• All three new tests pass on H100 and GB200, for both use_reentrant values.
• The wider -k activation_recompute family passes on both: 218 tests on H100, 416
  on GB200 (the Blackwell recipe list is longer, so MXFP8, block-scaling and NVFP4
  recompute paths are exercised there too).
• black, pylint 3.3.1 (10.00/10 on the changed source file) and the license check
  pass locally.

Known limitation, not addressed here

The stash is gated on self.training and the restores are not, so a module left in
.eval() inside a grad-enabled checkpointed region still hits the same KeyError.
That is the other half of the same asymmetry and needs a different fix — pairing the
restore to a decision recorded at stash time rather than re-deriving it in a later
autograd phase — so it is left for a follow-up rather than bundled here.

[greptile-apps]

Greptile Summary

Updates FP8 activation-recompute tracking to use the live autocast state within checkpoint regions.

• Separates checkpoint-region and recompute-phase state from FP8 enablement.
• Preserves delayed-scaling metadata when FP8 autocast begins inside a checkpointed callable.
• Adds reentrant and non-reentrant regression coverage for inner-autocast and mixed FP8 regions in the existing L0 numerics suite.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed.py	Separates activation-recompute region tracking from live FP8 enablement so inner autocast contexts receive matching stash and restore behavior.
tests/pytorch/test_numerics.py	Adds inner-autocast and mixed-region regression tests to a numerics file that is executed wholesale by L0 PyTorch QA.

Sequence Diagram

sequenceDiagram
  participant Checkpoint
  participant Region as Recompute region state
  participant Callable
  participant Autocast as Inner FP8 autocast
  participant Module as FP8 module
  Checkpoint->>Region: Enter checkpoint region
  Checkpoint->>Callable: Execute forward
  Callable->>Autocast: Enable FP8
  Autocast->>Module: Run module
  Module->>Region: Query region and live FP8 state
  Region-->>Module: FP8 recompute enabled
  Module->>Module: Stash FP8 metadata
  Checkpoint->>Region: Enter recompute phase
  Checkpoint->>Callable: Replay forward
  Callable->>Autocast: Enable FP8
  Autocast->>Module: Run module
  Module->>Module: Restore stashed FP8 metadata

Loading

_{Reviews (14): Last reviewed commit: "Merge branch 'main' into fix/fp8-inner-a..." | Re-trigger Greptile}

[pggPL]

Could the tests go into tests/pytorch/test_numerics.py instead of a new file? test_gpt_full_activation_recompute (line 729) already parametrizes use_reentrant and the full fp8_recipes matrix, and asserts numerical equality against the non-recomputed run — the new tests only check isfinite plus the presence of "global_fp8_buffer_pos_fwd_recompute", so they'd also pass for a fix that merely suppressed the KeyError while leaving the scales inconsistent. The change to _test_e2e_full_recompute (line 650) would be small: an inner_autocast flag moving the autocast at line 687 inside the checkpointed callable, with the reference run unchanged. To avoid blowing up an already large matrix, I'd rather not parametrize the existing test on that flag — a separate narrow test calling the same helper with inner_autocast=True for one dtype/recipe (plus use_reentrant) is enough. That would also drop the need for the new qa/L0_pytorch_unittest/test.sh entry, since test_numerics.py is already in L0.

[nvegesna-netizen]

Done — moved into test_numerics.py, no new file and no qa/L0_pytorch_unittest/test.sh entry, so the diff is two files. (This comment consolidates a few I posted earlier today; it is the current state.)

What changed in the tests

• _test_e2e_full_recompute gained an inner_autocast flag that moves the autocast inside the checkpointed callable, with autocast(enabled=fp8 and not inner_autocast) outside. At the default inner_autocast=False it is byte-equivalent to before, so every pre-existing caller is unchanged.
• test_gpt_full_activation_recompute_with_inner_autocast is narrow as you asked: bf16, DelayedScaling/HYBRID, 126m, bs=1, parametrized only on use_reentrant. test_gpt_full_activation_recompute is untouched, so the big matrix doesn't grow.
• It compares against the non-recomputed reference with assert_close at rtol=0.125, atol=0.0675, plus per-module assertions that every stashed forward scale is restored exactly once.

I also restored two tests that my earlier move had silently dropped:

• test_checkpoint_with_mixed_fp8_regions_saves_only_fp8_recompute_state — the mixed FP8/non-FP8 region case. This is the genuine regression test: it fails on main with the KeyError during backward() for both use_reentrant values.
• test_checkpoint_without_fp8_does_not_save_fp8_recompute_state — to be accurate, this one is not regression coverage; it passes on main too, since the stash is already double-guarded by delayed_scaling_recipe. Its value is pinning the design, i.e. that the gate moved into the getter rather than being deleted. Kept deliberately without an FP8 skipif so it runs on any bf16-capable GPU.

Validation

Ran on an H100, pinned to the branch commit:

test_gpt_full_activation_recompute_with_inner_autocast[True]                  PASSED
test_gpt_full_activation_recompute_with_inner_autocast[False]                 PASSED
test_checkpoint_without_fp8_does_not_save_fp8_recompute_state[True]           PASSED
test_checkpoint_without_fp8_does_not_save_fp8_recompute_state[False]          PASSED
test_checkpoint_with_mixed_fp8_regions_saves_only_fp8_recompute_state[True]   PASSED
test_checkpoint_with_mixed_fp8_regions_saves_only_fp8_recompute_state[False]  PASSED
6 passed, 1189 deselected

pytest tests/pytorch/test_numerics.py -k activation_recompute
218 passed, 977 deselected

Locally: black clean, pylint 3.3.1 --rcfile=pylintrc on distributed.py = 10.00/10, license check passes, all commits DCO signed.

Note the last commit (assertion rework + a is_bf16_available skip + a docstring fix) landed after that run, so it needs a re-run; everything else above was validated at the commit stated.

One limitation worth stating

The numeric comparison in the inner-autocast test cannot currently catch a wrong-scale restore — only a missing one. autocast_enter increments autocast_depth unconditionally, and autocast_exit only updates scales when the depth returns to 0 with FP8 enabled; the outer autocast(enabled=False) pins it at >= 1, so reduce_and_update_fp8_tensors(forward=True) is never called and the forward scale stays at 1.0 throughout. Stash and no-stash are therefore numerically identical, and the load-bearing checks are the stash/restore bookkeeping assertions plus the fact that the pre-fix bug crashes outright.

The minimal way to give the numeric arm teeth is contextlib.nullcontext() instead of autocast(enabled=False) when inner_autocast is set: the inner autocast then exits at depth 0 and, in the non-reentrant case, a scale update fires between the two forward phases. I haven't folded that in because it makes the update counts asymmetric between the two use_reentrant arms, which is the same class of asymmetry that made the original outer-autocast reference a latent false positive. Happy to add it here or as a follow-up — your call.

Separately, could you kick off /te-ci when convenient? The GitHub Actions jobs here are build + sanity only, so they never execute these tests.

---

Scope and a known remaining gap

Writing this out because the change is easy to under-describe.

What it makes reachable. Previously, delayed-scaling recipes crashed with KeyError: 'global_fp8_buffer_pos_fwd_recompute' whenever a checkpointed callable opened its own autocast and no enabled autocast was active at the te.checkpoint call site. That hit both use_reentrant values, not just reentrant. Non-delayed recipes (current scaling, MXFP8, block scaling, NVFP4) never entered this path at all — both stash helpers early-return via _has_delayed_scaling_state — so they already ran, and their numerics are unchanged here. Note the trigger is "no enabled autocast at the call site", not "an inner autocast exists": outer-enabled plus inner-enabled, the documented nested-recipe pattern, always worked.

Known remaining gap. With use_reentrant=True (the default, distributed.py:63/:703, no env override) and delayed scaling and no enclosing enabled autocast, the forward amax reduction and scale update never fire. autocast_exit is the only forward-update site (quantization.py:769-772) and needs autocast_depth == 0 and torch.is_grad_enabled(); the reentrant phase-1 forward runs under torch.no_grad() (distributed.py:375), and during recompute ctx.fp8 is captured after the inner autocast has already exited (distributed.py:393), so the wrapper is disabled yet still holds depth at 1. Two independent gates, one per phase.

This gating predates the change and is untouched by it — it was simply unreachable while the configuration crashed. use_reentrant=False is unaffected and updates normally at distributed.py:813. Backward scales are unaffected; they update through the separate is_first_fp8_module path.

It also needs qualifying: because reduce_and_update_fp8_tensors iterates a process-global buffer, "scales stay at their initial value" only holds while every delayed-scaling module lives inside such a region. One FP8 module running outside flushes the update for all of them.

How much it costs, if hit. The quantization stays self-consistent — scale_inv is derived from the same scale at cast time (float8_tensor.py:141) and handed to cuBLASLt, so magnitudes are preserved and there is no gain error. The cost is lost E4M3 range on the low side only; scale=1.0 is the most overflow-safe scale, so nothing saturates. Simulated E4M3 cast+descale, RMS relative error vs an ideal 448/amax scale:

tensor sigma	err @ scale=1	err @ ideal	ratio
1.0 (activations)	2.65%	2.65%	1.00x
0.03 (typical weight init)	2.87%	2.65%	1.08x
0.02	3.30%	2.65%	1.25x
0.005 (deep-model init)	11.3%	2.65%	4.26x

So: negligible for activations, a few percent for typical weights, and it degrades sharply only for very small-magnitude tensors. I would call it suboptimal FP8 range utilization rather than an accuracy bug — but it is entirely silent, with no warning or diagnostic anywhere.

I have not fixed it here. The phase-1 gate is deliberate, and the update it guards triggers a cross-rank amax collective, so relaxing it changes when a collective fires — not something to guess at without multi-GPU validation. I would rather disclose it than paper over it. Happy to follow up with either the gate relaxation or a one-time warning when a delayed-scaling module is stashed in a recompute region that cannot reach depth 0, alongside two other pre-existing items I found while reviewing this (the non-restoring __exit__ under nested te.checkpoint, and the self.training push vs ungated pop asymmetry).

One more scope note. is_fp8_activation_recompute_enabled() is now evaluated dynamically rather than sampled at checkpoint entry, so in the inner-autocast pattern it returns True for all recipes. That also feeds columnwise weight-quantizer usage (module/linear.py:441-445, module/grouped_linear.py:569-573). It is memory/perf-shaped rather than numerics-shaped, and monotone — the flag can only turn on where it was previously off — but it is worth a reviewer's eye.

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.