Skip to content

save unnecessary matmul#30

Merged
skyw merged 4 commits into
mainfrom
skyw/optimize_soap
Oct 1, 2025
Merged

save unnecessary matmul#30
skyw merged 4 commits into
mainfrom
skyw/optimize_soap

Conversation

@skyw

@skyw skyw commented Sep 30, 2025

Copy link
Copy Markdown
Contributor

Only diagonal of approx_eigenvalue_matrix is needed for _orthogonal_iteration, added logic to skip the last matrix multiply.

Frobenius norm is unitarily invariant, that is for any orthogonal matrix Q and square matrix A , $||Q^T A Q||_F = ||A||_F$. So we can calculate norm on kronecker_factor instead of approx_eigenvalue_matrix.

@skyw skyw requested a review from mkhona-nvidia September 30, 2025 22:39
@copy-pr-bot

copy-pr-bot Bot commented Sep 30, 2025

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Hao Wu <skyw@nvidia.com>
@skyw skyw force-pushed the skyw/optimize_soap branch from 990a47a to e3bf176 Compare September 30, 2025 22:43
@skyw

skyw commented Sep 30, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test e3bf176

Comment thread emerging_optimizers/soap/soap_utils.py Outdated
# (i.e. the approximated eigenvectors diagonalize the kronecker factor)
approx_eigenvalue_matrix = eigenbasis.T @ kronecker_factor @ eigenbasis
# Update eigenbasis when necessary. Update is skipped only when adaptive update criteria is met.
if _adaptive_criteria_met(

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

_adaptive_criteria_met also extracts the diagonal and computes diagonal norm, in addition to matrix frobenius norm. So I think this will need the full approx_eigenvalue_matrix diagonal for diagonal norm and the kronecker factor for frobenius norm, separately.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

right. Added it back.

Comment thread emerging_optimizers/soap/soap_utils.py
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Comment thread emerging_optimizers/soap/soap_utils.py
@skyw

skyw commented Sep 30, 2025

Copy link
Copy Markdown
Contributor Author

/ok to test 7265ff4

@mkhona-nvidia mkhona-nvidia left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

New change

@skyw skyw enabled auto-merge (squash) October 1, 2025 00:00
@skyw skyw merged commit 2ebdb41 into main Oct 1, 2025
12 checks passed
@skyw skyw deleted the skyw/optimize_soap branch October 1, 2025 00:02
mkhona-nvidia pushed a commit to mkhona-nvidia/Emerging-Optimizers that referenced this pull request Oct 3, 2025
* save unnecessary matmul

Signed-off-by: Hao Wu <skyw@nvidia.com>

* simplify criteria logic

Signed-off-by: Hao Wu <skyw@nvidia.com>

* remove max precondition dim

Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: mikail <mkhona@nvidia.com>
mkhona-nvidia pushed a commit to mkhona-nvidia/Emerging-Optimizers that referenced this pull request Oct 3, 2025
* save unnecessary matmul

Signed-off-by: Hao Wu <skyw@nvidia.com>

* simplify criteria logic

Signed-off-by: Hao Wu <skyw@nvidia.com>

* remove max precondition dim

Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: mikail <mkhona@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants