[HaochenYuan]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

CheckpointFunction.backward() re-runs the checkpointed function from inside autograd's backward, i.e. outside the
fp8_autocast the forward ran under; it only restored the RNG state. A Transformer Engine module inside the region
therefore recomputes with FP8 disabled: with FP8 parameters TE dequantizes the weight and the wgrad-fusion binding fails (AttributeError: 'Tensor' object has no attribute 'main_grad'); with BF16 parameters the recompute silently runs in BF16 while the forward ran in FP8.

This affects the sites that call tensor_parallel.checkpoint directly: --recompute-modules core_attn (attention.py,
absorbed_mla.py) and --recompute-modules gdn (GatedDeltaNet / KimiDeltaAttention). The block-level sites route to
te_checkpoint under FP8, and CheckpointWithoutOutput already re-enters the FP8 state itself.

Fix

forward() records FP8GlobalStateManager.is_fp8_enabled() and the recipe and runs under activation_recompute_forward(recompute_phase=False); backward() re-enters fp8_autocast(enabled=True, fp8_recipe=<recorded>) with activation_recompute_forward(recompute_phase=True) for the recompute — the same handling as CheckpointWithoutOutput._recompute and TE's checkpoint. Without TE, or when the forward ran outside an FP8 autocast, nothing changes.

Tests

tests/unit_tests/tensor_parallel/test_random.py: the recompute must see the forward's FP8 state (and none when there was none); a TE Linear with FP8 parameters and fused wgrad accumulation inside tensor_parallel.checkpoint must match the un-checkpointed pass bit for bit. Both fail on dev as described above and pass with this change (torchrun --nproc-per-node 2 -m pytest tests/unit_tests/tensor_parallel/test_random.py -k "fp8 or test_checkpoint").

Validation (1 GB200 node, mxfp8, TE 2.20)

Case	dev @ 7b29912	this PR
DeepSeek-V3.2-style GPT (dsa attention variant), FP8 parameters, selective core_attn recompute	crashes in iteration 1 (main_grad AttributeError)	runs
Qwen3.5-shaped GDN model, gdn recompute vs the same run without recompute	diverges from iteration 2	identical to every digit over 6 iterations

Behaviour change

Any tensor_parallel.checkpoint region whose forward runs under FP8 now recomputes in FP8 instead of BF16. This matches te_checkpoint and CheckpointWithoutOutput; for FP8 + `recompute the gradients change (to
match the forward) and the FP8-parameter crash disappears. Delayed scaling reuses the forward's scaling factors through TE's recompute phases.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact @NVIDIA/mcore-oncall.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• If this PR adds or changes a GPU kernel (Triton, jit_fuser/torch.compile, CUDA extension, TE or external-library dispatch, or a scatter/index accumulation), I have added or updated its bit-exact determinism test and registered it in tests/unit_tests/determinism/kernels/manifest.py (guide)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[HaochenYuan]

/ok to test 913513f

[buptzyb]

/review model=claude mode=strict

[svcnvidia-nemo-ci]

Strict review passed: I found no significant issues. LGTM.

Findings: CRITICAL 0 · IMPORTANT 0 · SUGGESTION 0

What I checked

• The fix copies an existing pattern. CheckpointFunction.forward records whether FP8 is on and which recipe is active. It then runs under activation_recompute_forward(recompute_phase=False). backward re-enters fp8_autocast with that recipe and uses recompute_phase=True. This is the same handling, in the same context order, as CheckpointWithoutOutputFunction.forward and CheckpointWithoutOutput._recompute.

• The contexts are entered in the right order. In both phases fp8_ctx is entered before activation_recompute_forward, so the recompute flag sees FP8 as enabled.

• Non-FP8 paths are unchanged. bool(HAVE_TE and ...) stops early when TE is missing, and runs that are not under FP8 use nullcontext(). Behaviour without TE or without FP8 is therefore the same as before.

• Every new field is used. ctx.fp8 and ctx.fp8_recipe are read in backward, and fwd_ctx, fp8_ctx and recompute_ctx are each entered.

• Only the intended call sites change. These sites already switch to te_checkpoint under config.fp8 or config.fp4:

  • block/full recompute in recompute.py
  • MLP recompute in transformer_layer.py
  • MoE and shared-expert recompute in moe_layer.py (also L887-L901)

  The change therefore only reaches the sites that always call tensor_parallel.checkpoint, which are the ones the PR targets:

  • core_attn in attention.py (L579)
  • absorbed_mla.py
  • GDN in gdn.py

• No reachable nesting. TE resets its recompute flags when the context exits instead of restoring the previous values, so nesting one of these contexts inside another would clear the outer region's flags. I found no reachable path where that happens:

  • The selective sites need recompute_granularity == 'selective', so they can't run inside a full-recompute te_checkpoint.
  • The CheckpointWithoutOutput() nested inside the GDN region (the gdn_norm_out path) is built with fp8=False, so it never touches those flags.

• The tests match the change. The new tests in test_random.py cover three things:

  • the recompute sees the forward's FP8 state, with a no-FP8 control;
  • the FP8 state is restored when backward is called directly;
  • a TE Linear with FP8 parameters and fuse_wgrad_accumulation gives the same output, x.grad and main_grad with and without checkpointing, under both MXFP8 and delayed scaling.

  Without the fix, the FP8-parameter case fails with the main_grad AttributeError whatever the tolerances, so the test does guard against a regression.

I checked this by reading the source only; I did not run the tests.

Risk: low. The change is small and local. It mirrors existing in-repo handling, and it only changes behaviour for tensor_parallel.checkpoint regions whose forward ran under FP8. Those regions now recompute the way the forward ran, which is what te_checkpoint and CheckpointWithoutOutput already do.

Review generated by claude