[PyTorch] Pair delayed-scaling FP8 recompute metadata per module

[nvegesna-netizen]

Description

Fix delayed-scaling FP8 metadata stash/restore pairing when a checkpointed module is in eval mode or changes mode between the original forward and recompute forward.

The original forward stashed metadata only when self.training was true, while recompute restored metadata from every FP8 module in the recompute phase. An eval module could therefore try to restore a stash it never created. Conversely, deriving the two decisions independently makes mode changes capable of leaking or mispairing FIFO entries.

Each module now records its pending recompute stashes and whether the current invocation actually restored one. Every delayed-scaling FP8 module in checkpoint phase 1 stashes: reentrant checkpointing runs that forward under no_grad, so an eval module receiving an intermediate tensor has no reliable module-local autograd signal even though backward will replay it. Phase 2 consumes only a stash recorded by that module, and end_forward restores live metadata only after a matching consume.

Type of change

• Bug fix (non-breaking change which fixes an issue)
• Documentation change (change only to the documentation, either a fix or a new content)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Track pending delayed-scaling recompute stashes per module.
• Pair phase-2 consumption and end-of-forward restoration with the stash that produced them.
• Stash eval modules in checkpoint phase 1, including reentrant intermediate-input modules under no_grad.
• Reset pending bookkeeping when loading FP8 extra state invalidates the saved buffer position.
• Add multi-iteration regressions for train/eval modes, both checkpoint implementations, mode changes, FIFO drainage, and the reentrant intermediate-input case.

Validation

• Ten focused cases cover both reentrant and non-reentrant checkpoint implementations.
• Reverting only module/base.py produces the expected eval-module and reentrant-intermediate failures.
• The broader PyTorch recompute/checkpoint test selection passes on FP8-capable GPU hardware.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation (not applicable: no user-facing API changes)
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing relevant unit tests pass with my changes

[greptile-apps]

Greptile Summary

The PR adjusts delayed-scaling FP8 activation-recompute bookkeeping so checkpointed modules stash metadata independently of training mode and avoids creating unreachable recompute state when autograd is disabled.

• Stashes delayed-scaling metadata for every FP8 module participating in checkpoint phase one.
• Runs checkpointed TE functions directly under no_grad, preserving the user’s forward context.
• Adds train/eval, mode-transition, reentrant/non-reentrant, FIFO-drainage, and no-backward regression coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Makes phase-one delayed-scaling metadata stashing independent of module training mode.
transformer_engine/pytorch/distributed.py	Bypasses checkpoint recompute bookkeeping when autograd is disabled while retaining the user-provided forward context.
tests/pytorch/test_numerics.py	Adds focused regression coverage for eval modules, mode changes, both checkpoint implementations, FIFO drainage, intermediate inputs, and no-backward execution.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[TE checkpoint invocation] --> B{Autograd enabled?}
  B -- No --> C[Enter user forward context]
  C --> D[Run function directly]
  B -- Yes --> E[Checkpoint phase 1]
  E --> F[Each delayed-scaling FP8 module stashes metadata]
  F --> G[Backward requests recompute]
  G --> H[Checkpoint phase 2]
  H --> I[Module consumes matching FIFO metadata]
  I --> J[Forward completes and live metadata is restored]

Loading

_{Reviews (2): Last reviewed commit: "Bypass FP8 recompute bookkeeping without..." | Re-trigger Greptile}

[pggPL]

After a longer discussion with Codex, we came to the following conclusion:

Could this be simplified by making the checkpoint phase the sole source of truth?

It looks like the valid cases covered here are fixed by removing the self.training condition:

- if self.training and is_fp8_activation_recompute_enabled():
+ if is_fp8_activation_recompute_enabled():
      FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta)

After that, every delayed-scaling module encountered in phase 1 creates a stash, and every such module in phase 2 should consume one. This covers eval modules, reentrant forwards running under no_grad, and train/eval mode changes without adding separate module state.

fp8_recompute_stashes appears to duplicate the state of the existing per-module deque and can diverge from it. For example, set_extra_state() resets the counter but does not remove the corresponding snapshots from the old global deque, potentially leaving them orphaned.

I am also concerned about silently continuing when fp8_recompute_stashes == 0. A delayed-scaling module appearing in recompute without a matching stash seems like an invariant violation—divergent checkpoint execution, state replacement between forward and backward, or a bookkeeping bug. Continuing with live FP8 metadata may produce an incorrect recompute instead of a clear failure.

Is there a supported execution path where phase 2 legitimately has no matching phase-1 stash after removing the self.training guard? If not, could we keep the strict one-to-one stash/consume behavior and reduce this PR to the guard removal plus the regression tests? If such a path does exist, it may need a checkpoint-frame token rather than a second per-module counter.

What do you think?