[Traceable FSDP2] Add all_gather_into_tensor out variant #126334
Closed
Conversation
This file contains 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
pytorch-bot
bot
added
oncall: distributed
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/126334
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 6722c5c with merge base c312cd8 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
yf225
force-pushed
the
inductor_inplace_ag
branch
from
5f18b32
to
172bc95
Compare
wanchaol
reviewed
May 15, 2024
| @@ -321,6 +334,13 @@ TORCH_LIBRARY(_c10d_functional, m) { | |||
| c10::DispatchKey::CompositeExplicitAutograd, ::all_reduce_coalesced_), | |||
| {at::Tag::pt2_compliant_tag}); | |||
|
|
|||
| m.def( | |||
| "all_gather_into_tensor_(Tensor(a!) output, Tensor input, int group_size, str group_name) -> Tensor(a!)", | |||
There was a problem hiding this comment.
I would really love this to be an out-variant API, this is not a inplace allgather from operator prospective, it's a out variant op, we should follow the aten naming convention to make the op be an actual out-variant
There was a problem hiding this comment.
Sounds great! Updated.
yf225
changed the title
yf225
changed the title
yf225
changed the title
yifuwang
reviewed
May 16, 2024
| @@ -321,6 +334,13 @@ TORCH_LIBRARY(_c10d_functional, m) { | |||
| c10::DispatchKey::CompositeExplicitAutograd, ::all_reduce_coalesced_), | |||
| {at::Tag::pt2_compliant_tag}); | |||
|
|
|||
| m.def( | |||
| "all_gather_into_tensor_out(Tensor input, int group_size, str group_name, *, Tensor(a!) out) -> Tensor(a!)", | |||
There was a problem hiding this comment.
Nice! We can establish the convention that:
- Collectives that modify the input are postfixed with
_ - Out-variant of collectives are postfixed with
_out, and theoutargument is keyword-only
|
@pytorchbot merge |
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 |
ZelboK
pushed a commit
to ZelboK/pytorch
that referenced
this pull request
May 19, 2024
) This PR adds `torch.ops._c10d_functional.all_gather_into_tensor_out`. It's important for tracing FSDP2, because FSDP2 pre-allocates the output buffer of AllGather, and makes input buffer an alias of the output buffer, and expects both of them to be used to achieve lower memory usage. If we don't preserve this behavior and instead functionalize the AllGather op, AllGather op will then create a brand-new output buffer (instead of reusing), thus significantly increasing the memory usage. The expectation is that we will "re-inplace" the AllGather op by switching to the out variant in Inductor post-grad stage via an FX pass, so this API is not expected to be directly used by users. Pull Request resolved: pytorch#126334 Approved by: https://github.com/yifuwang, https://github.com/wanchaol
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.
This PR adds
torch.ops._c10d_functional.all_gather_into_tensor_out.It's important for tracing FSDP2, because FSDP2 pre-allocates the output buffer of AllGather, and makes input buffer an alias of the output buffer, and expects both of them to be used to achieve lower memory usage. If we don't preserve this behavior and instead functionalize the AllGather op, AllGather op will then create a brand-new output buffer (instead of reusing), thus significantly increasing the memory usage.
The expectation is that we will "re-inplace" the AllGather op by switching to the out variant in Inductor post-grad stage via an FX pass, so this API is not expected to be directly used by users.
cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @chauhang @d4l3k