Skip to content

[simplefsdp auto-bucketing] add ir node reorder helper function #158098

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 8 commits into
base: gh/ruisizhang123/5/base
Choose a base branch
from
Open

[simplefsdp auto-bucketing] add ir node reorder helper function #158098

wants to merge 8 commits into from
+277 −6

Conversation

ruisizhang123
Copy link
Contributor

@ruisizhang123 ruisizhang123 commented Jul 11, 2025 •
edited
Loading

This pr adds reordering to pre-fetch bucketed AG/RS nodes.

The logic is as follows:

In forward pass:

  • reorder ag_i before wait_i-1_ag

In backward pass:

  • reorder ag_i after wait_i-1_ag
  • reorder wait_i_rs before rs_i+1

Stack from ghstack (oldest at bottom):

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov

@ruisizhang123
Update
febda47 
[ghstack-poisoned]
ruisizhang123 added a commit that referenced this pull request Jul 11, 2025
@ruisizhang123
add ir node reorder helper function
cad5f29 
ghstack-source-id: 7323307
Pull-Request: #158098
@pytorch-bot PyTorch Bot
Copy link

pytorch-bot bot commented Jul 11, 2025 •
edited
Loading

🔗 Helpful Links

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

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:

❌ 3 New Failures, 2 Unrelated Failures

As of commit 87afeb9 with merge base 33c9b41 (image):

NEW FAILURES - The following jobs have failed:

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

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

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

This was referenced Jul 11, 2025
Draft
Open
@github-actions GitHub Actions
Copy link
Contributor

This PR needs a release notes: label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@ruisizhang123 ruisizhang123 changed the title add ir node reorder helper function [simplefsdp auto-bucketing] add ir node reorder helper function Jul 11, 2025
@ruisizhang123 ruisizhang123 reopened this Jul 11, 2025
@ruisizhang123
Update
a2238db 
[ghstack-poisoned]
@ruisizhang123
Update
bbe01c2 
[ghstack-poisoned]
@ruisizhang123
Update
ff2ad43 
[ghstack-poisoned]
@ruisizhang123
Update
e269db6 
[ghstack-poisoned]
ruisizhang123 added a commit that referenced this pull request Jul 12, 2025
@ruisizhang123
add ir node reorder helper function
079e810 
ghstack-source-id: 2cd4c28
Pull-Request: #158098
@ruisizhang123
Update
3b7bac1 
[ghstack-poisoned]
@ruisizhang123 ruisizhang123 mentioned this pull request Jul 15, 2025
Open
@ruisizhang123
Update
a2e3b3d 
[ghstack-poisoned]
@ruisizhang123 ruisizhang123 mentioned this pull request Jul 17, 2025
Open
@ruisizhang123
Update
87afeb9 
[ghstack-poisoned]
@ruisizhang123 ruisizhang123 mentioned this pull request Aug 10, 2025
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers
No reviews
Assignees
No one assigned
Labels
ciflow/inductor module: inductor
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@ruisizhang123