[wujingyue]

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

What does this PR do?

I'm aware of #8485 and will rebase this PR on top of it after it lands.

When MFSDP modules execute in a different order from their registration, construction-order prefetch can gather the wrong next layer. Choose fully_shard_context(FsdpContext(prefetch_order=PrefetchOrder.RECORDED)) to record the first ordinary eager forward/backward pass and use those orders for subsequent prefetching.

with fully_shard_context(FsdpContext(prefetch_order=PrefetchOrder.RECORDED)) as context:
    fully_shard(model, mesh=mesh, placements=placements)

loss = model(inputs).square().mean()
loss.backward()  # context.post_backward() ends recording.

Record forward entry in pre_forward() and backward entry in pre_backward(). Only modules participating in backward enter its trace. Checkpoint recomputation does not add another forward occurrence. Repeated calls to a shared layer are skipped while it is already in BACKWARD and its weights remain available. Recording and replay use the same existing guard, so backward calls A → A → B record A → B in this case.

PrefetchOrder.REGISTRATION remains the default; PrefetchOrder.RECORDED selects the observed order. This reuses the existing ExecutionOrder objects and prefetch budgets; parameters are gathered on demand while recording. Discard the context if recording fails. Subsequent passes must repeat the recorded call order consistently across data-parallel peers.

The final autograd callback ends recording in context.post_backward(), alongside its existing weight cleanup and gradient synchronization.

Validation

Rebased onto origin/main at 79ee5c98a, using #8479's explicit
FsdpContext construction API.

On two RTX A6000 GPUs with NGC 26.08, the context, execution-order, and fully-shard
suites passed 65 tests per rank, with 13 skips requiring four GPUs.
The context regression registers two layers in the opposite order from execution,
runs a forward/backward pass, and checks both recorded orders.

Black, isort, Pylint (10/10), Ruff, and git diff --check passed.
Canonical CI uses NGC 26.09 and is pending on the rebased head.

Issue tracking

Related to #7764 and #7998. Original runtime-prefetch work by @shjwudp in #7998; this implements the reduced scope independently.

Code provenance

No code copied from other repositories.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[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.

[XiaodaNV]

Following up on the discussion in #8432: this doesn't record recomputation yet, so re-entered recomputation still gathers MXFP8 rowwise planes on demand.

I will take that as a follow-up after this lands: record (module, FORWARD) in backward order when recomputing, prefetch each successor for its recorded phase, and advance/prefetch from the recomputation's pre_forward. Does that sound right to you?

[wujingyue]

I will take that as a follow-up after this lands: record (module, FORWARD) in backward order when recomputing, prefetch each successor for its recorded phase, and advance/prefetch from the recomputation's pre_forward. Does that sound right to you?

You'll probably also have to wait for #8485, which I plan to merge next.

[XiaodaNV]

I will take that as a follow-up after this lands: record (module, FORWARD) in backward order when recomputing, prefetch each successor for its recorded phase, and advance/prefetch from the recomputation's pre_forward. Does that sound right to you?

You'll probably also have to wait for #8485, which I plan to merge next.

Sounds good, I'll build it on top of #8485.

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/41241652389

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/41257787183