Skip to content

[simplefsdp auto-bucketing] manual bucketing with plan #158321

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 3 commits into
base: gh/ruisizhang123/6/base
Choose a base branch
from
Open

[simplefsdp auto-bucketing] manual bucketing with plan #158321

wants to merge 3 commits into from

Conversation

ruisizhang123
Copy link
Contributor

@ruisizhang123 ruisizhang123 commented Jul 15, 2025 •
edited
Loading

This manual plan takes a list of module definition as input, and bucket all-gather/reduce-scatter based on the definition.

An example would be:

  • The plan is set to betorch._inductor.config.simplefsdp.bucketing_plan= ["tok_embeddings", "layers.[0-31]", "norm+output"].
  • The inductor backend will group nodes from the same module as a bucket. Here, the bucket groups are tok_embeddings, layers.0, layers.1, etc.

Stack from ghstack (oldest at bottom):

@ruisizhang123
Update
d9c7b45 
[ghstack-poisoned]
@pytorch-bot PyTorch Bot
Copy link

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

🔗 Helpful Links

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

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

ruisizhang123 added a commit that referenced this pull request Jul 15, 2025
@ruisizhang123
[simplefsdp auto-bucketing] manual bucketing with plan
103d7db 
ghstack-source-id: b7dbc9b
Pull-Request: #158321
This was referenced Jul 15, 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 mentioned this pull request Jul 15, 2025
Open
@ruisizhang123
Update
d3e46fb 
[ghstack-poisoned]
@ruisizhang123 ruisizhang123 mentioned this pull request Jul 17, 2025
Open
@ruisizhang123
Update
af8c4ce 
[ghstack-poisoned]
@ruisizhang123 ruisizhang123 mentioned this pull request Jul 25, 2025
Closed
@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