Muon

[eliebak]

What does this PR do?

draft of the draft don't mind

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guidelines?
• Did you write any new necessary tests?
• Did you log the throughput and loss you get to ensure the PR works as expected in actual training?
• Did you log the memory usage? you can use this tool to understand the memory usage breakdown in nanotron.
• If you modified anything related to checkpoints, did you verify that saving and reloading checkpoints still works correctly?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.