Improve shampoo update scale by skyw · Pull Request #306 · NVIDIA-NeMo/Emerging-Optimizers

skyw · 2026-10-05T20:44:01Z

Support different scaling method to match RMS of AdamW and other optimizer.

copy-pr-bot · 2026-10-05T20:44:05Z

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.

greptile-apps · 2026-10-05T20:47:52Z

Confidence Score: 5/5

[Medium risk] Refactors the shampoo optimizer's scaling calculation.

The PR appears safe to merge, though the previously noted scaling guidance remains absent.

Findings

Unbalanced KL factors change scaling ▶
Scaling guidance removed ▶

Summary

The PR adds selectable scaling for Shampoo updates and moves KL-correction metadata onto preconditioner classes. Since the previous review, it has updated the step loop to read that metadata from the selected preconditioner and adjusted the tests. No new actionable issue was established.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    G[Gradient] --> F[Update factor state]
    F --> P[Precondition scalar update]
    F --> S[Compute selected scale]
    P --> U[Scale and apply update]
    S --> U

_{Reviews (6) · Last reviewed commit: "qol improvement" · Reviewed by Greptile}

greptile-apps · 2026-10-05T21:08:45Z

+    trace_exp = 1 / p_root_inv - (0.5 if kl_corrected else 0.25)
+    scale: float | torch.Tensor = (size_out * size_in) ** (0.5 - 1 / p_root_inv)
+    if trace_exp != 0:
+        scale = scale * (L.trace().clamp_min(eps) ** trace_exp) * (R.trace().clamp_min(eps) ** trace_exp)


Unbalanced KL factors change scaling If KlShampoo(p_root_inv=4) has one factor trace below eps while the product of the traces remains finite, clamping each trace separately changes the update scale. KL factors can represent the same gradient statistics with different left/right splits, so this can make an otherwise equivalent update substantially smaller than the intended normalization.

Knowledge Base Used: Second-order preconditioning systems

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

@greptile It is valid but too rare of a case, ignore.

Understood — I’ll treat this as an acceptable edge case and won’t pursue the finding further for this PR.

skyw · 2026-10-07T21:22:58Z

/ok to test 4fdc833

copy-pr-bot · 2026-10-07T21:23:02Z

/ok to test 4fdc833

@skyw, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

skyw · 2026-10-07T22:51:56Z

/ok to test 4fdc833

copy-pr-bot · 2026-10-07T22:52:00Z

/ok to test 4fdc833

@skyw, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

skyw · 2026-10-07T23:05:45Z

/ok to test 15aa571

copy-pr-bot · 2026-10-07T23:05:48Z

/ok to test 15aa571

@skyw, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

skyw · 2026-10-07T23:10:50Z

/ok to test 9952428

copy-pr-bot · 2026-10-07T23:10:54Z

/ok to test 9952428

@skyw, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

greptile-apps · 2026-10-07T23:13:19Z

+    kl_corrected: bool = False,
+    eps: float,
+) -> float | torch.Tensor:
+    """Get the scale for the Shampoo update.


Scaling guidance removed The deleted explanation was the only guidance that the default spectral scale targets unit RMS for AdamW-like learning rates and that trace dependence disappears at root order 4 for Shampoo or 2 for KL-Shampoo. The remaining docstring names the modes but does not explain this behavior, making it harder to choose a mode or root order without unintentionally changing update magnitude. Please retain the practical selection guidance, even if the full derivation is too long.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

skyw · 2026-10-08T00:31:02Z

/ok to test 9952428

github-actions · 2026-10-08T00:45:00Z

Test Results

99 files ± 0 200 suites +2 2m 21s ⏱️ +11s
1 547 tests + 83 1 547 ✅ + 83 0 💤 ±0 0 ❌ ±0
3 673 runs +166 3 673 ✅ +166 0 💤 ±0 0 ❌ ±0

Results for commit 9952428. ± Comparison against base commit 8b775b5.

This pull request removes 1 and adds 84 tests. Note that renamed tests count towards both.

__main__.ShampooTest ‑ test_zero_grad_applies_only_weight_decay

__main__.GetShampooScaleFactorTest ‑ test_invalid_mode_raises
__main__.GetShampooScaleFactorTest ‑ test_smoke0 (shape=(8, 16), p_root_inv=1, mode='shape_scaling', kl_corrected=False)
__main__.GetShampooScaleFactorTest ‑ test_smoke1 (shape=(8, 16), p_root_inv=1, mode='shape_scaling', kl_corrected=True)
__main__.GetShampooScaleFactorTest ‑ test_smoke10 (shape=(8, 16), p_root_inv=2, mode='unit_rms_norm', kl_corrected=False)
__main__.GetShampooScaleFactorTest ‑ test_smoke11 (shape=(8, 16), p_root_inv=2, mode='unit_rms_norm', kl_corrected=True)
__main__.GetShampooScaleFactorTest ‑ test_smoke12 (shape=(8, 16), p_root_inv=4, mode='shape_scaling', kl_corrected=False)
__main__.GetShampooScaleFactorTest ‑ test_smoke13 (shape=(8, 16), p_root_inv=4, mode='shape_scaling', kl_corrected=True)
__main__.GetShampooScaleFactorTest ‑ test_smoke14 (shape=(8, 16), p_root_inv=4, mode='spectral', kl_corrected=False)
__main__.GetShampooScaleFactorTest ‑ test_smoke15 (shape=(8, 16), p_root_inv=4, mode='spectral', kl_corrected=True)
__main__.GetShampooScaleFactorTest ‑ test_smoke16 (shape=(8, 16), p_root_inv=4, mode='unit_rms_norm', kl_corrected=False)
…

codecov · 2026-10-08T00:45:22Z

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

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

Signed-off-by: Hao Wu <skyw@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Hao Wu <skyw@nvidia.com>

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

skyw · 2026-10-08T16:38:46Z

/ok to test 5501414

copy-pr-bot · 2026-10-08T16:38:50Z

/ok to test 5501414

@skyw, there was an error processing your request: E2

See the following link for more information: https://docs.gha-runners.nvidia.com/cpr/e/2/

greptile-apps Bot reviewed Oct 5, 2026

View reviewed changes

Comment thread emerging_optimizers/shampoo/shampoo.py Outdated

Comment thread emerging_optimizers/shampoo/shampoo.py Outdated

Comment thread emerging_optimizers/shampoo/shampoo.py Outdated

skyw force-pushed the skyw/shampoo_spectrum_scale branch from 8835476 to 4fdc833 Compare October 5, 2026 21:05

greptile-apps Bot reviewed Oct 5, 2026

View reviewed changes

skyw marked this pull request as draft October 5, 2026 21:12

skyw marked this pull request as ready for review October 7, 2026 19:57

skyw requested a review from mkhona-nvidia October 7, 2026 21:05

skyw force-pushed the skyw/shampoo_spectrum_scale branch from 4fdc833 to 15aa571 Compare October 7, 2026 23:04

greptile-apps Bot reviewed Oct 7, 2026

View reviewed changes

copy-pr-bot Bot deployed to public October 8, 2026 00:31 Active

copy-pr-bot Bot deployed to test October 8, 2026 00:31 Active

copy-pr-bot Bot deployed to public October 8, 2026 00:33 Active

skyw and others added 4 commits October 8, 2026 09:38

improve shampoo update scale

2d2a4d8

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

qol improvment

be6dff3

Signed-off-by: Hao Wu <skyw@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Hao Wu <skyw@nvidia.com>

remove some comment

ab3f779

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

qol improvement

5501414

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

skyw force-pushed the skyw/shampoo_spectrum_scale branch from 56729f9 to 5501414 Compare October 8, 2026 16:38

skyw merged commit a459fc1 into main Oct 8, 2026
5 checks passed

skyw deleted the skyw/shampoo_spectrum_scale branch October 8, 2026 16:42

Conversation

skyw commented Oct 5, 2026

Uh oh!

copy-pr-bot Bot commented Oct 5, 2026

Uh oh!

greptile-apps Bot commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Findings

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

skyw Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

skyw commented Oct 7, 2026

Uh oh!

copy-pr-bot Bot commented Oct 7, 2026

Uh oh!

skyw commented Oct 7, 2026

Uh oh!

copy-pr-bot Bot commented Oct 7, 2026

Uh oh!

skyw commented Oct 7, 2026

Uh oh!

copy-pr-bot Bot commented Oct 7, 2026

Uh oh!

skyw commented Oct 7, 2026

Uh oh!

copy-pr-bot Bot commented Oct 7, 2026

Uh oh!

greptile-apps Bot Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

skyw commented Oct 8, 2026

Uh oh!

github-actions Bot commented Oct 8, 2026

Test Results

Uh oh!

codecov Bot commented Oct 8, 2026

Codecov Report

Uh oh!

skyw commented Oct 8, 2026

Uh oh!

copy-pr-bot Bot commented Oct 8, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

greptile-apps Bot commented Oct 5, 2026 •

edited

Loading