improvements from Moonlight

[ehartford]

MoonshotAI suggested some improvements to Muon in training Moonlight.

https://github.com/MoonshotAI/Moonlight/blob/master/examples/toy_train.py

Perhaps they would make a good addition to Muon?

[toothacher17]

hi, @ehartford thanks for mentioning it! To be fair, I think the core ideas of training Moonlight are:

1. weight decay
2. adjusting update rms by matrix shape
3. matching to AdamW RMS

For point 1 and 2, Keller's current impl should already contain it (as we mentioned in the paper) under the setting of nanogpt. For point 3, it is mostly designed for large scale over-train setting and might not be the best setting under the nanogpt speedrun (small scale of model, small scale of tokens)

[lin72h]

@toothacher17 Nice suggestion!