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 duplicate FP8 state updates during activation-checkpoint recomputation by replacing the shared class-level _is_first_fp8_module list with a per-checkpoint _ActivationRecomputeState dataclass, adding reserve_first_fp8_module to suppress ownership claims during the original forward, and adding a _recompute flag on autocast that blocks the forward reduce_and_update_fp8_tensors call during recompute. The monotonic merge in set_autocast_state additionally prevents a nested autocast exit from reviving already-consumed ownership.

• distributed.py: activation_recompute_forward now carries per-checkpoint shared state; the __enter__ try/except rolls back all global flags on any failure, and __exit__ restores is_first_fp8_module only for the recompute phase or on forward-phase exceptions.
• quantization.py: autocast_enter skips the is_first_fp8_module = True reset at depth-0 when _recompute=True; autocast_exit skips the forward FP8 update under the same flag; set_autocast_state uses a boolean AND merge to monotonically preserve consumed ownership.
• Module changes (linear.py, layernorm_linear.py, layernorm_mlp.py): the per-module save/restore/conditional-undo pattern for is_first_fp8_module is removed or narrowed, delegating ownership bookkeeping entirely to the context managers.

Confidence Score: 5/5

• Safe to merge; the ownership logic is correct, all rollback paths are covered, and the test suite exercises all major code paths including nested autocast, forward exceptions, and recompute failures.
• The refactoring is well-scoped, the new _ActivationRecomputeState dataclass cleanly encapsulates per-checkpoint ownership, and the _recompute flag on autocast neatly prevents the forward FP8 update from firing during recompute without touching any hot paths. The only open question is whether the retained is_recomputation restore in _LayerNormMLP can trigger multiple backward FP8 updates when multiple LayerNormMLP modules share a single recomputed segment — but this is a pre-existing pattern narrowed (not widened) by this PR, and all 49 audit-suite tests pass.
• transformer_engine/pytorch/module/layernorm_mlp.py — the is_recomputation save/restore survives the cleanup and is worth a clarifying comment about the invariant that makes it safe.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/distributed.py	Core change: replaces the old class-level list-based _is_first_fp8_module stack with a per-checkpoint _ActivationRecomputeState dataclass. Adds reserve_first_fp8_module to suppress ownership during original forward, validates state before recompute entry, and rolls back all global flags on exception. Logic is sound and edge cases (nested autocast, forward failure, recompute failure) are correctly handled.
transformer_engine/pytorch/quantization.py	Two focused additions: (1) monotonic merge in set_autocast_state prevents a nested-autocast exit from reviving an already-consumed is_first_fp8_module; (2) _recompute flag on autocast / autocast_enter / autocast_exit skips the is_first_fp8_module reset at depth-0 entry and the forward reduce_and_update_fp8_tensors call at depth-0 exit, so the recompute pass does not trigger a duplicate forward FP8 update.
transformer_engine/pytorch/module/linear.py	Removes the old save/restore/conditional-undo pattern in _check_fp8_reduce_and_update entirely. The context manager in distributed.py now owns the invariant, so Linear modules simply consume is_first_fp8_module on both forward and recompute, with the correct value guaranteed by the surrounding activation_recompute_forward context.
transformer_engine/pytorch/module/layernorm_linear.py	Removes the in_fp8_activation_recompute_phase import and its save/restore block around is_first_fp8_module. Unlike layernorm_mlp.py, _LayerNormLinear has no is_recomputation code path, so the removal is complete and clean.
transformer_engine/pytorch/module/layernorm_mlp.py	Removes in_fp8_activation_recompute_phase() from the restore condition (fixing duplicate backward updates under checkpoint recompute), but retains the is_recomputation branch that saves and restores is_first_fp8_module. The retained restore is intentional for the TE-internal non-reentrant recompute path, but creates a subtle asymmetry with _LayerNormLinear and _Linear.
tests/pytorch/test_reentrant_fp8_updates.py	New regression test covering seven (checkpoint_mode, segments, num_layers) combinations plus a numeric-equivalence check. Uses monkeypatching to count reduce_and_update_fp8_tensors calls and asserts exactly (1 forward, 1 backward) across all paths.
tests/pytorch/test_fp8_checkpoint_state_edges.py	New targeted edge-case tests for nested autocast, forward exceptions, recompute-enter validation failure, and recompute-body exceptions. The clean_fp8_state fixture asserting depth/phase on teardown provides a useful global sanity check after every test.

Sequence Diagram

sequenceDiagram
    participant OA as te.autocast (outer)
    participant ARF as activation_recompute_forward<br/>(forward, reserve=True)
    participant M1 as Module 1 (Linear)
    participant M2 as Module 2 (Linear)
    participant ARR as activation_recompute_forward<br/>(recompute, state shared)
    participant RA as autocast(_recompute=True)

    Note over OA: is_first_fp8_module = True
    OA->>ARF: "__enter__ → saves True in state, sets flag=False"
    ARF->>M1: "forward → is_first_fp8_module()=False → bwd_update=False"
    ARF->>M2: "forward → is_first_fp8_module()=False → bwd_update=False"
    ARF->>OA: "__exit__ → forward_completed=True (flag stays False)"
    OA->>OA: autocast_exit → reduce_and_update_fp8_tensors(fwd) ×1 ✓

    Note over ARR: backward / recompute phase
    ARR->>ARR: "__enter__ → restores is_first_fp8_module=True (from state)"
    ARR->>RA: "__enter__ (_recompute=True) → depth++ but no flag reset"
    RA->>M1: "recompute → is_first_fp8_module()=True → bwd_update=True"
    RA->>M2: "recompute → is_first_fp8_module()=False → bwd_update=False"
    RA->>ARR: "autocast __exit__ → monotonic merge keeps flag=False, NO fwd update"
    ARR->>ARR: __exit__ → restores _previous_is_first_fp8_module
    Note over M1: backward → reduce_and_update_fp8_tensors(bwd) ×1 ✓

Loading

_{Reviews (5): 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?

[pggPL]

@AlbertYang514 sorry for delayed response.

I started reviewing this PR and I think this issue may be solved together with other issue, I need one or few days for investigation and will reach out to you soon.

[AlbertYang514]

@pggPL
Thanks for taking a look. I tested the current NVIDIA/main (172bd93), which already contains #3284, using black-box regressions extracted from this PR without applying the #3213 production fix.
The original duplicate FP8 update ownership issue is still reproducible. The expected update count is forward=1, backward=1, but I see:
no checkpoint fwd=1 bwd=1
non-reentrant/single fwd=2 bwd=1
non-reentrant/per-layer fwd=4 bwd=1
reentrant/single/1 layer fwd=2 bwd=1
reentrant/single/3 layers fwd=2 bwd=3
reentrant/per-layer fwd=4 bwd=3
nested reentrant fwd=3 bwd=2
Losses and gradients remain finite; the failures are specifically duplicate ownership/update bookkeeping. I also still reproduce the nested-autocast ownership and exception/recompute rollback issues covered here.
For comparison, the #3284 inner-autocast/recompute metadata regression coverage passes on the same current main (6 passed), so #3284 appears to address the metadata stash/restore/recompute-region issue rather than the first-module update ownership issue here.
I'm also running an additional cross-check now: applying the #3213 production fix to the pre-#3284 main state and running #3284's regression coverage against it. I'll post the result once that check is complete.

[AlbertYang514]

@pggPL
I finished the cross-check against #3284 as well.
The result indicates that the two fixes address distinct regressions:
Current main with #3284: #3284 regression tests pass (6/6), while the #3213 ownership regressions still fail (6/8 fail, with duplicate forward/backward updates).
Pre-#3284 main with only the #3213 production patch applied: all six #3284 regression tests fail, primarily with unbalanced recompute metadata stash/restore and KeyError: 'global_fp8_buffer_pos_fwd_recompute'.
So #3284 fixes the inner-autocast recompute-region / metadata bookkeeping issue, while #3213 fixes the first-module ownership / duplicate update and rollback issue. Neither fix subsumes the other.
For the reverse test I used pre-#3284 SHA 8260f49, applied only the #3213 production diff (3969485), and did not apply #3284 production changes.

[AlbertYang514]

@pggPL
I’ve updated #3213 onto current main with the post-#3284-compatible implementation.
The PR is now mergeable and the focused validation is green:
#3213 ownership regressions: 8 passed
#3213 edge/rollback tests: 6 passed
#3284 regressions: 6 passed
nearby checkpoint tests: passed
All tested checkpoint configurations now perform exactly one forward and one backward FP8 update, while preserving #3284’s inner-autocast recompute-region behavior.
New head: aa415db