Skip to content

[FSDP] Allow to use TorchDispatch with FSDP #88014

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

[FSDP] Allow to use TorchDispatch with FSDP #88014

fegin wants to merge 5 commits into from

Conversation

@fegin
Copy link
Contributor

@fegin fegin commented Oct 28, 2022 •
edited
Loading

Stack from ghstack (oldest at bottom):

Add _no_dispatch_record_stream to disable TorchDispatch before calling record_stream().

@pytorch-bot
Copy link

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

🔗 Helpful Links

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

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

✅ No Failures

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

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

@pytorch-bot pytorch-bot bot added the release notes: distributed (fsdp) release notes category label Oct 28, 2022
@awgu
Copy link
Collaborator

awgu commented Oct 28, 2022

Is this for landing, or is this just for experimentation? I am concerned that memory logs from using this will be inaccurate because memory will be incorrectly freed early without record_stream().

zhaojuanmao
zhaojuanmao reviewed Oct 28, 2022
View reviewed changes
Copy link
Contributor

@zhaojuanmao zhaojuanmao left a comment

Choose a reason for hiding this comment

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

look good to me, will let Andrew or Rohan to accept

@fegin
Copy link
Contributor Author

fegin commented Oct 28, 2022 •
edited
Loading

@awgu record_stream() is still called with this PR. What _no_dispatch_record_stream() does is to disable TorchDispatch before calling record_stream() and reenable it after record_stream() is done. So the memory won't be freed early.

awgu
awgu approved these changes Oct 28, 2022
View reviewed changes
Copy link
Collaborator

@awgu awgu left a comment

Choose a reason for hiding this comment

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

Thanks for the explanation! I definitely misunderstood.

This looks good to me. Do you know if Core team is planning to add support for record_stream() so that this is a temporary fix?

@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Oct 28, 2022
@fegin
Copy link
Contributor Author

fegin commented Oct 31, 2022 •
edited
Loading

This looks good to me. Do you know if Core team is planning to add support for record_stream() so that this is a temporary fix?

Saw a TODO for TorchDispatch to support class block list filter. If that function is implemented, users can just use the block list to filter out record_stream(). Not sure when that feature will be implemented though.

@awgu
Copy link
Collaborator

awgu commented Oct 31, 2022

This looks good to me. Do you know if Core team is planning to add support for record_stream() so that this is a temporary fix?

Saw a TODO for TorchDispatch to support class block list filter. If this function is implemented, users can just the block list to filter out record_stream(). Not sure when that feature will be implemented though.

Sounds good to me!

Also, sorry, I am landing some PRs from my stack, so there may be some rebase conflicts :(

@fegin
Update on "[FSDP] Allow to use TorchDispatch with FSDP"
e011b6c 
Add `_no_dispatch_record_stream` to disable TorchDispatch before calling `record_stream()`.

[ghstack-poisoned]
@fegin
Update on "[FSDP] Allow to use TorchDispatch with FSDP"
453084f 
Add `_no_dispatch_record_stream` to disable TorchDispatch before calling `record_stream()`.

[ghstack-poisoned]
fegin added a commit that referenced this pull request Oct 31, 2022
@fegin
[FSDP] Allow to use TorchDispatch with FSDP
a9c726c 
ghstack-source-id: 6786e65
Pull Request resolved: #88014
@fegin
Update on "[FSDP] Allow to use TorchDispatch with FSDP"
e481089 
Add `_no_dispatch_record_stream` to disable TorchDispatch before calling `record_stream()`.

[ghstack-poisoned]
@fegin
Update on "[FSDP] Allow to use TorchDispatch with FSDP"
b594741 
Add `_no_dispatch_record_stream` to disable TorchDispatch before calling `record_stream()`.

[ghstack-poisoned]
fegin added a commit that referenced this pull request Nov 3, 2022
@fegin
[FSDP] Allow to use TorchDispatch with FSDP
a175ccb 
ghstack-source-id: c850496
Pull Request resolved: #88014
@fegin
Copy link
Contributor Author

fegin commented Nov 3, 2022

@pytorchbot merge

@pytorchmergebot
Copy link
Collaborator

pytorchmergebot commented Nov 3, 2022

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

kulinseth pushed a commit to kulinseth/pytorch that referenced this pull request Nov 5, 2022
@fegin @kulinseth
[FSDP] Allow to use TorchDispatch with FSDP (pytorch#88014)
a0ef18c 
Add `_no_dispatch_record_stream` to disable TorchDispatch before calling `record_stream()`.
Pull Request resolved: pytorch#88014
Approved by: https://github.com/awgu
kulinseth pushed a commit to kulinseth/pytorch that referenced this pull request Dec 10, 2022
@fegin @kulinseth
[FSDP] Allow to use TorchDispatch with FSDP (pytorch#88014)
47d5732 
Add `_no_dispatch_record_stream` to disable TorchDispatch before calling `record_stream()`.
Pull Request resolved: pytorch#88014
Approved by: https://github.com/awgu
@facebook-github-bot facebook-github-bot deleted the gh/fegin/37/head branch June 8, 2023 17:15
@ezyang ezyang mentioned this pull request Aug 11, 2023
Closed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@zhaojuanmao zhaojuanmao zhaojuanmao left review comments

@awgu awgu awgu approved these changes

@mrshenli mrshenli Awaiting requested review from mrshenli

@pritamdamania87 pritamdamania87 Awaiting requested review from pritamdamania87

@rohan-varma rohan-varma Awaiting requested review from rohan-varma

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

@kwen2501 kwen2501 Awaiting requested review from kwen2501

Assignees

No one assigned

Labels

ciflow/trunk Trigger trunk jobs on your pull request Merged release notes: distributed (fsdp) release notes category

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@fegin @awgu @pytorchmergebot @zhaojuanmao