Add Muon Optimizer to contrib

[leloykun]

Adds support for @KellerJordan Muon optimizer detailed here: https://github.com/KellerJordan/modded-nanogpt

The Muon optimizer does steepest descent under approximately Spectral norm--or more accurately, under Schatten-p norm for some large p (with Schatten-infty norm = Spectral norm).

[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!