[PyTorch] Preserve FP8 recompute state for inner autocast

[nvegesna-netizen]

Description

Activation checkpointing currently samples the global FP8 state before invoking
the checkpointed callable. If that callable opens its own FP8 autocast context,
the initial forward is incorrectly treated as a non-FP8 checkpoint region.
Delayed-scaling modules therefore do not preserve the scale and amax metadata
needed by the recompute forward.

Track the activation-recompute region independently of the FP8 state at context
entry, and gate the public query on the current FP8 state instead. This keeps
non-FP8 execution unchanged while allowing an inner FP8 autocast context to
participate in activation-recompute bookkeeping.

Changes

• Preserve FP8 delayed-scaling metadata when FP8 autocast begins inside a
  checkpointed callable.
• Add GPU regression coverage for inner-FP8, fully non-FP8, and mixed
  FP8/non-FP8 checkpoint regions with both reentrant and non-reentrant
  checkpoint implementations.

Validation

• The regression matrix checks that backward completes with finite loss and
  gradients, FP8 recompute metadata is preserved only for FP8 modules, and a
  fully non-FP8 checkpoint does not create FP8 recompute metadata.
• Validated with multi-GPU FP8 training using activation checkpointing and an
  FP8 autocast context opened inside the checkpointed callable. Training
  completed with finite loss and gradients, and the missing recompute metadata
  failure did not recur.
• git diff --check passes.
• Both modified Python files compile with Python 3.10.
• The repository-wide PyTorch lint script was attempted locally, but the host
  Python 3.9 interpreter cannot parse syntax already present on the current
  main branch. The reported errors were confined to unchanged files.

[greptile-apps]

Regression tests omitted from CI

The PyTorch QA job enumerates individual test files, but this new standalone regression suite is not included, so its FP8 activation-recompute coverage does not execute in CI.

Knowledge Base Used: Tests and QA

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[greptile-apps]

Greptile Summary

The PR updates FP8 activation-recompute bookkeeping and moves its regression coverage into the established numerics suite.

• Tracks activation-recompute regions independently from the FP8 state at checkpoint entry.
• Gates the public recompute query on the current FP8 autocast state.
• Exercises inner-autocast recomputation with both reentrant and non-reentrant checkpointing.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the prior CI-coverage issue is resolved because the regression now resides in test_numerics.py, which the PyTorch CI jobs execute directly without a test-name filter.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed.py	Separates checkpoint-region tracking from FP8-at-entry while retaining dynamic FP8 gating at module query time.
tests/pytorch/test_numerics.py	Adds inner-FP8-autocast numerical regression coverage to a test file already executed by the PyTorch CI jobs.

Sequence Diagram

sequenceDiagram
  participant Checkpoint
  participant Region as Recompute region
  participant Autocast as Inner FP8 autocast
  participant Module as FP8 module
  Checkpoint->>Region: Enter activation-recompute context
  Region->>Region: Mark checkpoint region active
  Checkpoint->>Autocast: Invoke checkpointed callable
  Autocast->>Module: Enable FP8 and run module
  Module->>Region: Query FP8 recompute bookkeeping
  Region-->>Module: Enabled while FP8 is currently active
  Module->>Module: Preserve scale and amax metadata

Loading

_{Reviews (5): Last reviewed commit: "test: compare inner FP8 autocast recompu..." | Re-trigger Greptile}

[pggPL]

With the gate moved to the getter, _FP8_ACTIVATION_RECOMPUTE_ENABLED no longer has anything to do with FP8 — it now just means "inside an activation recompute region". Worth renaming to _IN_ACTIVATION_RECOMPUTE_REGION; it's free, the global only appears in this file, and is_fp8_activation_recompute_enabled() keeps its name since it now returns the conjunction.

[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.