Optional auxiliary learning rate for Adam within Moun

[RaphaelRe]

#1563
This PR adds an option for an additional learning rate for the tree leaves that are handled by Adam(W) in the Muon API.

I also added a line to highlight, that the API uses AdamW (and not Adam) under the hood with a weight decay of 0. A minor thing, but may impact performance if one is not aware of it while using a positive weight decay for Adam. This is not directly written in the docstrings and can only be seen either in the code or in the contrib examples.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[rdyro]

Great, thank you!

I think you need to sign the contributor agreement (CLA), but I can merge then

[RaphaelRe]

Yep. Forgot this. Should be signed now.
Maybe you can rerun the check if necessary.

[RaphaelRe]

Sorry, it looks like there was a delay with the agreement or I missed something. Now the CLA is checked as signed.

[rdyro]

Can you try to force push or push an empty commit? I think this should refresh the github CI

[RaphaelRe]

Made an empty commit. Maybe you can resolve it now.
If not, I can make a new pull request.
Sorry for the issue. The next time, I will be a bit more patient. Promised 😅

[rdyro]

No worries, thanks for the PR :)