[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. Module training mode is not a valid pairing signal: reentrant checkpointing deliberately runs the original forward under no_grad, and a module may also change mode before recompute.

Every delayed-scaling FP8 module in checkpoint phase 1 now stashes its metadata, independent of module training mode. Phase 2 retains the existing strict FIFO restore behavior, so an execution mismatch remains visible rather than being silently skipped.

When te.checkpoint() is entered with outer autograd disabled, no backward recompute is normally possible. TE-containing callables therefore execute directly under the supplied forward context, avoiding FP8 recompute snapshots that could never be consumed. Non-TE callables continue to use native PyTorch checkpointing.

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

• Stash delayed-scaling metadata for every FP8 module participating in checkpoint phase 1, including eval modules and reentrant intermediate-input modules.
• Keep strict phase-2 FIFO restoration and the existing delayed-scaling end-of-forward restore invariant.
• Bypass TE checkpoint/recompute bookkeeping when checkpoint is entered with outer autograd disabled, while preserving the user-provided forward context.
• Add multi-iteration regressions for train/eval modes, both checkpoint implementations, mode changes, FIFO drainage, no-backward execution, and explicit nested gradient enablement.

Validation

• Eighteen focused cases cover reentrant and non-reentrant checkpointing, train/eval modes, mode changes, no-backward forwards, and explicit nested gradients enabled through either context_fn or the callable.
• The explicit-gradient cases compare output, input gradient, and weight gradient against direct execution and verify that no FP8 recompute snapshots remain.
• Reverting the relevant production changes reproduces the expected eval/autograd, no-backward, and explicit-gradient failures.
• The broader PyTorch recompute/checkpoint selection passes on FP8-capable GPU hardware (450 passed, 90 capability skips).

Grad-mode boundary

Checkpoint-entry grad mode is authoritative. If te.checkpoint() is called under outer torch.no_grad() but its context_fn or callable explicitly re-enables gradients internally, execution remains numerically correct, but the direct-forward path bypasses activation checkpointing and may retain more activations.

Callers that need checkpoint recomputation for such a gradient-enabled region should enable gradients around the checkpoint call itself:

with torch.no_grad():
    # Evaluation work...
    with torch.enable_grad():
        output = te.checkpoint(function, input, use_reentrant=False)

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 new public API)
• My changes generate no new warnings
• I have added tests that prove my fix is effective
• New and existing relevant unit tests pass with my changes

[greptile-apps]

Greptile Summary

The PR fixes delayed-scaling FP8 metadata pairing for checkpointed modules whose training mode differs between the original and recompute forwards.

• Stashes recompute metadata for every delayed-scaling FP8 module participating in checkpoint phase one.
• Bypasses checkpoint bookkeeping when autograd is disabled while preserving the requested forward context.
• Adds coverage for eval modules, mode transitions, reentrant and non-reentrant checkpointing, FIFO drainage, and nested explicit gradient contexts.

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	Broadens phase-one FP8 metadata stashing to include eval modules so recompute consumption remains paired across module mode changes.
transformer_engine/pytorch/distributed.py	Directly executes checkpointed functions when autograd is disabled, preserving the forward context and avoiding unreachable recompute state.
tests/pytorch/test_numerics.py	Adds focused regression coverage for checkpoint mode changes, eval modules, no-grad execution, explicit inner gradients, and recompute FIFO drainage.

Sequence Diagram

sequenceDiagram
  participant User
  participant Checkpoint as TE checkpoint
  participant Module as FP8 module
  participant Buffer as Recompute metadata FIFO
  User->>Checkpoint: checkpointed forward
  Checkpoint->>Module: phase 1 forward
  Module->>Buffer: stash scale and amax metadata
  User->>Checkpoint: backward
  Checkpoint->>Module: phase 2 recompute
  Module->>Buffer: consume matching metadata
  Buffer-->>Module: original scale and amax
  Module->>Module: recompute forward
  Module->>Module: restore live metadata

Loading

_{Reviews (3): Last reviewed commit: "Test explicit grad inside no-grad checkp..." | 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?

[nvegesna-netizen]

Thank you for the careful review. I agree with the core conclusion.

I could not identify a supported execution path where a delayed-scaling module legitimately appears in checkpoint phase 2 without having appeared and stashed in phase 1. Such a path would require divergent checkpoint execution, FP8 state replacement, or another invariant violation, and it should fail rather than silently recompute using live metadata.

I revised the PR accordingly:

• removed fp8_recompute_stashes and fp8_recompute_meta_restored;
• removed the set_extra_state() counter reset;
• restored strict, unconditional phase-2 FIFO consumption for delayed scaling;
• retained only the checkpoint-phase condition in the module:

if is_fp8_activation_recompute_enabled():
    FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta)

Testing that minimal change independently confirmed your analysis: all ten original eval/mode-change cases pass without the additional module state.

That discriminator also exposed a separate phase-1-without-phase-2 case. If te.checkpoint() is entered while outer autograd is disabled, phase 1 previously saved FP8 recompute snapshots even though no backward/recompute could consume them. Repeated forwards then accumulated unreachable snapshots. This already affected train-mode modules and removing the training guard would extend it to eval-mode modules.

I addressed that at the checkpoint boundary rather than in the module. For TE-containing callables, te.checkpoint() now executes the function directly under the supplied forward context when checkpoint-entry grad mode is disabled. This is the only location where the caller's actual grad state is still available; checking torch.is_grad_enabled() inside a module would be incorrect because a valid reentrant checkpoint deliberately executes its original forward under no_grad(). Non-TE callables still route through native PyTorch checkpointing.

The revised validation covers 18 focused cases across:

• reentrant and non-reentrant checkpointing;
• train and eval modules;
• train/eval mode changes between forward and recompute;
• repeated no-backward forwards and FIFO drainage;
• explicit nested gradient enablement through either context_fn or the callable.

For the explicit-gradient cases, output, input gradient, and weight gradient match direct execution, and the recompute FIFO remains empty. The broader recompute/checkpoint selection completes with 450 passes and 90 expected capability skips.

One boundary is now documented in the PR description: if checkpoint is entered under outer torch.no_grad() but gradients are re-enabled only inside the context/callable, execution is numerically correct but takes the direct-forward path, so checkpoint memory savings do not apply. A caller that needs recomputation for that region should enable gradients around the te.checkpoint() call itself.

The current head is 966baaa2. Please let me know if you would prefer the checkpoint-boundary cleanup separated from the one-line module fix, but I kept them together because the no-backward leak is the direct edge introduced for eval modules by applying the otherwise-correct simplification.