[RFC] Don't materialize ignored modules for FSDP #108032
Closed
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
Per title. This seems needed for cases where I have a large embedding I want to separately manage, but FSDP would initialize it and thus consume the memory. Currently the interaction with torchdistX materialize_module is not tested, this can be done as follow up work. Differential Revision: [D48722046](https://our.internmc.facebook.com/intern/diff/D48722046/) [ghstack-poisoned]
rohan-varma
requested review from
H-Huang,
awgu,
d4l3k,
fduwjj,
fegin,
kiukchung,
kwen2501,
mrshenli,
wanchaol,
wz337 and
zhaojuanmao
as code owners
pytorch-bot
bot
added
the
release notes: distributed (fsdp)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/108032
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Merge Blocking SEVsThere is 1 active merge blocking SEVs. Please view them below:
If you must merge, use ❌ 4 New Failures, 1 Unrelated FailureAs of commit 017d4ae with merge base a20fac8 ( NEW FAILURES - The following jobs have failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
awgu
approved these changes
Aug 31, 2023
| device_id=self.rank, | ||
| ignored_modules=[m.a], | ||
| use_orig_params=True, | ||
| param_init_fn=lambda m: m.cuda(), |
There was a problem hiding this comment.
In general, this lambda m: m.cuda() would do some repeated checks trying to move modules to CUDA since param_init_fn would be called on every module. This should just lead to some CPU overhead since copying to CUDA is a no-op if already on CUDA.
As a tiny nit, the variable shadowing of m is also a bit precarious.
awgu
requested changes
Aug 31, 2023
|
That's weird, I feel like I ran the test before sending the PR... |
Per title. This seems needed for cases where I have a large embedding I want to separately manage, but FSDP would initialize it and thus consume the memory. Currently the interaction with torchdistX materialize_module is not tested, this can be done as follow up work. Differential Revision: [D48722046](https://our.internmc.facebook.com/intern/diff/D48722046/) [ghstack-poisoned]
Per title. This seems needed for cases where I have a large embedding I want to separately manage, but FSDP would initialize it and thus consume the memory. Currently the interaction with torchdistX materialize_module is not tested, this can be done as follow up work. Differential Revision: [D48722046](https://our.internmc.facebook.com/intern/diff/D48722046/) [ghstack-poisoned]
awgu
approved these changes
Sep 1, 2023
Per title. This seems needed for cases where I have a large embedding I want to separately manage, but FSDP would initialize it and thus consume the memory. Currently the interaction with torchdistX materialize_module is not tested, this can be done as follow up work. Differential Revision: [D48722046](https://our.internmc.facebook.com/intern/diff/D48722046/) [ghstack-poisoned]
pytorchmergebot
pushed a commit
that referenced
this pull request
Sep 5, 2023
Since these are ignored by FSDP, don't move them. Differential Revision: [D48727044](https://our.internmc.facebook.com/intern/diff/D48727044/) Pull Request resolved: #108033 Approved by: https://github.com/awgu ghstack dependencies: #108032
pytorchmergebot
pushed a commit
that referenced
this pull request
Sep 5, 2023
…NTRANT (#108435) We should use no_reentrant. There are a lot of users of this API, but it is in a prototype state so should be fine to change. Differential Revision: [D48898148](https://our.internmc.facebook.com/intern/diff/D48898148/) Pull Request resolved: #108435 Approved by: https://github.com/awgu ghstack dependencies: #108032, #108033
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 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.
Stack from ghstack (oldest at bottom):
Per title. This seems needed for cases where I have a large embedding
I want to separately manage, but FSDP would initialize it and thus consume the
memory.
Currently the interaction with torchdistX materialize_module is not tested,
this can be done as follow up work.
Differential Revision: D48722046