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

[nvegesna-netizen]

Description

Fix delayed-scaling FP8 metadata stash and 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 because a module may 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 inputs.
• Keep strict phase 2 FIFO restoration and the existing delayed-scaling end-of-forward restore invariant.
• Bypass TE checkpoint bookkeeping when checkpoint is entered with outer autograd disabled, while preserving the supplied forward context.
• Add focused regressions for eval pairing, mode transitions, both checkpoint implementations, no grad entry, FIFO drainage, and forward context preservation.

Validation

• Six focused cases cover reentrant and non-reentrant eval pairing, with the module either remaining in eval or switching to train before recompute, plus no grad entry and forward context preservation.
• The eval cases use an intermediate input and verify finite input and weight gradients, exact FIFO drainage, and restoration of the live delayed-scaling metadata.
• The no grad cases verify that no per-module recompute key is created and that a supplied forward context preserves input and weight gradients.
• Prior broader validation of the same production changes completed with 450 passes and 90 expected capability skips on FP8-capable GPU hardware.

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 pairs delayed-scaling FP8 metadata stashes with recomputation per module and avoids creating unconsumable recompute state when checkpointing begins with autograd disabled.

• Stashes delayed-scaling metadata for checkpointed FP8 modules regardless of training mode.
• Directly executes TE-containing checkpoint callables entered without autograd while preserving the supplied forward context.
• Adds coverage for eval-mode recomputation, mode changes, FIFO drainage, no-backward execution, and explicit gradient enablement.

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	Removes training mode as the metadata-pairing signal so every delayed-scaling FP8 module in checkpoint phase one creates a matching FIFO entry.
transformer_engine/pytorch/distributed.py	Runs TE-containing checkpoint callables directly when checkpoint entry has autograd disabled, preserving the supplied forward context and avoiding unreachable recompute state.
tests/pytorch/test_numerics.py	Adds focused FP8 checkpoint regressions covering eval modules, mode changes, no-autograd calls, buffer drainage, and explicit gradient enablement.

Sequence Diagram

sequenceDiagram
  participant Caller
  participant Checkpoint as te.checkpoint
  participant Module as FP8 Module
  participant Buffer as Recompute FIFO
  Caller->>Checkpoint: forward with autograd enabled
  Checkpoint->>Module: checkpoint phase 1
  Module->>Buffer: stash delayed-scaling metadata
  Caller->>Checkpoint: backward
  Checkpoint->>Module: recompute phase 2
  Module->>Buffer: restore oldest matching metadata
  Module->>Module: execute recompute forward
  Module->>Module: restore updated forward state

Loading

_{Reviews (4): Last reviewed commit: "Address FP8 recompute review feedback" | 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.

[pggPL]

overall looks good, left some comments about the tests and nits

[pggPL]

New line not needed

[nvegesna-netizen]

Removed.

[pggPL]

ai leftover

[nvegesna-netizen]

Removed.

[pggPL]

Do you think we need so much tests for one line fix? My agent says yes, but I'm sceptical about it.

[nvegesna-netizen]

Agreed. I reduced this to six focused cases covering eval pairing with and without a mode transition across both checkpoint paths, no grad entry, and forward context preservation.

[pggPL]

/te-ci pytorch

[nvegesna-netizen]

/te-ci pytorch

[pggPL]

/te-ci pytorch

[pggPL]

My agent says: this path can still create unreachable FP8 recompute stashes for reentrant checkpoints when grad mode is enabled but none of the autograd Function inputs require gradients. In that case this Function cannot receive backward, so phase 2 cannot happen; stale FIFO entries can later be consumed by a valid training recompute.

ctx.needs_input_grad is PyTorch autograd's decision for the inputs passed to .apply(), so this preserves the current reentrant semantics while disabling only bookkeeping that cannot be consumed.

The relevant line is outside the current PR diff, so GitHub cannot attach an inline suggestion to it. Suggested source-only change:

diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
@@ -373,7 +373,9 @@ class _CheckpointFunction(torch.autograd.Function):
         torch_gpu_amp_ctx, torch_cpu_amp_ctx = _get_active_autocast_contexts()
 
         with torch.no_grad(), forward_ctx:
-            with activation_recompute_forward(activation_recompute=True, recompute_phase=False):
+            with activation_recompute_forward(
+                activation_recompute=any(ctx.needs_input_grad), recompute_phase=False
+            ):
                 outputs = run_function(*args, **kwargs)