-
Notifications
You must be signed in to change notification settings - Fork 3.4k
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.
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
Sharded Plugin 3/n: Expose step input to DDP plugin #4686
Conversation
@@ -61,21 +61,23 @@ def train(self): | |||
return self.ddp_train(process_idx=self.task_idx, mp_queue=None, model=model) | |||
|
|||
def training_step(self, args): | |||
return self._step(args) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nice. this is better haha.
Codecov Report
@@ Coverage Diff @@
## master #4686 +/- ##
======================================
Coverage 93% 93%
======================================
Files 117 117
Lines 8939 8941 +2
======================================
+ Hits 8306 8308 +2
Misses 633 633 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM !
Co-authored-by: Jirka Borovec <[email protected]>
Hello @SeanNaren! Thanks for updating this PR. There are currently no PEP 8 issues detected in this Pull Request. Cheers! 🍻 Comment last updated at 2020-11-18 14:16:08 UTC |
return self._step(args) | ||
|
||
def _step(self, args): | ||
args = self.ddp_plugin.on_before_forward(self.trainer.get_model(), *args) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
love it
* Allow ddp plugin to move the input to a different device if needed * Swapped name to on_before_forward to align with hooks in the future * Update pytorch_lightning/plugins/ddp_plugin.py Co-authored-by: Jirka Borovec <[email protected]> * Pass variable arg type to hook, add example * Remove blank space (pep check) * Added blank line Co-authored-by: William Falcon <[email protected]> Co-authored-by: Jirka Borovec <[email protected]>
What does this PR do?
Relates to #4178
ShardedDDP needs to move the device to the GPU before passing through to the model. This is not the case for DDP which handles this internally. Allows other custom DDP parallel plugins an opportunity to modify the input before passing into the wrapper!
cc @williamFalcon @tchaton @awaelchli @Borda
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 in short, see following bullet-list:
Did you have fun?
Make sure you had fun coding 🙃