Pass to split all_gather prologue and reduce_scatter prologue from fsdp graph #201
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
IvanKobzarev
added a commit
that referenced
this pull request
Oct 10, 2025
stack-info: PR: #201, branch: IvanKobzarev/stack/9
IvanKobzarev
force-pushed
the
IvanKobzarev/stack/9
branch
from
0159d8a to
68e7956
Compare
meta-cla
bot
added
the
CLA Signed
IvanKobzarev
added a commit
that referenced
this pull request
Oct 10, 2025
stack-info: PR: #201, branch: IvanKobzarev/stack/9
IvanKobzarev
force-pushed
the
IvanKobzarev/stack/9
branch
from
68e7956 to
a6dd593
Compare
IvanKobzarev
added a commit
that referenced
this pull request
Oct 10, 2025
stack-info: PR: #201, branch: IvanKobzarev/stack/9
IvanKobzarev
force-pushed
the
IvanKobzarev/stack/9
branch
2 times, most recently
from
7eb06ad to
0571cef
Compare
IvanKobzarev
added a commit
that referenced
this pull request
Oct 10, 2025
stack-info: PR: #201, branch: IvanKobzarev/stack/9
IvanKobzarev
added a commit
that referenced
this pull request
Oct 10, 2025
stack-info: PR: #201, branch: IvanKobzarev/stack/9
IvanKobzarev
force-pushed
the
IvanKobzarev/stack/9
branch
from
0571cef to
415b736
Compare
fmassa
reviewed
Oct 14, 2025
Comment on lines
+26
to
+63
| if len(n.users) != 1: | ||
| break |
Contributor
There was a problem hiding this comment.
This assumes that the partitioner has been changed to that we don't recompute the all-gather collectives in the backward pass.
@sanketpurandare you'll need to keep this in mind for your PR
IvanKobzarev
added a commit
that referenced
this pull request
Oct 23, 2025
…dp graph stack-info: PR: #201, branch: IvanKobzarev/stack/9
IvanKobzarev
force-pushed
the
IvanKobzarev/stack/9
branch
from
415b736 to
c098c92
Compare
IvanKobzarev
changed the title
…dp graph stack-info: PR: #201, branch: IvanKobzarev/stack/9
IvanKobzarev
force-pushed
the
IvanKobzarev/stack/9
branch
from
c098c92 to
18ccbec
Compare
bdhirsh
reviewed
Oct 29, 2025
| g_ins = g.find_nodes(op="placeholder") | ||
| prefetch_g_outs_map = [] | ||
|
|
||
| for g_in in g_ins: |
Contributor
There was a problem hiding this comment.
one thing I found a bit confusing when running this locally is that we are currently moving "views of inputs" into the prefetch subgraph. In my local example:
full backward graph: P2013183321
prefetch subgraph: P2013183365
remaining subgraph: P2013183449
you can see that tangents_1 is not an input in the "remaining subgraph", which is surprising because we should not be performing any FSDP collectives directly on tangents_1. And it looks like this is because there is a view(tangents_1) in the main backward, that we end up moving into the prefetch subgraph (technically harmless but confusing).
xmfan
pushed a commit
that referenced
this pull request
Oct 29, 2025
…dp graph stack-info: PR: #201, branch: IvanKobzarev/stack/9
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stacked PRs:
Pass to split all_gather prologue and reduce_scatter prologue from fsdp graph