-
Notifications
You must be signed in to change notification settings - Fork 3.6k
Sharded Plugin 2/n: Allow ddp plugin to modify optimizer state saving #4675
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
Conversation
Codecov Report
@@ Coverage Diff @@
## master #4675 +/- ##
======================================
Coverage 93% 93%
======================================
Files 117 117
Lines 8941 8949 +8
======================================
+ Hits 8311 8319 +8
Misses 630 630 |
tchaton
left a comment
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 !
Borda
left a comment
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
williamFalcon
left a comment
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.
dope
# Conflicts: # pytorch_lightning/plugins/ddp_plugin.py
…#4675) * Allow ddp plugin to modify optimizer state saving * Rely on the accelerator for optimizer states * Ensure we init the accelerator for the saving function * Better comment for optim state dump * Revert "Ensure we init the accelerator for the saving function" This reverts commit af65eff * Added accelerator check to initialize tuner before saving model checkpoint * Simplify comment * Revert "Added accelerator check to initialize tuner before saving model checkpoint" This reverts commit f9929c0 * Return single optimizer state to reduce duplication * Fixed docstring * Fixed typing * Fixed comment * Added CHANGELOG.md Co-authored-by: chaton <[email protected]>
What does this PR do?
Ties to #4178
Allows custom ddp plugin access to saving optimizer state. This is required for optimizer state sharding as we need to consolidate the optimizer state from all processes before saving.
cc @ananthsub @awaelchli
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 🙃