[fix][FSDP] fix weight init when using apply() (fixes #490 and #444) #543
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
46f422f to
a31d5dd
Compare
|
Doesn't work with SyncBN, will fix |
min-xu-ai
approved these changes
Mar 20, 2021
|
The test failures might be related to the cast_buffer change? |
It's very weird, tests only seem to fail on 1.6, but pass in 1.7.1 and 1.8 (and on my local machine)... will dig a bit |
facebook-github-bot
pushed a commit
to facebookresearch/vissl
that referenced
this pull request
Apr 15, 2021
Summary: Before this PR (facebookresearch/fairscale#543) was merged, we used to need the extra cuda() calls. Now, they are not needed. Unfortunately, this doesn't solve the long model init time issue we have. A FSDP model init still take >20 mins for me. This is really bad for debugging the regnet128 conv layer crash problem I am debugging. The following debugging output shows that most delays are in FSDP wrapping, some in BN wrapping and some in the layer wrapping. ``` INFO 2021-04-14 12:18:35,883 regnet_2.py: 159: block created INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:18:35,884 regnet_2.py: 161: cpu INFO 2021-04-14 12:19:07,388 regnet_2.py: 163: block bn wrapped INFO 2021-04-14 12:19:18,388 regnet_2.py: 166: block wrapped ``` In any case, this PR is pretty safe and should go in so that we don't need to do an extra `cuda()` call before wrapping. Pull Request resolved: fairinternal/ssl_scaling#75 Reviewed By: prigoyal Differential Revision: D27776285 Pulled By: min-xu-ai fbshipit-source-id: 3e43c6fe750fd6ee35933400b03a069d62040d8a
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.
compute_deviceso that we can usesummon_full_paramsimmediately afterFSDP.__init__, even if the params are still on CPU (this also fixes [FSDP] improve robustness to mismatch between torch.cuda.current_device and model's device #444)FSDP.applyso that it callssummon_full_paramsfirst. This makes it possible to do weight inits viamodel.apply(custom_weight_init_fn)without segfaulting, and should give identical results to not using FSDPThis also required reworking
_all_buffers_toto no longer useapply, since it is called from within_lazy_initand created some circular logic:apply -> summon_full_params -> _lazy_init -> _all_buffers_to -> apply.