[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

This PR fixes FP8 metadata preservation when activation checkpointing contains an inner FP8 autocast region.

• Tracks activation-recompute regions independently from FP8 state at checkpoint entry while keeping the public query gated by current FP8 enablement.
• Adds regression coverage for inner-FP8, non-FP8, and mixed checkpoint regions with reentrant and non-reentrant implementations.
• Registers the new regression suite in the L0 PyTorch unittest job, resolving the previously reported CI omission.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously omitted regression suite is now explicitly invoked by the L0 PyTorch unittest script.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed.py	Separates checkpoint-region tracking from FP8 enablement at context entry and applies the FP8 gate when recompute state is queried.
tests/pytorch/test_fp8_activation_recompute.py	Adds GPU regression tests covering FP8 metadata preservation and exclusion across checkpoint modes.
qa/L0_pytorch_unittest/test.sh	Explicitly invokes the new activation-recompute regression suite in the curated L0 PyTorch CI job.

_{Reviews (3): Last reviewed commit: "test: run FP8 recompute coverage in PyTo..." | Re-trigger Greptile}