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
Allow frozen data classes in optimizer state dict #16656
Allow frozen data classes in optimizer state dict #16656
Conversation
The frozen data class should be explicitly allowed to exist in the optimizer's `state`.
Converted to a draft for safety regarding the other merge happening first. This PR itself is ready, though. |
Other merge happened, this is good to go from my side. |
Tests fail because it expects the |
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## master #16656 +/- ##
========================================
- Coverage 82% 82% -0%
========================================
Files 438 418 -20
Lines 31744 31534 -210
========================================
- Hits 25991 25786 -205
+ Misses 5753 5748 -5 |
Thanks for your work on fixing everything up! |
Thanks, and this was a great addition. 2.0 here we come! |
What does this PR do?
We use the (now merged) PR in
lightning-utilities
that allows frozen dataclasses to be ignored inapply_to_collection
. With this, we can now allow saving optimizers for which thestate
includes frozen dataclasses that we just want to pass through.Any optimizer that includes a frozen dataclass in its
state
, is not supported in Lightning without this, instead causing errors.This should be backward-compatible, since previously, frozen dataclasses in optimizers simply caused an error. Other than that, nothing changed.
Before submitting
I tried, tests were broken for me even without the changes. Manually running
tests/tests_fabric/utilities/test_optimizer.py
does not complain.Did not seem necessary.