[PT] support custom all_gather and reduce_scatter comms #155189
Closed
+287
−80
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
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/155189
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 1cacb20 with merge base f79689b ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/inductor
module: inductor
oncall: distributed
|
This pull request was exported from Phabricator. Differential Revision: D75714362 |
xunnanxu
force-pushed
the
export-D75714362
branch
2 times, most recently
from
2b428e6 to
569a4d6
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75714362 |
1 similar comment
|
This pull request was exported from Phabricator. Differential Revision: D75714362 |
xunnanxu
force-pushed
the
export-D75714362
branch
2 times, most recently
from
e806f3a to
b03004b
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75714362 |
xunnanxu
force-pushed
the
export-D75714362
branch
from
b03004b to
8fc9010
Compare
xunnanxu
force-pushed
the
export-D75714362
branch
from
8fc9010 to
8d1e2cf
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75714362 |
xunnanxu
force-pushed
the
export-D75714362
branch
from
8d1e2cf to
6a7b530
Compare
xunnanxu
force-pushed
the
export-D75714362
branch
from
6a7b530 to
37d962b
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75714362 |
xunnanxu
force-pushed
the
export-D75714362
branch
from
37d962b to
995e205
Compare
|
This pull request was exported from Phabricator. Differential Revision: D75714362 |
xunnanxu
force-pushed
the
export-D75714362
branch
from
94dc77f to
f9ebe84
Compare
xunnanxu
force-pushed
the
export-D75714362
branch
from
f9ebe84 to
a292323
Compare
Summary: Pull Request resolved: pytorch#155189 This change introduces 2 comm override APIs: `set_custom_all_gather` and `set_custom_reduce_scatter` to allow for custom behavior respectively. This allow users to control how the comm buffers are allocated and the exact comm implementation for flexibility. For details, see docstring in `Comm`, `BaseComm` in `_fsdp_api.py` Test Plan: CI Differential Revision: D75714362
|
This pull request was exported from Phabricator. Differential Revision: D75714362 |
xunnanxu
force-pushed
the
export-D75714362
branch
from
a292323 to
1cacb20
Compare
weifengpy
approved these changes
Jul 1, 2025
|
@pytorchbot merge |
Merge failedReason: This PR has internal changes and must be landed via Phabricator! Please try reimporting/rexporting the PR! Details for Dev Infra teamRaised by workflow job |
|
@pytorchbot merge (Initiating merge automatically since Phabricator Diff has merged) |
Merge startedYour 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 |
xunnanxu
added a commit
to xunnanxu/pytorch
that referenced
this pull request
Jul 2, 2025
…ther with (pytorch#157487) Summary: This is a follow up after the PR to add comm override support: pytorch#155189 The previous PR loosely checks the allocation mixin classes, which isn't really safe as the actual hook may still override the behavior. This may lead to unnecessary confusion for no good use case. So for now we just make the 2 sets of APIs largely incompatible: 1. setting custom comms after `set_allocate_memory_from_process_group_for_comm()` is ok. 2. setting `set_allocate_memory_from_process_group_for_comm()` after custom comms is ko. Basically `set_allocate_memory_from_process_group_for_comm` is like a drop in hammer while the `set_custom_all_gather/reduce_scatter()` are like finer-grained scalpels that require more code crafted. We can revisit this if there's use case in between but for now they can be largely viewed independent from each other (even tho we do share some of the underlying pieces for now, that could be subject to change and should not be exposed to end users). Test Plan: added UT Reviewed By: weifengpy Differential Revision: D77681620
pytorchmergebot
pushed a commit
that referenced
this pull request
Jul 3, 2025
…ther with custom comm hooks (#157487) Summary: This is a follow up after the PR to add comm override support: #155189 The previous PR loosely checks the allocation mixin classes, which isn't really safe as the actual hook may still override the behavior. This may lead to unnecessary confusion for no good use case. So for now we just make the 2 sets of APIs largely incompatible: 1. setting custom comms after `set_allocate_memory_from_process_group_for_comm()` is ok. 2. setting `set_allocate_memory_from_process_group_for_comm()` after custom comms is ko. Basically `set_allocate_memory_from_process_group_for_comm` is like a drop in hammer while the `set_custom_all_gather/reduce_scatter()` are like finer-grained scalpels that require more code crafted. We can revisit this if there's use case in between but for now they can be largely viewed independent from each other (even tho we do share some of the underlying pieces for now, that could be subject to change and should not be exposed to end users). Test Plan: added UT Differential Revision: D77681620 Pull Request resolved: #157487 Approved by: https://github.com/weifengpy
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
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.
Summary:
This change introduces 2 comm override APIs:
set_custom_all_gatherandset_custom_reduce_scatterto allow for custom behavior respectively.This allow users to control how the comm buffers are allocated and the exact comm implementation for flexibility.
For details, see docstring in
Commin_fsdp_api.pyRelated PR:
#150564
Test Plan: CI
Differential Revision: D75714362
cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @kadeng @muchulee8 @amjames @chauhang @aakhundov