Add Muon Optimizer to contrib

[leloykun]

Adds support for @KellerJordan's Muon optimizer as detailed here: https://kellerjordan.github.io/posts/muon/

This optimizer allowed us to reduce the training time of a GPT-2 level model down to just ~3.5 minutes on 8xH100s. See: https://github.com/KellerJordan/modded-nanogpt

This implementation also supports the adaptive version of Muon that can adapt to the scale of the gradients as they change during training. To enable this feature, simply set muon(adaptive=True, ...). See more at: https://github.com/leloykun/adaptive-muon

---

Muon is equivalent to Shampoo without accumulation. I.e., take Shampoo's update rule, discard the left and right preconditioners, then simplify the matrix operations. Alternatively, Muon can also be thought of as steepest descent under Spectral norm. For more details, please refer to @jxbz's paper https://arxiv.org/pdf/2409.20325.

More generally, steepest descent under Schatten-p norm, for some large p, can be thought of as variants of Muon (note that Schatten-infty norm = Spectral norm). A (rough) proof of this claim can be found here: https://x.com/leloykun/status/1846842887839125941. This implementation supports these variants by allowing the user to pass in custom coefficients for the newton-schulz iterations.

[leloykun]

Hi all!

How do I restrict the tests to exclude vectors as weights?

[vroulet]

Hello @leloykun,

Looks like an interesting contribution thank you!

How do I restrict the tests to exclude vectors as weights?

I don't fully understand your question. Could you give more context?

Other questions/remarks:

How does this optimizer treat vector-shaped values? The"muon_iterator" could run on vectors but not return what you want, so how do you make the distinction? Should it take a mask to only apply on matrices? Should it raise value error?

Also could you add the mathematical description in the doc string? That would greatly help.

Finally, put references at the end of the docstring (we'll adopt this format with #1129)

Thank you!

[leloykun]

Hi @vroulet,

Muon is only defined for matrix-shaped values.

I'm thinking of raising an error when the input is vector-shaped, but where's the best place to put it? If there are other optimizers here that does this, can you point me to them?

[vroulet]

Hello @leloykun

• I would put a test in the init function.
• What is the algorithm for? Most architectures have non-matrix shaped parameters (for example just the scalar parameters of a layer norm). If the algorithm is specific to some problem, could you make a quick notebook explaining the setup? (Is it for example for finetuning with LoRa?)
• In any case, some doctest would be good.
• For the tests, you may add a synthetic linear classification problem with centered data to test the algorithm for example (again if you have a specific setup let me know).

Thanks again and sorry for the delay!

[vroulet]

Hey @leloykun,

Looking at the original repository, it seems that the best course of action is to:

• keep the scale_by_muon as you did
• make muon as a partition instance where the partition is inferred from the dimensions of the leaves in the optimizer
• add in the tests a simple mlp and make all optimizers pass it

Please let me know if you are willing to pursue the PR

[rdyro]

Hi @leloykun, I'll be happy to take over merging this PR on our end, let me know once you want me to take a look at the changes / start the merge process.

[leloykun]

Hi @rdyro! Thank you so much for the help. This PR should now be ready for review. I've also already addressed @vroulet's concerns. Please let me know if there is anything else I need to do on my end.

[rdyro]

Can you introduce these cosmetic changes? The PR looks really good!

[rdyro]

Hi @rdyro! Thank you so much for the help. This PR should now be ready for review. I've also already addressed @vroulet's concerns. Please let me know if there is anything else I need to do on my end.

Hey, I just reviewed the PR, it looks great, thanks for addressing @vroulet's comments. I left some cosmetic comments.

Once you get those in, can you squash your commits please?

[rdyro]

I'm sorry, I mislead you. Because we're supporting Python 3.9 Union and Tuple are necessary, exactly like you had it (having Optional would be ok too, but you can leave it as Union[A, None]). I'll start the merge once you can revert those Union/Tuple changes.

[leloykun]

@rdyro, I've just addressed your nits and squashed the commits.

I kept the Union & Optional type hints tho as support for | was only added in v3.10 instead of v3.9... See: https://peps.python.org/pep-0604/