[PP] Support OVERLAP_F_B computation type #158978

H-Huang · 2025-07-23T21:08:47Z

Stack from ghstack (oldest at bottom):

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see #159591)

The IR looks like:

[0F0, 0F1, 0F2, 0F3, 0F4, 0F5, 0F6, 7F0, 7I0, 7W0, 7F1, 7I1, 7W1, 7F2, 7I2, 7W2, 7F3, (0F7;7B3)OVERLAP_F_B, (7F4;0B0)OVERLAP_F_B, (0F8;7B4)OVERLAP_F_B, (7F5;0B1)OVERLAP_F_B, (0F9;7B5)OVERLAP_F_B, (7F6;0B2)OVERLAP_F_B, 7B6, (7F7;0B3)OVERLAP_F_B, 7B7, (7F8;0B4)OVERLAP_F_B, 7B8, (7F9;0B5)OVERLAP_F_B, 7B9, 0I6, 0W6, 0I7, 0W7, 0I8, 0W8, 0I9, 0W9]
[1F0, 1F1, 1F2, 1F3, 1F4, 6F0, 1F5, 6F1, 6I0, 6W0, 6F2, 6I1, 6W1, 6F3, (1F6;6B2)OVERLAP_F_B, (6F4;1B0)OVERLAP_F_B, (1F7;6B3)OVERLAP_F_B, (6F5;1B1)OVERLAP_F_B, (1F8;6B4)OVERLAP_F_B, (6F6;1B2)OVERLAP_F_B, (1F9;6B5)OVERLAP_F_B, (6F7;1B3)OVERLAP_F_B, 6B6, (6F8;1B4)OVERLAP_F_B, 6B7, (6F9;1B5)OVERLAP_F_B, 6B8, 1B6, 6I9, 1I7, 6W9, 1I8, 1W7, 1I9, 1W8, 1W9]
[2F0, 2F1, 2F2, 5F0, 2F3, 5F1, 2F4, 5F2, 5I0, 5W0, 5F3, (2F5;5B1)OVERLAP_F_B, (5F4;2B0)OVERLAP_F_B, (2F6;5B2)OVERLAP_F_B, (5F5;2B1)OVERLAP_F_B, (2F7;5B3)OVERLAP_F_B, (5F6;2B2)OVERLAP_F_B, (2F8;5B4)OVERLAP_F_B, (5F7;2B3)OVERLAP_F_B, (2F9;5B5)OVERLAP_F_B, (5F8;2B4)OVERLAP_F_B, 5B6, (5F9;2B5)OVERLAP_F_B, 5B7, 2B6, 5B8, 2I7, 5I9, 2I8, 2W7, 2I9, 5W9, 2W8, 2W9]
[3F0, 4F0, 3F1, 4F1, 3F2, 4F2, 3F3, 4F3, 3F4, 4B0, (4F4;3B0)OVERLAP_F_B, (3F5;4B1)OVERLAP_F_B, (4F5;3B1)OVERLAP_F_B, (3F6;4B2)OVERLAP_F_B, (4F6;3B2)OVERLAP_F_B, (3F7;4B3)OVERLAP_F_B, (4F7;3B3)OVERLAP_F_B, (3F8;4B4)OVERLAP_F_B, (4F8;3B4)OVERLAP_F_B, (3F9;4B5)OVERLAP_F_B, (4F9;3B5)OVERLAP_F_B, 4B6, 3B6, 4B7, 3B7, 4I8, 3I8, 4I9, 3I9, 4W8, 3W8, 4W9, 3W9]

In this PR, the schedule execution will just treat the OVERLAP_F_B as two separate operations of F and B (so there is no actual overlap). The next step is to allow users to create a custom function to plug in what this operation does.

pytorch/torch/distributed/pipelining/schedules.py

Lines 1205 to 1216 in 8146290

    
           # TODO: For now we are splitting OVERLAP_F_B into replacing it to 
        
           # its forward and backward components 
        
           # We need to figure out how to do the communication 
        
           for rank in compute_actions: 
        
               new_actions: list[_Action] = [] 
        
               for action in compute_actions[rank]: 
        
                   if action is not None and action.sub_actions is not None: 
        
                       # Replace OVERLAP_F_B action with its sub_actions 
        
                       new_actions.extend(action.sub_actions) 
        
                   else: 
        
                       new_actions.append(action) 
        
               compute_actions[rank] = new_actions

cc @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta

[ghstack-poisoned]

pytorch-bot · 2025-07-23T21:08:52Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/158978

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Failing to create ec2 instances -> lots of queueing

✅ You can merge normally! (1 Unrelated Failure)

As of commit 1943aac with merge base e4b123b ():

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

pull / linux-jammy-py3_9-clang9-xla / test (xla, 1, 1, lf.linux.12xlarge, unstable) (gh) (#158876)
sccache: error: couldn't connect to server

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ghstack-source-id: 71a7b99 Pull Request resolved: #158978

torch/distributed/pipelining/_utils.py

torch/distributed/pipelining/schedules.py

torch/distributed/pipelining/_schedule_visualizer.py

wconstab · 2025-07-24T21:07:46Z

torch/distributed/pipelining/schedules.py

+            # Check if action has sub_actions
+            if action.sub_actions is not None:
+                # Process each sub_action instead of the main action
+                for sub_action in action.sub_actions:


is the idea here that you first call F for one stage/mb, and then call B for the other stage/mb?

If so, how do you ensure that they are able to overlap? I thought the point of OVERLAP_F_B action w~as to call 1 user function, where inside the user function they take care of properly overlapping things.

edit: sorry, missed that this is validate_schedule. i think this is a good way to do it.

inside the user function they take care of properly overlapping things.

Yep, that was the plan. The custom user function is not implemented in this PR yet

wconstab

i think this looks like its on the right track, but a couple things i'd like to see

examples of how the schedule would be authored (is the IR notation ergonomic? iiuc it has commas inside the op's text. Please add some parsing tests that also illustrate how this looks
how is this going to be connected to the executor/runtime? Without seeing this piece i'm a little hesitant to land the scaffolding part.

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see pytorch/torchtitan#1447) cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

ghstack-source-id: d615be5 Pull Request resolved: #158978

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see pytorch/torchtitan#1447) cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

ghstack-source-id: 37ef4da Pull Request resolved: #158978

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see pytorch/torchtitan#1447) cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

ghstack-source-id: 5100aa8 Pull Request resolved: #158978

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see pytorch/torchtitan#1447) cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

H-Huang · 2025-07-31T21:57:12Z

examples of how the schedule would be authored (is the IR notation ergonomic? iiuc it has commas inside the op's text. Please add some parsing tests that also illustrate how this looks

DualPipeV is generated through #159591, but yes if users want to create it strictly from the IR its not well documented. Current it looks like this:

[0F0, 0F1, 0F2, 0F3, 0F4, 0F5, 0F6, 7F0, 7I0, 7W0, 7F1, 7I1, 7W1, 7F2, 7I2, 7W2, 7F3, (0F7;7B3)OVERLAP_F_B, (7F4;0B0)OVERLAP_F_B, (0F8;7B4)OVERLAP_F_B, (7F5;0B1)OVERLAP_F_B, (0F9;7B5)OVERLAP_F_B, (7F6;0B2)OVERLAP_F_B, 7B6, (7F7;0B3)OVERLAP_F_B, 7B7, (7F8;0B4)OVERLAP_F_B, 7B8, (7F9;0B5)OVERLAP_F_B, 7B9, 0I6, 0W6, 0I7, 0W7, 0I8, 0W8, 0I9, 0W9]
[1F0, 1F1, 1F2, 1F3, 1F4, 6F0, 1F5, 6F1, 6I0, 6W0, 6F2, 6I1, 6W1, 6F3, (1F6;6B2)OVERLAP_F_B, (6F4;1B0)OVERLAP_F_B, (1F7;6B3)OVERLAP_F_B, (6F5;1B1)OVERLAP_F_B, (1F8;6B4)OVERLAP_F_B, (6F6;1B2)OVERLAP_F_B, (1F9;6B5)OVERLAP_F_B, (6F7;1B3)OVERLAP_F_B, 6B6, (6F8;1B4)OVERLAP_F_B, 6B7, (6F9;1B5)OVERLAP_F_B, 6B8, 1B6, 6I9, 1I7, 6W9, 1I8, 1W7, 1I9, 1W8, 1W9]
[2F0, 2F1, 2F2, 5F0, 2F3, 5F1, 2F4, 5F2, 5I0, 5W0, 5F3, (2F5;5B1)OVERLAP_F_B, (5F4;2B0)OVERLAP_F_B, (2F6;5B2)OVERLAP_F_B, (5F5;2B1)OVERLAP_F_B, (2F7;5B3)OVERLAP_F_B, (5F6;2B2)OVERLAP_F_B, (2F8;5B4)OVERLAP_F_B, (5F7;2B3)OVERLAP_F_B, (2F9;5B5)OVERLAP_F_B, (5F8;2B4)OVERLAP_F_B, 5B6, (5F9;2B5)OVERLAP_F_B, 5B7, 2B6, 5B8, 2I7, 5I9, 2I8, 2W7, 2I9, 5W9, 2W8, 2W9]
[3F0, 4F0, 3F1, 4F1, 3F2, 4F2, 3F3, 4F3, 3F4, 4B0, (4F4;3B0)OVERLAP_F_B, (3F5;4B1)OVERLAP_F_B, (4F5;3B1)OVERLAP_F_B, (3F6;4B2)OVERLAP_F_B, (4F6;3B2)OVERLAP_F_B, (3F7;4B3)OVERLAP_F_B, (4F7;3B3)OVERLAP_F_B, (3F8;4B4)OVERLAP_F_B, (4F8;3B4)OVERLAP_F_B, (3F9;4B5)OVERLAP_F_B, (4F9;3B5)OVERLAP_F_B, 4B6, 3B6, 4B7, 3B7, 4I8, 3I8, 4I9, 3I9, 4W8, 3W8, 4W9, 3W9]

I'm thinking of just having it be (op1;op2;...) but not sure if that's general enough yet. Still thinking about it

how is this going to be connected to the executor/runtime? Without seeing this piece i'm a little hesitant to land the scaffolding part.

Right now DualPipeV is using the runtime. The workaround for OVERLAP since we dont have it implemented is to treat it as separate Forward and backward ops. This is working and able to run, the next steps would be to register a custom function for this operation.

wconstab

thanks, please add your IR snippet to the PR desc for future reference but i think thats a usable format.

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see #159591) The IR looks like: ``` [0F0, 0F1, 0F2, 0F3, 0F4, 0F5, 0F6, 7F0, 7I0, 7W0, 7F1, 7I1, 7W1, 7F2, 7I2, 7W2, 7F3, (0F7;7B3)OVERLAP_F_B, (7F4;0B0)OVERLAP_F_B, (0F8;7B4)OVERLAP_F_B, (7F5;0B1)OVERLAP_F_B, (0F9;7B5)OVERLAP_F_B, (7F6;0B2)OVERLAP_F_B, 7B6, (7F7;0B3)OVERLAP_F_B, 7B7, (7F8;0B4)OVERLAP_F_B, 7B8, (7F9;0B5)OVERLAP_F_B, 7B9, 0I6, 0W6, 0I7, 0W7, 0I8, 0W8, 0I9, 0W9] [1F0, 1F1, 1F2, 1F3, 1F4, 6F0, 1F5, 6F1, 6I0, 6W0, 6F2, 6I1, 6W1, 6F3, (1F6;6B2)OVERLAP_F_B, (6F4;1B0)OVERLAP_F_B, (1F7;6B3)OVERLAP_F_B, (6F5;1B1)OVERLAP_F_B, (1F8;6B4)OVERLAP_F_B, (6F6;1B2)OVERLAP_F_B, (1F9;6B5)OVERLAP_F_B, (6F7;1B3)OVERLAP_F_B, 6B6, (6F8;1B4)OVERLAP_F_B, 6B7, (6F9;1B5)OVERLAP_F_B, 6B8, 1B6, 6I9, 1I7, 6W9, 1I8, 1W7, 1I9, 1W8, 1W9] [2F0, 2F1, 2F2, 5F0, 2F3, 5F1, 2F4, 5F2, 5I0, 5W0, 5F3, (2F5;5B1)OVERLAP_F_B, (5F4;2B0)OVERLAP_F_B, (2F6;5B2)OVERLAP_F_B, (5F5;2B1)OVERLAP_F_B, (2F7;5B3)OVERLAP_F_B, (5F6;2B2)OVERLAP_F_B, (2F8;5B4)OVERLAP_F_B, (5F7;2B3)OVERLAP_F_B, (2F9;5B5)OVERLAP_F_B, (5F8;2B4)OVERLAP_F_B, 5B6, (5F9;2B5)OVERLAP_F_B, 5B7, 2B6, 5B8, 2I7, 5I9, 2I8, 2W7, 2I9, 5W9, 2W8, 2W9] [3F0, 4F0, 3F1, 4F1, 3F2, 4F2, 3F3, 4F3, 3F4, 4B0, (4F4;3B0)OVERLAP_F_B, (3F5;4B1)OVERLAP_F_B, (4F5;3B1)OVERLAP_F_B, (3F6;4B2)OVERLAP_F_B, (4F6;3B2)OVERLAP_F_B, (3F7;4B3)OVERLAP_F_B, (4F7;3B3)OVERLAP_F_B, (3F8;4B4)OVERLAP_F_B, (4F8;3B4)OVERLAP_F_B, (3F9;4B5)OVERLAP_F_B, (4F9;3B5)OVERLAP_F_B, 4B6, 3B6, 4B7, 3B7, 4I8, 3I8, 4I9, 3I9, 4W8, 3W8, 4W9, 3W9] ``` In this PR, the schedule execution will just treat the OVERLAP_F_B as two separate operations of F and B (so there is no actual overlap). The next step is to allow users to create a custom function to plug in what this operation does. https://github.com/pytorch/pytorch/blob/814629043a0c31441bc3749204c97f1e24fa3462/torch/distributed/pipelining/schedules.py#L1205-L1216 cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

H-Huang · 2025-08-01T19:59:24Z

@pytorchbot merge

pytorchmergebot · 2025-08-01T20:01:07Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Support OVERLAP_F_B computation type

cebdce3

[ghstack-poisoned]

H-Huang mentioned this pull request Jul 23, 2025

[PP] Refactor test_schedule_multiproc #158780

Closed

pytorch-bot bot added the oncall: distributed Add this issue/PR to distributed oncall triage queue label Jul 23, 2025

H-Huang added a commit that referenced this pull request Jul 23, 2025

Support OVERLAP_F_B computation type

0bef3ea

ghstack-source-id: 71a7b99 Pull Request resolved: #158978

H-Huang added the module: pipelining Pipeline Parallelism label Jul 23, 2025

H-Huang changed the title ~~Support OVERLAP_F_B computation type~~ [PP] Support OVERLAP_F_B computation type Jul 23, 2025

H-Huang added the release notes: distributed (pipeline) release notes category label Jul 23, 2025

Skylion007 reviewed Jul 24, 2025

View reviewed changes

torch/distributed/pipelining/_utils.py Outdated Show resolved Hide resolved

torch/distributed/pipelining/schedules.py Outdated Show resolved Hide resolved

wconstab reviewed Jul 24, 2025

View reviewed changes

torch/distributed/pipelining/_schedule_visualizer.py Show resolved Hide resolved

wconstab reviewed Jul 24, 2025

View reviewed changes

Update on "[PP] Support OVERLAP_F_B computation type"

5bd9a7f

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see pytorch/torchtitan#1447) cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

H-Huang added a commit that referenced this pull request Jul 28, 2025

Support OVERLAP_F_B computation type

cdde262

ghstack-source-id: d615be5 Pull Request resolved: #158978

Update on "[PP] Support OVERLAP_F_B computation type"

160c946

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see pytorch/torchtitan#1447) cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

H-Huang added a commit that referenced this pull request Jul 31, 2025

Support OVERLAP_F_B computation type

2edb174

ghstack-source-id: 37ef4da Pull Request resolved: #158978

Update on "[PP] Support OVERLAP_F_B computation type"

4522ca7

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see pytorch/torchtitan#1447) cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

H-Huang added a commit that referenced this pull request Jul 31, 2025

Support OVERLAP_F_B computation type

b041317

ghstack-source-id: 5100aa8 Pull Request resolved: #158978

Update on "[PP] Support OVERLAP_F_B computation type"

8146290

Some changes to validation code and visualizer to support a new computation type that will be used in DualPipeV (see pytorch/torchtitan#1447) cc awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta [ghstack-poisoned]

H-Huang mentioned this pull request Jul 31, 2025

[PP] Add DualPipeV schedule #159591

Closed

H-Huang requested a review from kwen2501 July 31, 2025 21:36

wconstab approved these changes Jul 31, 2025

View reviewed changes

H-Huang added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 1, 2025

pytorchmergebot added the merging label Aug 1, 2025

pytorchmergebot added the Merged label Aug 1, 2025

pytorchmergebot closed this in 5e8b956 Aug 1, 2025

pytorchmergebot removed the merging label Aug 1, 2025

github-actions bot deleted the gh/H-Huang/197/head branch September 1, 2025 02:20

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[PP] Support OVERLAP_F_B computation type #158978

[PP] Support OVERLAP_F_B computation type #158978

Uh oh!

H-Huang commented Jul 23, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jul 23, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wconstab Jul 24, 2025 •

edited

Loading

Uh oh!

H-Huang Jul 31, 2025

Uh oh!

wconstab left a comment

Uh oh!

H-Huang commented Jul 31, 2025

Uh oh!

wconstab left a comment

Uh oh!

H-Huang commented Aug 1, 2025

Uh oh!

pytorchmergebot commented Aug 1, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	# TODO: For now we are splitting OVERLAP_F_B into replacing it to
	# its forward and backward components
	# We need to figure out how to do the communication
	for rank in compute_actions:
	new_actions: list[_Action] = []
	for action in compute_actions[rank]:
	if action is not None and action.sub_actions is not None:
	# Replace OVERLAP_F_B action with its sub_actions
	new_actions.extend(action.sub_actions)
	else:
	new_actions.append(action)
	compute_actions[rank] = new_actions

[PP] Support OVERLAP_F_B computation type #158978

[PP] Support OVERLAP_F_B computation type #158978

Uh oh!

Conversation

H-Huang commented Jul 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/158978

❗ 1 Active SEVs

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

wconstab Jul 24, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

H-Huang Jul 31, 2025

Choose a reason for hiding this comment

Uh oh!

wconstab left a comment

Choose a reason for hiding this comment

Uh oh!

H-Huang commented Jul 31, 2025

Uh oh!

wconstab left a comment

Choose a reason for hiding this comment

Uh oh!

H-Huang commented Aug 1, 2025

Uh oh!

pytorchmergebot commented Aug 1, 2025

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

H-Huang commented Jul 23, 2025 •

edited

Loading

pytorch-bot bot commented Jul 23, 2025 •

edited

Loading

wconstab Jul 24, 2025 •

edited

Loading