-
Notifications
You must be signed in to change notification settings - Fork 175
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
Support PyTorch 2.0.0 #559
Conversation
PyTorch 2.0 is not supported yet, it's support will be added in #559 (there are multiple issues to resolve). Until then, we need to require `torch<2.0.0` (otherwise 2.0.0 is installed, so CI doesn't work right now). This PR also adds Python 3.10 to the CI.
I've rebased this PR to the master branch after #558 was merged. |
@@ -1,5 +1,5 @@ | |||
PyYAML | |||
torch>=1.9.0,<2.0.0 | |||
torch>=1.9.0 |
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.
We'll have torch 1.13.0 on Python 3.7, torch 2.0.0 on Python 3.8-3.10. This allows us to test both torch 1.13.0 and 2.0.0.
We can drop Python 3.7 and PyTorch < 2.0 support when necessary.
hivemind/optim/optimizer.py
Outdated
def zero_grad(self, set_to_none: bool = False): | ||
_SET_TO_NONE_DEFAULT = Version(torch.__version__).major >= 2 | ||
|
||
def zero_grad(self, set_to_none: bool = _SET_TO_NONE_DEFAULT): |
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.
So it's consistent with the torch-wide default.
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
tests/test_optimizer.py
Outdated
assert (model1.weight.grad is None or torch.all(model1.weight.grad == 0)) and ( | ||
model2.weight.grad is None or torch.all(model2.weight.grad == 0) | ||
), "zero grad did not trigger" |
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.
Can we also use Version(torch.__version__).major >= 2
here? (or at least import _SET_TO_NONE_DEFAULT)
This way, the test won't have a bug where we always set the grad to None regardless of the PyTorch version
Codecov Report
@@ Coverage Diff @@
## master #559 +/- ##
==========================================
- Coverage 76.01% 75.92% -0.09%
==========================================
Files 81 81
Lines 7995 8008 +13
==========================================
+ Hits 6077 6080 +3
- Misses 1918 1928 +10
|
- Fix LRSchedulerBase - Handle None after .zero_grad() in torch 2.0.0 - Use set_to_none=True by default in torch>=2.0 - Add set_to_none param to TrainingStateAverager.step() Co-authored-by: Aleksandr Borzunov <[email protected]> (cherry picked from commit 98531ce)
@borzunov: