Add branch condition for calling move to device in prefetch (FSDP 3/n) #6342
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
SeanNaren
added
feature
SeanNaren
requested review from
awaelchli,
Borda,
carmocca,
justusschock,
tchaton and
williamFalcon
as code owners
SeanNaren
commented
Mar 4, 2021
|
|
||
| @pytest.mark.parametrize('move_to_device_pre_dispatch_enabled', [True, False]) | ||
| @mock.patch('pytorch_lightning.plugins.DDPPlugin.model_to_device') | ||
| def test_move_to_device_in_pre_dispatch(mock_model_to_device, tmpdir, move_to_device_pre_dispatch_enabled): |
There was a problem hiding this comment.
Two things:
- Would've liked to add a spawn test but mocks don't seem to carry across to new processes?
- Can I combine patch and parametrize?
There was a problem hiding this comment.
- Nope, you'd have to create the mock in each process
- What do you mean? Having different mocks for each parametrize? I don't think so
There was a problem hiding this comment.
You might be able to do so with a callback, applying the patch context manager.
There was a problem hiding this comment.
You can make patch conditional on parametrize. with the return object.
Codecov Report
@@ Coverage Diff @@
## master #6342 +/- ##
=======================================
+ Coverage 91% 92% +1%
=======================================
Files 160 161 +1
Lines 11408 11476 +68
=======================================
+ Hits 10384 10547 +163
+ Misses 1024 929 -95 |
carmocca
reviewed
Mar 4, 2021
Comment on lines
+317
to
+323
| @property | ||
| def call_move_to_device_hook_in_pre_dispatch(self) -> bool: | ||
| """ | ||
| Call the ``model_to_device`` function within pre_dispatch if this is set to True. | ||
| Useful for when plugin would like to call model_to_device at another time, or skip the call. | ||
| """ | ||
| return True |
There was a problem hiding this comment.
Should this be in ParallelPlugin?
| ): | ||
| model = BoringModel() | ||
| trainer = Trainer( | ||
| default_root_dir=tmpdir, fast_dev_run=True, accelerator='ddp', plugins=DDPPlugin(), num_processes=1 |
There was a problem hiding this comment.
Why do you need to pass DDPPlugin?
|
|
||
| @pytest.mark.parametrize('move_to_device_pre_dispatch_enabled', [True, False]) | ||
| @mock.patch('pytorch_lightning.plugins.DDPPlugin.model_to_device') | ||
| def test_move_to_device_in_pre_dispatch(mock_model_to_device, tmpdir, move_to_device_pre_dispatch_enabled): |
There was a problem hiding this comment.
- Nope, you'd have to create the mock in each process
- What do you mean? Having different mocks for each parametrize? I don't think so
Co-authored-by: Carlos Mocholí <[email protected]>
|
Hey @SeanNaren. Any updates ? |
|
Given how little now exists in the pre_dispatch function after #6506, we don't need this anymore :) |
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.
What does this PR do?
Adds an additional property to the DDP/DDP Spawn Plugin that allows plugins that extend DDP to disable moving model to device. In FSDP this is required since we can only move to device after configure_ddp is called, which would require overriding this function pre_dispatch. Also fixed an issue where spawning the DDP Plugin was not possible because we missed copying a default setter.
Fixes #<issue_number>
Before submitting
PR review
Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:
Did you have fun?
Make sure you had fun coding 🙃