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

For reentrant checkpointing entered with autograd enabled, phase 1 bookkeeping is skipped when no custom Function input needs gradients. An ineligible nested call inherits any enclosing checkpoint phase so outer recomputation remains paired.

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.
• Skip standalone reentrant bookkeeping when autograd cannot invoke the custom Function backward, while preserving an enclosing recompute phase.
• Add focused regressions for eval pairing, mode transitions, both checkpoint implementations, no grad entry, reentrant input eligibility, nested phase preservation, FIFO drainage, and forward context preservation.

Validation

• Eight 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, forward context preservation, standalone reentrant input eligibility, and nested phase 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.
• The reentrant eligibility cases verify that standalone ineligible calls leave no FIFO state and that nested ineligible calls preserve and drain the enclosing checkpoint phase.
• Focused H100 validation of the reentrant eligibility additions passed all eight cases with zero failures, errors, or skips on the hash-pinned rc8 backport.
• Prior broader validation of the eval pairing and no grad boundary changes completed with 450 passes and 90 expected capability skips on FP8-capable GPU hardware.

Grad mode boundary

Checkpoint entry grad mode determines whether TE checkpointing is used. For reentrant checkpointing, PyTorch autograd input eligibility also determines whether standalone recompute bookkeeping is enabled. 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 changes or remains in evaluation mode.

• Stashes recompute metadata for every delayed-scaling FP8 module participating in the initial checkpoint phase.
• Avoids unreachable recompute bookkeeping when checkpointing cannot produce a backward recomputation.
• Preserves an enclosing checkpoint phase for nested reentrant calls without gradient-eligible inputs.
• Adds focused tests for evaluation mode, mode transitions, nested checkpoints, disabled autograd, context preservation, FIFO drainage, and gradients.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains within the eligible follow-up review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed.py	Adjusts checkpoint boundaries to bypass bookkeeping without outer autograd and preserve enclosing recompute state for ineligible nested reentrant calls.
transformer_engine/pytorch/module/base.py	Removes the training-mode condition from delayed-scaling metadata stashing so phase-one and phase-two module executions remain paired.
tests/pytorch/test_numerics.py	Adds targeted FP8 checkpoint regressions covering evaluation mode, mode changes, no-grad entry, nested eligibility, context preservation, gradients, and FIFO drainage.

Sequence Diagram

sequenceDiagram
  participant C as te.checkpoint
  participant F as Original forward
  participant M as FP8 module
  participant B as Metadata FIFO
  participant R as Backward recompute

  C->>F: Enter checkpoint phase 1
  F->>M: Execute delayed-scaling FP8 forward
  M->>B: Stash pre-forward metadata
  C-->>R: Backward requests recomputation
  R->>M: Execute checkpoint phase 2
  M->>B: Consume matching FIFO entry
  B-->>M: Restore phase-1 metadata
  M->>M: Run recompute and restore live metadata

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix/fp8-recompu..." | 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]

/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)

[nvegesna-netizen]

/te-ci pytorch

[nvegesna-netizen]

Good catch. ctx.needs_input_grad is the correct ownership signal. I used nullcontext() for the ineligible case because activation_recompute_forward(False, False) would mask an enclosing checkpoint phase. I added standalone and nested regressions. The focused H100 run passed all eight cases.