Skip to content

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

Open
wants to merge 20 commits into
base: main
Choose a base branch
from

Conversation

runame
Copy link
Contributor

@runame runame commented May 30, 2025

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 and EigenvalueCorrectedShampooPreconditionerList.

@runame runame requested review from hjmshi and tsunghsienlee May 30, 2025 05:49
@runame runame self-assigned this May 30, 2025
@runame runame added the enhancement New feature or request label May 30, 2025
@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 30, 2025
Copy link
Contributor

@tsunghsienlee tsunghsienlee left a 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:

  1. Limit this config so we don't need to worry about this for ShampooPreconditionerConfig.
  2. No need to modify shampoo_preconditioner_list.py.
  3. Keep _eigenvalues_estimate_criterion_below_or_equal_tolerance() as it is now, and people who use matrix_eigendecomposition() could still use this always.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. enhancement New feature or request
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants