[shjwudp]

Summary

Record actual MFSDP v2 parameter demands over a complete iteration, preserving repeated module calls and mixed forward/backward phases. This builds on the merged ExecutionOrder change (#8207) and brings the iteration-scope implementation from #8284 into this PR.

Changes

• Keep ExecutionOrder generic: callers store weak module references or (weakref(module), phase) pairs. Traversal uses advance() and an explicit reset, without module lookup or _unique_index.
• Add iteration(context, key=..., rerecord=...). The first successful scope records every demand; subsequent scopes validate replay and prefetch through a separate monotonic cursor. Failed or divergent iterations invalidate their cached order.
• Preserve static forward/backward prefetch without a recording scope. Root hooks reset the static cursors at the first root; active iteration scopes own their cursor lifecycle.
• Scope the VPP and combined 1F1B entry points, and issue phase-aware demands from the combined hooks.
• Retain --megatron-fsdp-record-prefetch-order for other eager training schedules. Training and validation share the same scope wrapper, with separate cache keys for their schedule configurations.
• Record repeated backward invocations while opening a shared module's backward phase and NVTX range once, preserving the current gradient-completion lifecycle.

Usage and limits

For ordinary eager training, enable --use-megatron-fsdp --megatron-fsdp-version 2 --megatron-fsdp-record-prefetch-order --cuda-graph-impl none. Static forward/backward prefetch continues to work without this flag. Eager VPP and combined 1F1B enter the scope at their schedule boundaries.

Recorded schedules must be repeatable and collective-compatible across peers. Explicit scope keys or rerecord=True handle changed execution patterns. CUDA graph recording/replay is outside this change.

Work in progress: an earlier combined VPP experiment stalled during replay, including in its eager reference. That issue has not been resolved or revalidated on this revision; the presence of iteration scopes does not establish combined VPP correctness.

Validation

• Eight local source-level checks passed for occurrence traversal, prefetch budgets, mixed-phase recording/replay, invalidation, and weak ownership. These checks isolate the order logic and mock communication; they are not distributed GPU tests.
• Import sorting, Black 26.3.0 (core/tests), Ruff, syntax parsing, and whitespace checks passed.
• Preserved the distributed shared-projection regression comparing outputs, gradients, and SGD updates across recording and replay. Added iteration/checkpoint coverage and retained existing context regressions.
• GPU validation of this revision is pending. The source transfer to CW timed out, after which both CW and EOS login hostnames stopped resolving. No GPU allocation was created.

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

[wujingyue]

Hey @shjwudp,

It's amazing to see how quickly you can prototype it on V2!

If you want this to be more than a prototype, I'd recommend have your first PR targeting just recording the prefetch module order. That should be easier and the PR should be much simpler, and we know recording module orders is needed right now. Then, on top of that, you can draft-add the "compile" and "replay", and verify performance benefit on your customer models.

How does that sound to you?

[shjwudp]

• Add --megatron-fsdp-record-prefetch-order to record the first executed training iteration, including all microbatches. Shared contexts across model chunks are entered once.

Hi @wujingyue , I’ve updated the code to only record the module execution order during the forward and backward passes. What do you think of this approach? It still needs testing, but this is the general direction I’m planning to take.

[wujingyue]

Thanks for the PR!

[wujingyue]

Inlined

[wujingyue]

Thanks for the PR!

Several changes are a bit controversial at this very moment. Can we extract the non-controversial part into a separate PR so it's not helped up in review for too long? https://artsy.github.io/blog/2021/03/09/strategies-for-small-focused-pull-requests/

For example, I'd first target the standard (non-PP, non-VPP, non-1f1b) use case where static order mismatches execution.

[wujingyue]

This appears to be an over-engineering that agents are "great" at :)

I don't think inheritance is the right design pattern. We could store the states in FsdpContext or if that becomes too complicated wrap them in a class with ExecutionOrder as an attribute.

[shjwudp]

For the "store the states in FsdpContext". I’m wondering whether FsdpContext grow so much state could itself be a concern.

My intention was to introduce an IterationOrder abstraction, so I had the agent draft this class. Are you suggesting that IterationOrder should encapsulate an ExecutionOrder instance instead?

[wujingyue]

It's OK to keep IterationOrder, but I'd prefer composition over inheritance. Something like

class IterationOrder:
  execution_order: ExecutionOrder
  cursors

[shjwudp]

Hi Jingyue,

In the latest implementation, I’ve moved it into the context. Does this placement look reasonable to you? Relevant changes

I’ve also merged forward_order and backward_order into execution_order. The IterationOrder implementation is now split into two parts: #8485 is the first part, and #8486 is the second.

[wujingyue]

Make an enum

[wujingyue]

What's "demand" and should it be a separate PR so this PR stays small and focused? https://google.github.io/eng-practices/review/developer/small-cls.html#what-is-small

[shjwudp]

To clarify, demand mode disables prefetching entirely. This differs from static, where the module execution order is specified at initialization, and record_and_replay, where the order is captured at runtime and replayed for prefetching. It serves as a fallback when prefetching is not used.

Currently, this mode is intended for validation and forward-only execution.

[wujingyue]

Got it. We have arguments for prefetch sizes. Can they be used instead to disable prefetching?

[wujingyue]

I understand the English but I don't understand the problem it tries to solve. Again, a common problem of agent-generated code comments :)

I'd let the agent start with the simple alternative of letting one context stay in one mode, and then try to prove the complication is needed on top of that.

[shjwudp]

Thanks for the feedback!

The comment was intended to explain the logic in the following three lines code: when using static prefetching with VPP, we need to call reset_cursor. Since VPP involves multiple model chunks, we need to identify the root module of the first VPP model chunk and call reset_cursor there.

[wujingyue]

Again, I understand the English but not the reason behind it.

I got your earlier point that recording needs a stop signal, but this iteration scope seems to do something else that I've yet to understand the motivation of.

[shjwudp]

Currently, there are two modes require stop signals at different boundaries.

In module.py, the signal is used for static prefetching. In mfsdp_iteration.py, I define a separate iteration scope that spans multiple module hook invocations, with the stop signal tied to the end of that iteration scope.