Skip to content

[FSDP] summon_full_params() in computation stream #86836

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
Closed
wants to merge 5 commits into from
Closed

[FSDP] summon_full_params() in computation stream #86836

wants to merge 5 commits into from

Conversation

awgu
Copy link
Collaborator

@awgu awgu commented Oct 12, 2022 •
edited
Loading

Stack from ghstack:

This should help with memory usage. In particular, this allows FSDP to use caching allocator blocks from the computation stream for the summon_full_params() all-gathers, which should help avoid over-allocating blocks to the unshard stream.

@awgu awgu mentioned this pull request Oct 12, 2022
Closed
@pytorch-bot pytorch-bot bot added the release notes: distributed (sharded) release notes category label Oct 12, 2022
This was referenced Oct 12, 2022
Closed
Closed
@pytorch-bot
Copy link

pytorch-bot bot commented Oct 12, 2022 •
edited
Loading

🔗 Helpful Links

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

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit c777cde:
💚 Looks good so far! There are no failures yet. 💚

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

awgu pushed a commit that referenced this pull request Oct 12, 2022
[FSDP] summon_full_params() in computation stream
b421f68 
ghstack-source-id: 53f4fcf
Pull Request resolved: #86836
@awgu awgu added release notes: distributed (fsdp) release notes category and removed release notes: distributed (sharded) release notes category labels Oct 12, 2022
@awgu awgu marked this pull request as ready for review October 12, 2022 21:37
@awgu awgu marked this pull request as draft October 13, 2022 01:10
@awgu awgu mentioned this pull request Oct 13, 2022
Closed
@awgu awgu marked this pull request as ready for review October 13, 2022 02:16
@awgu awgu mentioned this pull request Oct 13, 2022
Closed
awgu pushed a commit to awgu/pytorch that referenced this pull request Oct 15, 2022
[FSDP] summon_full_params() in computation stream
fe1d22d 
ghstack-source-id: 2bc70f8
Pull Request resolved: pytorch#86836
Rick0317 pushed a commit to Rick0317/pytorch that referenced this pull request Oct 18, 2022
[FSDP] summon_full_params() in computation stream
ce3e24a 
ghstack-source-id: 9463308
Pull Request resolved: pytorch/pytorch#86836
awgu pushed a commit to awgu/pytorch that referenced this pull request Oct 21, 2022
[FSDP] summon_full_params() in computation stream
db96ec2 
ghstack-source-id: 2bc70f8
Pull Request resolved: pytorch#86836
rohan-varma
rohan-varma approved these changes Oct 24, 2022
View reviewed changes
Copy link
Contributor

@rohan-varma rohan-varma left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, but curious why we need this change.

torch/distributed/fsdp/fully_sharded_data_parallel.py
free_unsharded_flat_params = [handle.needs_unshard() for handle in self._handles]
self._unshard(self._handles)
self._streams["computation"].wait_stream(self._streams["unshard"])
# No need to call `wait_stream()` since we unshard in the computation
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Curious, why would we want to move this to computation stream?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This allows us to use caching allocator blocks from the computation stream for these all-gathers, which should help avoid over-allocating blocks to the unshard stream.

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 24, 2022
Update on "[FSDP] summon_full_params() in computation stream"
d060183 
This should help with memory usage.

[ghstack-poisoned]
Update on "[FSDP] summon_full_params() in computation stream"
c777cde 
This should help with memory usage.

[ghstack-poisoned]
awgu pushed a commit that referenced this pull request Oct 24, 2022
[FSDP] summon_full_params() in computation stream
2bec650 
ghstack-source-id: 7982746
Pull Request resolved: #86836
@awgu
Copy link
Collaborator Author

awgu commented Oct 24, 2022

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your 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

Advanced Debugging
Check the merge workflow status
here

@github-actions
Copy link
Contributor

Hey @awgu.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

@awgu awgu added the topic: not user facing topic category label Oct 24, 2022
kulinseth pushed a commit to kulinseth/pytorch that referenced this pull request Nov 5, 2022
@kulinseth
[FSDP] summon_full_params() in computation stream (pytorch#86836)
2657fe3 
This should help with memory usage. In particular, this allows FSDP to use caching allocator blocks from the computation stream for the `summon_full_params()` all-gathers, which should help avoid over-allocating blocks to the unshard stream.
Pull Request resolved: pytorch#86836
Approved by: https://github.com/rohan-varma
kulinseth pushed a commit to kulinseth/pytorch that referenced this pull request Dec 10, 2022
@kulinseth
[FSDP] summon_full_params() in computation stream (pytorch#86836)
ee7f2ab 
This should help with memory usage. In particular, this allows FSDP to use caching allocator blocks from the computation stream for the `summon_full_params()` all-gathers, which should help avoid over-allocating blocks to the unshard stream.
Pull Request resolved: pytorch#86836
Approved by: https://github.com/rohan-varma
@facebook-github-bot facebook-github-bot deleted the gh/awgu/125/head branch June 8, 2023 15:22
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@mrshenli mrshenli Awaiting requested review from mrshenli

@zhaojuanmao zhaojuanmao Awaiting requested review from zhaojuanmao

@pritamdamania87 pritamdamania87 Awaiting requested review from pritamdamania87

@mingzhe09088 mingzhe09088 Awaiting requested review from mingzhe09088

@H-Huang H-Huang Awaiting requested review from H-Huang

@kwen2501 kwen2501 Awaiting requested review from kwen2501

1 more reviewer

@rohan-varma rohan-varma rohan-varma approved these changes

Reviewers whose approvals may not affect merge requirements
Assignees
No one assigned
Labels
ciflow/trunk Trigger trunk jobs on your pull request Merged release notes: distributed (fsdp) release notes category topic: not user facing topic category
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@awgu @pytorchmergebot @rohan-varma