Skip to content

Introduce new shampoo and soap - #281

Merged
skyw merged 13 commits into
mainfrom
skyw/introduce_new_shampoo_and_soap
Aug 10, 2026
Merged

Introduce new shampoo and soap#281
skyw merged 13 commits into
mainfrom
skyw/introduce_new_shampoo_and_soap

Conversation

@skyw

@skyw skyw commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Improve abstraction. clear path for more variant of Shampoo and Soap.

@copy-pr-bot

copy-pr-bot Bot commented Aug 7, 2026

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.

skyw added 4 commits August 7, 2026 16:25
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
@skyw
skyw force-pushed the skyw/introduce_new_shampoo_and_soap branch from f050252 to a9b2fcf Compare August 7, 2026 23:25
@skyw

skyw commented Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

@greptileai review the draft

@skyw

skyw commented Aug 7, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test a9b2fcf

@greptile-apps

greptile-apps Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR promotes the experimental SOAP redesign into a dedicated Shampoo package and expands its preconditioner abstraction.

  • Adds shared Shampoo and SOAP preconditioner protocols.
  • Separates first-step initialization from subsequent factor and eigenbasis updates.
  • Adds ReKLS and per-parameter-group scalar optimizer settings.
  • Adds numerical parity, state-layout, projection, and smoke tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
emerging_optimizers/shampoo/shampoo_base.py Introduces shared tensor-pair and preconditioner protocols for Shampoo-family optimizers.
emerging_optimizers/shampoo/soap_v3.py Refactors SOAP preconditioning, adds explicit first-step handling and ReKLS, and fixes the previously reported factor-update precision gap.
tests/test_soap_v3.py Adds state, projection, legacy-parity, ReKLS-parity, and optimizer smoke coverage.

Reviews (8): Last reviewed commit: "relax threshold for rekls" | Re-trigger Greptile

Comment thread emerging_optimizers/shampoo/soap_v3.py Outdated
Comment thread emerging_optimizers/shampoo/soap_v3.py
@skyw

skyw commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test a9b2fcf

@github-actions

github-actions Bot commented Aug 8, 2026

Copy link
Copy Markdown

Test Results

   91 files  + 2    169 suites  +8   1m 53s ⏱️ +3s
1 257 tests +27  1 257 ✅ +27  0 💤 ±0  0 ❌ ±0 
2 936 runs  +54  2 936 ✅ +54  0 💤 ±0  0 ❌ ±0 

Results for commit be08912. ± Comparison against base commit f507fc9.

♻️ This comment has been updated with latest results.

skyw added 2 commits August 7, 2026 19:34
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
@skyw
skyw marked this pull request as ready for review August 8, 2026 03:23
Signed-off-by: Hao Wu <skyw@nvidia.com>
@skyw

skyw commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 715ce11

skyw added 2 commits August 7, 2026 20:33
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
@skyw

skyw commented Aug 8, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 4f516ff

skyw added 2 commits August 10, 2026 09:25
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
@skyw

skyw commented Aug 10, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test 74ac1aa

@skyw

skyw commented Aug 10, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test f3fe539

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

skyw commented Aug 10, 2026

Copy link
Copy Markdown
Contributor Author

/ok to test be08912

@codecov

codecov Bot commented Aug 10, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@skyw
skyw merged commit 65e08e6 into main Aug 10, 2026
30 checks passed
@skyw
skyw deleted the skyw/introduce_new_shampoo_and_soap branch August 10, 2026 20:07
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.

1 participant