Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix Muon implementation #1170

Open
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

leloykun
Copy link
Contributor

@leloykun leloykun commented Jan 9, 2025

Apply orthogonalization on momentum term instead of the original updates.

This should now match the original implementation by @KellerJordan


Postmortem: There was an operation before this line that takes in mu_hat and outputs updates. Thus, this line made sense then. I have since removed that operation. And thus, this line should now take in mu_hat instead. My fault for not double-checking.

What we should do so this never happens again: we should add harder toy benchmarks so mistakes like this would be spotted more quickly.

Tagging @rdyro for review.

@rdyro
Copy link
Collaborator

rdyro commented Jan 9, 2025

No worries, thanks for the quick fix!

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