Adaptive learning rate for Muon: NorMuon and AdaMuon #76
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
skyw
requested changes
Nov 13, 2025
Contributor
There was a problem hiding this comment.
How many will need this 2nd momentum other than OrthogonalizedOptimizer? Looks like none (or very few) because Adam based optimizer won't use it. In which case, should inherit from OrthogonalizedOptimizer and override step() function.
The choices are between copying code and keeping adding more functionalities to one class. In this case, I think copying some code in step() is the best trade off.
One another need to consider is if it turned out to be not very useful, how painful to remove the code, which would also lead to inheriting and override step function.
emerging_optimizers/orthogonalized_optimizers/orthogonalized_optimizer.py
Outdated
Show resolved
Hide resolved
mkhona-nvidia
commented
Nov 13, 2025
emerging_optimizers/orthogonalized_optimizers/adaptive_orthogonalized_optimizer.py
Outdated
Show resolved
Hide resolved
skyw
requested changes
Nov 13, 2025
Contributor
skyw
left a comment
skyw
left a comment
There was a problem hiding this comment.
I'm also debating whether we should define a class for 2nd momentum and passing an object, that way argument list can be better managed. Can to it in further PR when opinion form.
emerging_optimizers/orthogonalized_optimizers/adaptive_orthogonalized_optimizer.py
Outdated
Show resolved
Hide resolved
emerging_optimizers/orthogonalized_optimizers/adaptive_orthogonalized_optimizer.py
Outdated
Show resolved
Hide resolved
emerging_optimizers/orthogonalized_optimizers/adaptive_orthogonalized_optimizer.py
Outdated
Show resolved
Hide resolved
emerging_optimizers/orthogonalized_optimizers/adaptive_orthogonalized_optimizer.py
Outdated
Show resolved
Hide resolved
emerging_optimizers/orthogonalized_optimizers/adaptive_orthogonalized_optimizer.py
Outdated
Show resolved
Hide resolved
mkhona-nvidia
commented
Nov 18, 2025
skyw
requested changes
Nov 18, 2025
Contributor
skyw
left a comment
skyw
left a comment
There was a problem hiding this comment.
Some minor changes before we can merge, otherwise LGTM.
mkhona-nvidia
force-pushed
the
adaptive_orthogonalized_optimizer
branch
from
9b86617 to
a9f8d2c
Compare
skyw
reviewed
Nov 18, 2025
Contributor
Author
|
/ok to test 75aa10b |
skyw
previously approved these changes
Nov 18, 2025
Contributor
Author
|
/ok to test c720455 |
Signed-off-by: mikail <[email protected]>
…d second moment computations within same code Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
This reverts commit 9d9ddf2. Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
This reverts commit 66f9196. Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
mkhona-nvidia
force-pushed
the
adaptive_orthogonalized_optimizer
branch
from
d4bc2ec to
cc72e34
Compare
Contributor
Author
|
/ok to test cc72e34 |
Signed-off-by: mikail <[email protected]>
Signed-off-by: mikail <[email protected]>
Contributor
Author
|
/ok to test a677a3a |
skyw
approved these changes
Nov 18, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
This MR extends Muon to support adaptive learning rates for Muon-like optimizers, namely NorMuon (https://arxiv.org/pdf/2510.05491) and AdaMuon (https://arxiv.org/pdf/2507.11005).
Normuon recently was incorporated into the speedrun record (https://github.com/KellerJordan/modded-nanogpt/blob/master/train_gpt.py#L595)