MFSDP v2: separate backward completion from gradient readiness #7764

Is your feature request related to a problem? Please describe.

MFSDP v2 currently combines three events in FsdpModule.post_backward(): releasing full weights, reducing gradients, and resetting the module phase. For trainable parameters, its custom register_post_backward_hook() is actually driven by parameter-gradient readiness, not PyTorch's module post-backward hook.

Shared-layer MTP exposes the distinction. Two prediction depths invoke the same FSDP block twice in one forward. One loss.backward() enters that wrapper twice, but each shared parameter's post-accumulate hook fires once after both contributions are combined. #7757 fixes the resulting BACKWARD -> BACKWARD assertion with an idempotent entry guard. This issue proposes separating the lifetimes instead of keeping the execution phase open until wgrads are ready.

Describe the solution you'd like

Introduce distinct callbacks with explicit responsibilities:

Do not assume a fixed ordering between post_backward and grads_ready: parameter gradients can become ready before the module post-hook. Weight storage and gradient storage must have independent lifetimes.

Two details must be resolved before this is safe:

1. PyTorch's post-hook is not always the end of weight consumption. When no inputs require gradients, it fires at output-gradient readiness. A local probe of forward(x) = P.square() * x, with trainable P and x.requires_grad=False, observed:

   pre_backward -> post_backward -> backward reads saved P -> grads_ready
   

   Releasing P at the raw post-hook is premature. Define and test a fallback completion point for this case, including modules with no owned trainable parameters. Also audit delayed wgrad consumers rather than assuming that every delayed operation needs, or does not need, full weights. See PyTorch's hook firing rules.

2. Static prefetch order cannot represent repeated MTP invocations. FsdpContext.finalize() builds a module order and IndexedOrder rejects duplicate module objects. For a repeated block containing two separately sharded units:

   Actual backward: B(depth 2) -> A(depth 2) -> B(depth 1) -> A(depth 1) -> decoder
   Static order:    B          -> A          -> decoder
   

   At A(depth 2), static prefetch gathers decoder instead of B(depth 1). If the earlier invocation released B's weights, B needs an on-demand gather. The current mandatory gather-and-wait in pre_backward() provides a correctness fallback; the wrong prediction loses overlap and extends the decoder's weight residency. Preserve consistent collective ordering across ranks.

Describe alternatives you've considered

• Retain shared weights until gradient completion, as the ordinary-reuse guard in Fix MFSDP v2 backward with shared-layer MTP #7757 permits. This avoids gathers between depths but holds weights longer.
• Release per invocation with conservative or disabled backward prefetch for repeated units; retain the mandatory on-demand gather as a first implementation.
• Make prefetch follow invocation order, using an explicit scheduler or a validated runtime trace. A list of unique module objects is insufficient. Handle changing execution order without relying on a stale trace.

Additional context

Implementation starting points, pinned to #7757 commit 62ca2f9943f03fd93a73023b483db7b00828dbb9:

• Hook registration, phase transitions, weight release, and gradient reduction in module.py.
• Static successor lookup and duplicate rejection in indexed_order.py.
• TinySharedMTP and its gradient-parity regression. It combines normalized hidden states with shifted token embeddings, reuses a sharded residual block at two depths, and sums next-token cross-entropy losses. The parent owns the embedding, projection, outer normalizations, and output head.

On that PR's branch, run:

uv run python -m torch.distributed.run --standalone --nproc-per-node=2 -m pytest -q \
  tests/unit_tests/distributed/mfsdp_v2/test_fully_shard.py::test_shared_mtp_backward_matches_baseline

The existing guard and test passed on two RTX A6000 GPUs with PyTorch 2.13.0a0+8145d630e8.nv26.06; the context/fully-shard suites reported 44 passed and 6 skips requiring four ranks. The proposed callback split and prefetch changes have not been implemented or validated. The no-grad-input hook trace above was measured separately in that environment.

Suggested completion criteria:

• Define the callback contract and safe weight-release fallback, allowing either callback order and preserving gradient buffers across repeated gathers/releases.
• Extend the MTP regression to two separately sharded units repeated across depths plus a preceding decoder unit, so the static-order mismatch is exercised. Check losses, input gradients, every parameter-gradient shard, and optimizer updates against an unsharded reference across multiple steps/microbatches.
• Cover no-grad inputs, delayed TE wgrad, and changing reuse counts. Preserve existing reentrant and non-reentrant recomputation tests. Shared modules with nested reentrant checkpointing are a separate extension tracked in Fix MFSDP v2 backward for repeated module calls #7750; ordinary MTP should not require doubled callback counts.
• Check all-gather targets/counts and weight residency, including cleanup of prefetched storage. Report step time and peak active/reserved memory against retaining shared weights, and explain any additional gathers.
• Keep the context-level stream join correct after splitting reduction from weight release; verify that all reductions finish before gradients are consumed.

Related: #7757 (small ordinary-reuse fix), #7750 (shared-module checkpoint lifetimes), #7584 (phase-specific MXFP8 weight materialization).