Fix FSDP activation recompute prefetch

[shjwudp]

What

• Track whether the experimental FSDP root context is executing backward.
• Disable forward-order all-gather prefetch during activation recomputation.
• Keep recomputed parameters materialized until their matching post-backward reshard.
• Add regression coverage for both reentrant and non-reentrant activation checkpointing.

Bug

Activation recomputation executes module forward hooks while autograd is already running backward. Megatron FSDP v2 did not distinguish these recomputed forwards from the original forward pass, so FsdpModule.pre_forward() continued using the static forward order to prefetch the next module.

For two consecutive FSDP modules, the failure sequence is:

1. The later module is recomputed, completes backward, and is correctly resharded.
2. Autograd then recomputes the earlier module.
3. The earlier module's forward prefetch unshards the later module again.
4. The later module's backward has already completed, so no subsequent hook reshards it.
5. Backward therefore finishes with the later module exposing unsharded parameters.

In the downstream optimizer path, this can leave stale gradients on empty local DTensor shards. The next zero_grad() clears nonempty optimizer shards but cannot reach those empty sharded parameter objects while the model exposes its unsharded parameters. Gradient reduction then observes a mixture of set and missing gradients and raises:

FSDP sharded gradients must be either all set or all None.

Fix

The root FSDP context now enters backward state from its root pre-backward hook until the autograd final callback runs.

While in this state:

• recomputed forwards do not issue forward-order prefetch;
• recomputed parameters remain materialized for their immediately following backward, avoiding an unnecessary second all-gather;
• post_backward() performs the matching gradient reduction and reshard;
• the autograd final callback clears the backward state so normal prefetch resumes on the next forward.

Regression test

The test activation-checkpoints two independently sharded linear layers and is parameterized over use_reentrant=True and use_reentrant=False.

Without the fix, both modes leave the second layer as an unsharded Parameter after backward. With the fix, both layers return to sharded DTensor parameters and the root context exits backward state.

Tests

• Unpatched main, focused 4-GPU regression: 2 failed as expected (reentrant and non-reentrant).
• Patched focused 4-GPU regression: 2 passed.
• Final non-profiler test_fully_shard.py suite: 22 passed, 3 deselected.
• Downstream ND-parallel suite using the exact 4-GPU command: 4 passed, 5 skipped.
• Black, isort, Ruff, Pylint, and MyPy passed on the changed files.

The complete test_fully_shard.py run reaches an unrelated test-container profiler incompatibility: FunctionEvent.linked_correlation_id is unavailable. The activation-recompute regression and all non-profiler cases pass.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.