[wujingyue]

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

What does this PR do?

Depends on #8721. Review the scheduling-only change in commit 53b34027d. The pull-request/8721 mirror is not available yet, so this draft targets main and includes its prerequisite commits. Both branches are rebased onto main at 88312648a; #8478 has merged.

A PP/VPP schedule issues several autograd backward calls through one shared FSDP context. Ending recording after the first backward captures only part of that schedule. With caller-managed gradient synchronization, finish recording at finish_grad_sync() after all microbatches, and check that subsequent replay has consumed the complete trace. Per-autograd deferred module cleanup remains in place.

Select recorded mode for PP/VPP and EP overlap, reject registration-mode ambient contexts for these schedules, and record compute events from the combined-1F1B hooks. Evaluation gathers on demand without consuming the training trace. Ordinary eager training retains registration mode. Both modes use the single cursor and lookahead from #8721, including lookahead across forward/backward phase changes.

Validation

Eos job 6418873, exact commit 53b34027d766b57f5c3f3c435fa30aea82022daa: 34 focused tests passed on each of eight H100 GPUs.

Coverage includes grouped/interleaved shared-context PP/VPP call sequences, recording across multiple backward calls, incomplete-window rejection, recorded replay, evaluation between training passes, combined-1F1B owner hooks, and activation recomputation. The interleaved VPP regression checks that the next module's parameters are already materialized across a phase change, using the default one-event prefetch. The tests drive shared-context wrapping, local call sequences, and hooks directly.

Black, isort, Pylint, Ruff, and git diff --check passed. The cached dev CI image uses NGC 26.08; the repository pins 26.09.

Code provenance

Refactors this repository's implementation and tests. No code copied from other repositories. Existing copyright and license notices are preserved.

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