Fix the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo #111
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
facebook-github-bot
added
the
CLA Signed
|
This pull request was exported from Phabricator. Differential Revision: D72093421 |
tsunghsienlee
added a commit
to tsunghsienlee/optimizers
that referenced
this pull request
Mar 29, 2025
…nvalue-corrected Shampoo (facebookresearch#111) Summary: This diff fixes the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo. For eigendecompose Shampoo, this changes include initializing the eigenvectors matrices as identity matrices and initializing the eigenvalues as all ones. For eigenvalue-corrected Shampoo, this changes include initializing the eigenvectors matcies as identity matrices, and keeping corrected eigenvalues as all 0s. Extend `distributed_state_dict()`-related tests to include testing eigendecompose Shampoo and eigenvalue-corrected Shampoo to verify the initializations. Differential Revision: D72093421
tsunghsienlee
force-pushed
the
export-D72093421
branch
from
e8cd97f to
292672b
Compare
|
This pull request was exported from Phabricator. Differential Revision: D72093421 |
tsunghsienlee
added a commit
to tsunghsienlee/optimizers
that referenced
this pull request
Mar 29, 2025
…nvalue-corrected Shampoo (facebookresearch#111) Summary: This diff fixes the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo. For eigendecompose Shampoo, this changes include initializing the eigenvectors matrices as identity matrices and initializing the eigenvalues as all ones. For eigenvalue-corrected Shampoo, this changes include initializing the eigenvectors matcies as identity matrices, and keeping corrected eigenvalues as all 0s. Extend `distributed_state_dict()`-related tests to include testing eigendecompose Shampoo and eigenvalue-corrected Shampoo to verify the initializations. Differential Revision: D72093421
tsunghsienlee
force-pushed
the
export-D72093421
branch
from
292672b to
cad128b
Compare
|
This pull request was exported from Phabricator. Differential Revision: D72093421 |
tsunghsienlee
added a commit
to tsunghsienlee/optimizers
that referenced
this pull request
Mar 29, 2025
…nvalue-corrected Shampoo (facebookresearch#111) Summary: Pull Request resolved: facebookresearch#111 This diff fixes the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo. For eigendecompose Shampoo, this changes include initializing the eigenvectors matrices as identity matrices and initializing the eigenvalues as all ones. For eigenvalue-corrected Shampoo, this changes include initializing the eigenvectors matcies as identity matrices, and keeping corrected eigenvalues as all 0s. Extend `distributed_state_dict()`-related tests to include testing eigendecompose Shampoo and eigenvalue-corrected Shampoo to verify the initializations. Differential Revision: D72093421
tsunghsienlee
force-pushed
the
export-D72093421
branch
from
cad128b to
390e7a7
Compare
runame
reviewed
Mar 31, 2025
Comment on lines
+1154
to
+1122
| # Initialize factor_matrices_eigenvalues all ones. | ||
| for t in factor_matrices_eigenvalues: | ||
| block_info.get_tensor(t).fill_(1.0) |
There was a problem hiding this comment.
This could also be implemented with a new block_info.allocate_ones_tensormethod.
Comment on lines
+1142
to
+1111
| # Initialize factor_matrices_eigenvectors as identity matrices. | ||
| for t in factor_matrices_eigenvectors: | ||
| block_info.get_tensor(t).fill_diagonal_(1.0) | ||
|
|
There was a problem hiding this comment.
This could be refactored as a new block_info.allocate_eye_tensor method that calls block_info.allocate_zeros_tensor followed by this code.
The advantage is a unified interface for all initialization options (zeros, ones, eye) and it will be easier to change the implementation of each initialization method later on.
There was a problem hiding this comment.
Although I really prefer not enhancing BlockInfo due to the desire of removing it, what you said here make sense because it is not increasing the overall complexity of BlockInfo.
tsunghsienlee
added a commit
to tsunghsienlee/optimizers
that referenced
this pull request
Apr 3, 2025
…nvalue-corrected Shampoo (facebookresearch#111) Summary: This diff fixes the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo. For eigendecompose Shampoo, this changes include initializing the eigenvectors matrices as identity matrices and initializing the eigenvalues as all ones. For eigenvalue-corrected Shampoo, this changes include initializing the eigenvectors matcies as identity matrices, and keeping corrected eigenvalues as all 0s. Extend `distributed_state_dict()`-related tests to include testing eigendecompose Shampoo and eigenvalue-corrected Shampoo to verify the initializations. Differential Revision: D72093421
tsunghsienlee
force-pushed
the
export-D72093421
branch
from
390e7a7 to
fd0e041
Compare
|
This pull request was exported from Phabricator. Differential Revision: D72093421 |
tsunghsienlee
added a commit
to tsunghsienlee/optimizers
that referenced
this pull request
Apr 3, 2025
…nvalue-corrected Shampoo (facebookresearch#111) Summary: Pull Request resolved: facebookresearch#111 This diff fixes the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo. For eigendecompose Shampoo, this changes include initializing the eigenvectors matrices as identity matrices and initializing the eigenvalues as all ones. For eigenvalue-corrected Shampoo, this changes include initializing the eigenvectors matcies as identity matrices, and keeping corrected eigenvalues as all 0s. Extend `distributed_state_dict()`-related tests to include testing eigendecompose Shampoo and eigenvalue-corrected Shampoo to verify the initializations. Differential Revision: D72093421
tsunghsienlee
force-pushed
the
export-D72093421
branch
from
fd0e041 to
daa1eaa
Compare
added 2 commits
April 3, 2025 15:51
Summary: 1. Adding a new method `_create_preconditioned_dims_selector()` to create a list of boolean values indicating whether each dimension should be preconditioned or not. 2. Modifying the `_create_preconditioned_dims_selector_list()` method to use the new `_create_preconditioned_dims_selector()` method to create the list of boolean values. Differential Revision: D72407362
Summary: This diff addresses a potential type mismatch issue between `factor_matrices` and `factor_matrix_eigenvectors` in the QR algorithm. It also ensures consistent output type conversions for eigendecomposed Shampoo and eigenvalue-corrected Shampoo. - Convert `factor_matrix_eigenvectors` to match the type of `factor_matrices` before applying the QR algorithm. - Update output type conversions for eigendecomposed Shampoo and eigenvalue-corrected Shampoo to align with Shampoo. Prevents errors caused by type mismatches in the QR algorithm and ensures consistency in output type conversions across different Shampoo variants. Differential Revision: D72334952 Reviewed By: anana10c
tsunghsienlee
force-pushed
the
export-D72093421
branch
from
daa1eaa to
cc27512
Compare
tsunghsienlee
added a commit
to tsunghsienlee/optimizers
that referenced
this pull request
Apr 4, 2025
…nvalue-corrected Shampoo (facebookresearch#111) Summary: This diff fixes the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo. For eigendecompose Shampoo, this changes include initializing the eigenvectors matrices as identity matrices and initializing the eigenvalues as all ones. For eigenvalue-corrected Shampoo, this changes include initializing the eigenvectors matcies as identity matrices, and keeping corrected eigenvalues as all 0s. Extend `distributed_state_dict()`-related tests to include testing eigendecompose Shampoo and eigenvalue-corrected Shampoo to verify the initializations. Reviewed By: anana10c Differential Revision: D72093421
…nvalue-corrected Shampoo (facebookresearch#111) Summary: Pull Request resolved: facebookresearch#111 This diff fixes the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo. For eigendecompose Shampoo, this changes include initializing the eigenvectors matrices as identity matrices and initializing the eigenvalues as all ones. For eigenvalue-corrected Shampoo, this changes include initializing the eigenvectors matcies as identity matrices, and keeping corrected eigenvalues as all 0s. Extend `distributed_state_dict()`-related tests to include testing eigendecompose Shampoo and eigenvalue-corrected Shampoo to verify the initializations. Reviewed By: anana10c Differential Revision: D72093421
|
This pull request was exported from Phabricator. Differential Revision: D72093421 |
tsunghsienlee
force-pushed
the
export-D72093421
branch
from
cc27512 to
e736eb4
Compare
facebook-github-bot
pushed a commit
that referenced
this pull request
Apr 7, 2025
Summary: Followed by the discussion in #111 (comment), add this function to unify the support of allocating an identiy matrix by leverging the existing `BlockInfo.allocate_zero_tensor()`. Reviewed By: anana10c Differential Revision: D72478672 fbshipit-source-id: 8c5d52d63e5bdaf74c77140f0fc5c97aff39f7fd
|
This pull request has been merged in b683e18. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
This diff fixes the statistics initializations of eigendecompose Shampoo and eigenvalue-corrected Shampoo. For eigendecompose Shampoo, this changes include initializing the eigenvectors matrices as identity matrices and initializing the eigenvalues as all ones. For eigenvalue-corrected Shampoo, this changes include initializing the eigenvectors matcies as identity matrices, and keeping corrected eigenvalues as all 0s.
Extend
distributed_state_dict()-related tests to include testing eigendecompose Shampoo and eigenvalue-corrected Shampoo to verify the initializations.Differential Revision: D72093421