-
Notifications
You must be signed in to change notification settings - Fork 43
Add adaptive eigendecomposition computation frequency for eigh
#177
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
base: main
Are you sure you want to change the base?
Conversation
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.
It seems the best place to implement this is inside matrix_eigendecomposition()
with only adding this field into EigendecompositionConfig
from the interface/engineering perspective (not sure this is the best from the algorithm/semantic level, need to discuss this with you). If this is possible, following are the benefits:
- Limit this config so we don't need to worry about this for
ShampooPreconditionerConfig
. - No need to modify
shampoo_preconditioner_list.py
. - Keep
_eigenvalues_estimate_criterion_below_or_equal_tolerance()
as it is now, and people who usematrix_eigendecomposition()
could still use this always.
Just like (approximate) eigendecompositions can be skipped in the warm-started QR algorithm, we can skip eigendecompositions with
torch.linalg.eigh
using the same criterion.This PR implements this for
EigendecomposedShampooPreconditionerList
andEigenvalueCorrectedShampooPreconditionerList
.