[shjwudp]

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

What does this PR do?

Add iteration-scoped MFSDP v2 demand recording and cursor-based replay for VPP chunk revisits and mixed forward/backward execution in combined 1F1B. The iteration scope owns recorded cursor resets; static F/B cursors reset in root hooks outside a scope. No explicit prefetch calls are added to the compute plan. This remains a work-in-progress: GPU replay has a synchronization stall described below.

Depends on #8207. The branch includes its unmerged commit; review the incremental implementation.

• A successful first iteration records (weakref(owner), phase) occurrences in an ExecutionOrder, including repeated materialized owners. Prefetch gathers never enter the recording.
• Replay reuses ExecutionOrder.advance(), next_item(offset), and reset() for execution progress. IterationOrder owns a separate monotonic prefetch position and the recording/cache lifecycle. The execution cursor validates each observed demand; the prefetch cursor advances monotonically, independently of buffer materialization. Missing, extra, or changed demand invalidates the cached order; failed recordings are discarded.
• VPP and combined 1F1B entry points add an iteration-scope decorator. Ordinary non-VPP F/B schedules retain static prefetch without recording. Nested combined calls join the outer scope, and training/validation and scheduling configurations use separate cache keys.
• Existing demand hooks retain gather waits; existing release and parameter-completion hooks retain stream and gradient-readiness ownership, including delayed wgrad.
• Static F/B prefetch remains available without recording. ExecutionOrder removes _unique_index and the item argument to advance(), with an explicit reset() instead of automatic wrapping. The caller validates each returned occurrence. ExecutionOrder stores generic payloads directly; static F/B callers supply weak module references, and iteration callers supply weak module/phase pairs. next_item(offset) looks ahead from the last consumed occurrence without advancing execution. Repeated calls to the same module retain distinct successors. Static prefetch size accounting, budgets, and all-gather calls remain in FsdpModule; ExecutionOrder has no dependency on that class. Outside a scope, only the first root in each static F/B order resets that cursor; later roots continue it. This preserves unscoped execution of independent roots sharing a context. Comments explain why resetting every root would restart the cursor mid-pass. Recomputation skips the forward reset, and validation leaves backward progress untouched.
• Inside an iteration scope, root hooks never reset static or recorded cursors. The scope caches ExecutionOrder instances containing paired weak module references and phase metadata. Execution progress is delegated to that order; strict replay validation and the separate prefetch position remain in IterationOrder. No static-order wrapper/cache is introduced.

Validation

Not ready for review: 8-H100 combined VPP replay stalls in the eager reference on the second iteration. With default batched P2P, the stack stops in device-wide synchronization; disabling that synchronization moves the stall to MoE token dispatch stream synchronization. Separate cursors do not resolve speculative collective ordering across communication boundaries.

Historical validation on 6ae867597 (before removing test_pipeline_schedulers.py): 44 focused tests on 8 H100 GPUs: static cursor/prefetch, iteration/checkpoint, context, SGD across sharding strategies and accumulation counts, mixed-precision Adam, optimizer synchronization, sibling-root NVTX, checkpoint roundtrips and packed chunks, and combined 1F1B numerical tests (delayed wgrad on/off). Both checkpoint variants cover ordinary and recorded execution. Reset spies verify five forward resets and three backward resets for unscoped training/validation, and zero root-hook resets inside iteration scopes. The original unscoped shared-root fixtures are restored, including three consecutive forward passes in the shared-context test. Static cases assert prefetch-before-demand, including on the first step. The two VPP tests were excluded from that passing selection. The scheduler test file, including its two combined 1F1B cases, has since been removed from this PR.

Profiler limitation: the two selected existing overlap/budget tests fail in the unchanged profiler helper because this container's FunctionEvent lacks linked_correlation_id. No overlap timing result is claimed. Repository-wide CI remains unvalidated.

Eight source-level order/recorder checks pass, including paired weak-reference payloads, repeated occurrences, budget handling, and prefetch progressing independently of execution. Four GPU checkpoint cases are excluded from this local harness and covered by the GPU selection above. Ruff, isort, formatting, and whitespace checks pass. Local formatting used Black 24.4.2; repository-wide formatting/CI remains pending.

Limits

Requires a stable compatible schedule across DP peers. A local divergence exception is not a distributed recovery protocol. Iteration recording is disabled for CUDA graph schedules; graph recording/replay lifecycle support and throughput/peak-memory validation are deferred.

• Relevant unit and numerical integration tests added
• Repository-wide autoformat/CI completed

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