Unify preset FSDP execution order traversal by shjwudp · Pull Request #8485 · NVIDIA/Megatron-LM

shjwudp · 2026-10-26T01:07:38Z

Summary

Replace separate forward and backward execution orders with one weakly held sequence of ExecutionEvent(module, phase) values and one cursor. Registration mode populates the sequence during construction; recorded mode populates it from execution. Both use the same traversal.

Implementation

Validate each observed module and phase against the event at the cursor, raising on divergence.
Use argument-free advance(), which automatically wraps at the end of the sequence.
Prefetch successive events using each event's phase and the existing prefetch budget. Forward/backward phase changes do not stop lookahead.
Preserve every recorded occurrence, including repeated calls to the same module.
Test registration-order divergence with a model that executes its two layers in the opposite order from registration.
Use recording mode for the shared-MTP regression, whose shared block runs twice per forward pass; retain its five-step comparison with unsharded training.

The execution-order implementation follows #8721. This PR changes three production files and three test files. Complete schedule-window recording, PP/VPP and EP-overlap integration, and evaluation handling are covered separately in #8778.

Validation

Rebased onto main at bfab98f4a, preserving the extracted MCore adapters and the updated communication-stream policies.

Exact commit 158b2b6ac3f426dce63e1d4ec4006b3418533869: Docker validation on two RTX A6000 GPUs passed 93 tests on each rank, with five tests skipped because they require four ranks. Ran the complete execution-order, context, fully-shard, and pipeline-context test files. Coverage includes shared-MTP training, registration-order divergence, recorded replay, activation recomputation, and HSDP/HFSDP loss comparisons.

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 FSDP implementation and tests, incorporating the simplification from #8721. Existing copyright and license notices are preserved. No code copied from other repositories.

copy-pr-bot · 2026-10-26T01:07:42Z

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

ExecutionOrder can hold phase-tagged items so preset forward and backward prefetch share one sequence. Preserve unique-item lookup and repeated-occurrence validation, replace next_item with indexed traversal, and keep prefetch within each preset phase.

Can you let your agent rewrite this? I've yet to figure out why self.forward_order and self.backward_order should be unified. Even considering just static registration, the post order of a module tree already deviates from its reverse pre order.

shjwudp · 2026-10-27T11:32:21Z

ExecutionOrder can hold phase-tagged items so preset forward and backward prefetch share one sequence. Preserve unique-item lookup and repeated-occurrence validation, replace next_item with indexed traversal, and keep prefetch within each preset phase.

Can you let your agent rewrite this? I've yet to figure out why self.forward_order and self.backward_order should be unified. Even considering just static registration, the post order of a module tree already deviates from its reverse pre order.

Hi Jingyue, apologies that the PR description wasn't clear before—I’ve updated it.

wujingyue

LGTM overall. Thanks for clarifying the motivation!

wujingyue · 2026-10-28T03:22:06Z

+        return len(self._items)
+
+    @property
+    def position(self) -> int:


I don't see a good reason to expose this. We pretty much use the order as a linked list instead.

The position allows us to maintain an unshardable interval within the context.

This addresses the issue with the two cursors: the unshard cursor and the prefetch cursor. A single FsdpModule.unshard call may trigger prefetching for zero or more modules, but the unshard cursor advances only once per call.

wujingyue · 2026-10-28T03:22:13Z

+        """Index of the next occurrence; len(order) means traversal is exhausted."""
+        return self._position
+
+    def __getitem__(self, index: int) -> ExecutionEvent:


wujingyue · 2026-10-28T03:25:06Z

+            raise RuntimeError("A module in the FSDP execution order was destroyed.")
+        return module, phase
+
+    def advance(self, item: ExecutionEvent | None = None) -> ExecutionEvent:


I remember asking about this in a previous PR: why item? Is it not yet the time to remove it?

Yes, we can remove it; a "reset cursor" has been added to the root module's pre-forward/backward hooks to ensure that we can advance from the correct state.

wujingyue · 2026-10-28T03:26:11Z

+        for index in range(len(self)):
+            yield self[index]
+
+    def reset_cursor(self, position: int = 0) -> None:


Suggested change

def reset_cursor(self, position: int = 0) -> None:

def reset_cursor(self) -> None:

Virtual Pipeline Parallelism need this position, as we have multiple "root" FsdpModules executing within a single execution_order.

Why does VPP need to reset the cursor to somewhere in the middle?

With VPP, I expect a context to cover multiple microbatches that are executed in an interleaved fashion. So I expect the context to capture the execution order of all, with each module appearing multiple times in it.

I think we don't need reset the cursor when the execution order is captured at runtime, but it is needed when the execution order is preset during initialization.

For example:

Preset: [L1-F, L2-F, L2-B, L1-B] Runtime: [L1-F, L1-F, L2-F, L2-F, L2-B, L2-B, L1-B, L1-B]

it is needed when the execution order is preset during initialization.

Got it. I'm OK with this temporarily, but should we eventually switch to only execution order for VPP? I don't see a good reason to use static order with VPP.

Yeah, I think we should enable record-and-replay by default for VPP—and, more generally, for any runtime-dependent execution schedules.

When record-and-replay isn’t enabled, falling back to no prefetching seems like a good approach.

wujingyue · 2026-10-29T06:28:34Z

        context = self.context
        current_stream = context.current_stream()
        if self.is_root():
+            context.reset_execution_order(cast(nn.Module, self), Phase.BACKWARD)


Why and why is it in this PR? The PR title suggests it's only refactoring. But changes like this don't seem to agree.

https://google.github.io/eng-practices/review/developer/small-cls.html#what-is-small

I can remove reset_execution_order, but for cases with multi-roots or where PP > 1, we will need to temporarily disable prefetching until the record-and-replay functionality is implemented.

The proposed change is in commit: 7231899.

Got it! I'll clean up this PR and merge it, resolving the conflicts with #8478 as well.

wujingyue

A question raised during my cleanup: given we consolidate forward and backward into one order, should prefetching stop at phase boundary? My first instinct says it shouldn't because forward and backward can interleave in PP and VPP.

Signed-off-by: Jack Chang <jianbinc@nvidia.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Use one automatically wrapping cursor for registration and recorded execution. Validate each observed event, prefetch successors using their phases, and simplify divergence coverage to a reordered model. Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Use recording mode for the shared-MTP regression because the shared block executes twice per forward pass. Preserve the existing five-step comparison with unsharded training. Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

nemo-automation-bot · 2026-11-02T04:31:15Z

🔄 Merge queue validation started!

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

nemo-automation-bot · 2026-11-02T08:34:23Z

🔄 Merge queue validation started!

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

shjwudp · 2026-11-02T08:52:29Z

A question raised during my cleanup: given we consolidate forward and backward into one order, should prefetching stop at phase boundary? My first instinct says it shouldn't because forward and backward can interleave in PP and VPP.

@wujingyue Yes, I agree—we shouldn't impose phase boundaries on prefetching. If execution_order spans multiple microbatches, prefetching should be able to cross those boundaries as well.

While comparing DSv4 Flash performance with ZeRO-3 (M-FSDP v2 + MXFP8) and ZeRO-1 (DistOpt + MXFP8), I found that the all-gather at the start of each microbatch and the reduce-scatter (RS) at the end account for a 10% gap in step time. Prefetching the next microbatch's module weights could help narrow this gap by unblock next microbatch executing.

shjwudp mentioned this pull request Oct 26, 2026

Capture and replay FSDP training execution order #8486

Draft

shjwudp marked this pull request as ready for review October 26, 2026 01:13

shjwudp requested a review from a team as a code owner October 26, 2026 01:13

shjwudp self-assigned this Oct 26, 2026

shjwudp added the module: megatron-fsdp label Oct 26, 2026

github-actions Bot added the complexity: medium label Oct 26, 2026

shjwudp mentioned this pull request Oct 26, 2026

Record and replay runtime prefetch order for MFSDP v2 #7998

Draft

wujingyue reviewed Oct 26, 2026

View reviewed changes

shjwudp requested a review from wujingyue October 27, 2026 11:32

wujingyue reviewed Oct 27, 2026

View reviewed changes

Comment thread megatron/core/distributed/fsdp/src/megatron_fsdp/experimental/execution_order.py Resolved

wujingyue mentioned this pull request Oct 28, 2026

Record MFSDP execution order for prefetch #8478

Merged

1 task done

shjwudp requested a review from wujingyue October 28, 2026 02:29

wujingyue reviewed Oct 28, 2026

View reviewed changes

shjwudp requested a review from wujingyue October 28, 2026 12:59

wujingyue reviewed Oct 29, 2026

View reviewed changes

ahmadki mentioned this pull request Oct 29, 2026

Fix stale MFSDP v2 weights after a mispredicted backward prefetch #8687

Open

4 of 7 tasks

shjwudp requested a review from wujingyue October 29, 2026 15:14

wujingyue assigned wujingyue and unassigned shjwudp Oct 30, 2026

wujingyue mentioned this pull request Oct 30, 2026

Refactor FSDP execution order to one cursor #8721

Draft

1 task

XiaodaNV mentioned this pull request Oct 30, 2026

Add SchedulePolicy.unshard_all_planes_in_backward for MXFP8 recomputation in MFSDP v2 #8432

Merged

1 task

wujingyue force-pushed the split/mfsdp-preset-order branch from a3c0982 to 53b8669 Compare November 1, 2026 03:58

copy-pr-bot Bot deployed to test November 1, 2026 03:59 Active

copy-pr-bot Bot deployed to public November 1, 2026 04:02 Active

wujingyue reviewed Nov 1, 2026

View reviewed changes

nemo-automation-bot Bot added the Approved All necessary approvals have been made label Nov 1, 2026

copy-pr-bot Bot deployed to test November 1, 2026 16:27 Active

copy-pr-bot Bot deployed to public November 1, 2026 16:29 Active

shjwudp and others added 10 commits November 1, 2026 18:44

Unify preset FSDP execution order traversal

847efa5

Signed-off-by: Jack Chang <jianbinc@nvidia.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Specialize ExecutionOrder for weakly held module-phase events

a7b0f27

Signed-off-by: Jack Chang <jianbinc@nvidia.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Reset execution cursors in root hooks and name execution events

b66e230

Signed-off-by: Jack Chang <jianbinc@nvidia.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Restore event-based advancement for shared module calls

4310c6d

Signed-off-by: Jack Chang <jianbinc@nvidia.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Use cursor-only FSDP execution order with shared MTP compatibility

b96c575

Signed-off-by: Jack Chang <jianbinc@nvidia.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Simplify execution cursor reset docstring

87a057d

Signed-off-by: Jack Chang <jianbinc@nvidia.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

Simplify preset FSDP traversal and gather pipeline parameters on demand

5941a2e

Signed-off-by: Jack Chang <jianbinc@nvidia.com> Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

fix: preserve recorded occurrences after execution-order rebase

1c506ef

Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

test: record repeated shared-MTP execution order

158b2b6

Use recording mode for the shared-MTP regression because the shared block executes twice per forward pass. Preserve the existing five-step comparison with unsharded training. Signed-off-by: Jingyue Wu <jingyuew@nvidia.com>

wujingyue force-pushed the split/mfsdp-preset-order branch from 18ee22d to 158b2b6 Compare November 2, 2026 02:53

copy-pr-bot Bot deployed to test November 2, 2026 02:55 Active

copy-pr-bot Bot deployed to public November 2, 2026 02:58 Active

wujingyue added this pull request to the merge queue Nov 2, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Nov 2, 2026

shjwudp added this pull request to the merge queue Nov 2, 2026

Merged via the queue into NVIDIA:main with commit 2fcef83 Nov 2, 2026
98 checks passed

shjwudp deleted the split/mfsdp-preset-order branch November 2, 2026 09:59

XiaodaNV mentioned this pull request Nov 2, 2026

Enable unshard_all_planes_in_backward in MFSDPv2's adaptor when necessary #8839

Open

1 task done

	def reset_cursor(self, position: int = 0) -> None:
	def reset_cursor(self) -> None:

Conversation

shjwudp commented Oct 26, 2026 • edited by wujingyue Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation

Validation

Code provenance

Uh oh!

copy-pr-bot Bot commented Oct 26, 2026

Uh oh!

wujingyue left a comment

Choose a reason for hiding this comment

Uh oh!

shjwudp commented Oct 27, 2026

Uh oh!

wujingyue left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wujingyue left a comment

Choose a reason for hiding this comment

Uh oh!

nemo-automation-bot Bot commented Nov 2, 2026

Uh oh!

Uh oh!

nemo-automation-bot Bot commented Nov 2, 2026

Uh oh!

shjwudp commented Nov 2, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

shjwudp commented Oct 26, 2026 •

edited by wujingyue

Loading