Fix duplicate FP8 state updates during activation checkpoint recomputation

[AlbertYang514]

Summary

This change fixes FP8 first-module ownership across activation-checkpoint forward and recompute. The checkpoint frame reserves ownership for the original forward and reuses the captured state during recompute, preventing duplicate FP8 update ownership while preserving the expected single update.

v2 edge-case handling

The revised candidate also covers three state-transition edges:

• nested autocast exit cannot revive first-module ownership that was consumed by the inner region;
• an exception in the original checkpointed forward restores ownership and invalidates the incomplete frame;
• recompute entry and body failures roll state back transactionally, allowing later checkpoint use in the same outer autocast.

Ownership merging is monotonic (available may become consumed, but consumed is not restored to available from an older snapshot). The nested-autocast path explicitly reacquires quantization state after restoring the outer autocast state rather than relying on object identity.

Validation

• Existing PR Fix duplicate FP8 state updates during activation checkpoint recomputation #3213 tests: 8/8 passed.
• Focused nested-autocast and exception-recovery tests: 6/6 passed.
• Full local checkpoint audit suite: 49/49 passed.
• CPU state-transition probes passed for all added edge cases.

Scope

This PR is limited to checkpoint/autocast state ownership and rollback. A residual SM120 RTC cast-transpose failure can still be reproduced with correct 1/1 ownership/update counts; its investigation and workaround belong to PR #3215 and are not part of this change.

[greptile-apps]

Greptile Summary

This PR fixes a real bug where DelayedScaling FP8 global state (is_first_fp8_module / forward amax update / backward reduction) was being advanced multiple times per logical microbatch when activation checkpointing was used. The root cause was that process-global FP8 state was not scoped to checkpoint frames: recompute autocasts created extra forward update boundaries, and per-module restoration of consumed ownership allowed every recomputed module to reclaim backward-update rights.

• Introduces _ActivationRecomputeState as a per-frame container that carries is_first_fp8_module across the original forward and its recompute, replacing the previous shared class-level list.
• Adds reserve_first_fp8_module=True to the reentrant no-grad forward to prevent any module from consuming backward-update ownership during the no-grad pass; adds _recompute=True to all recompute-phase autocasts to skip forward amax/scale advancement and ownership reset; removes the per-module in_fp8_activation_recompute_phase() restoration from linear.py and layernorm_linear.py, delegating that bookkeeping entirely to the context managers.
• Adds a monotonic AND-merge guard in autocast.__exit__ so restoring a nested autocast's saved configuration cannot revive ownership that was consumed while the nested scope was active.

Confidence Score: 5/5

The change is safe to merge: it fixes a well-documented bookkeeping bug without altering any tensor math, CUDA kernels, or distributed communication primitives.

The fix is carefully structured with validation before any global mutation, exception-safe rollback via try/finally throughout, and monotonic AND-merge to prevent ownership revival. The new test suite covers all six checkpoint configurations from the PR description plus five exception-path edge cases.

No files require special attention; the state transitions in distributed.py and quantization.py have been verified against all tested configurations.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed.py	Core fix: introduces _ActivationRecomputeState, per-frame state save/restore in activation_recompute_forward, reserve_first_fp8_module for reentrant no-grad pass, exception-safe rollback on failed original forward, and correct monotonic merge on recompute exit.
transformer_engine/pytorch/quantization.py	Adds _recompute flag to autocast_enter/autocast_exit to prevent ownership reset and spurious forward DelayedScaling update during recompute; adds monotonic AND-merge in autocast.exit to stop nested autocast restoration from reviving consumed ownership.
transformer_engine/pytorch/module/linear.py	Removes the in_fp8_activation_recompute_phase() restoration from _check_fp8_reduce_and_update; ownership management is now fully delegated to the activation_recompute_forward context manager.
transformer_engine/pytorch/module/layernorm_linear.py	Removes the in_fp8_activation_recompute_phase() guard from backward FP8 ownership check; activation_recompute_forward now manages the recompute-phase state.
transformer_engine/pytorch/module/layernorm_mlp.py	Removes only the in_fp8_activation_recompute_phase() condition from the combined check, correctly preserving the is_recomputation guard which handles a different internal MLP recomputation path.
tests/pytorch/test_reentrant_fp8_updates.py	New regression test covering all six checkpoint configurations; verifies exactly 1 forward and 1 backward FP8 update per autocast, finite/nonzero gradients, and clean global state after each run.
tests/pytorch/test_fp8_checkpoint_state_edges.py	New edge-case tests covering nested autocast ownership monotonicity, exception in original forward, failed recompute-enter state leak, and failed recompute body not poisoning a subsequent autocast scope.

Sequence Diagram

sequenceDiagram
    participant User as User training loop
    participant AC as te.autocast
    participant ARF as activation_recompute_forward
    participant RAC as autocast(_recompute=True)
    participant Mod as TE Module (forward)
    participant MB as Module backward

    User->>AC: "__enter__ depth 0 to 1, is_first=True"
    Note over ARF: Original forward (reentrant, no_grad)
    User->>ARF: "__enter__ reserve=True, state.is_first=True, global=False"
    User->>Mod: "forward() is_first()=False reserved, no bwd owner"
    User->>ARF: __exit__ success global stays False
    AC->>AC: "__exit__ depth 1 to 0, reduce_and_update(forward=True)"
    Note over ARF,RAC: Backward / recompute phase
    User->>ARF: "__enter__ recompute_phase=True, prev=False, global=True"
    User->>RAC: "__enter__ _recompute=True, depth 0 to 1, no reset"
    User->>Mod: "forward() is_first()=True, bwd_owner=True"
    User->>RAC: "__exit__ _recompute=True, no forward update"
    User->>ARF: "__exit__ merge prev(False) AND cur = False"
    User->>MB: "backward() reduce_and_update(forward=False)"

Loading

_{Reviews (3): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[AlbertYang514]

Hi @ksivaman and @ptrendx gentle ping when either of you has a chance.

This PR has been open for about a week and is currently mergeable. It fixes duplicate FP8 global state updates under activation checkpoint recomputation, with regression coverage for reentrant, non-reentrant, per-layer, and nested checkpoint paths.

Could either of you please take a look, or point me to the current owner of the activation-recompute / FP8 state-management path? I’m happy to adjust or split the change if preferred.

[AlbertYang514]

Hi @vthumbe1503 — I revised #3213 to cover nested-autocast ownership restoration and transactional recovery from original-forward, recompute-entry, and recompute-body exceptions.

The updated candidate passes the existing 8 tests, 6 focused edge tests, and all 49 cases in the local checkpoint audit suite. The residual SM120 RTC cast-transpose failure is still reproducible with correct 1/1 ownership/update counts and is being treated separately in #3215.

When convenient, could you please take another look at the revised ownership and rollback logic?