Dist_Muon optimizer support

[BoxiangW]

Features in this PR

• Supports local muon
• Supports distributed muon
• Tested on Dense and MoE models

Need https://github.com/NVIDIA-NeMo/Emerging-Optimizers

Note: In this PR, there is a example script that uses a custom model (llama3 8b with 8 expert), please change the model structure if needed.

To use dist_muon or muon optimizer, please run

# on compute node
cd <workspace>
git clone https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git
git clone https://github.com/BoxiangW/Megatron-LM.git
cd Megatron-LM/
git checkout boxiangw/llm-shower-repo
cd ..

bash Megatron-LM/muon.sh

To use example script, please

1. Add your wandb api key and entity to the script (optional)
2. Add your tokenizer model path and data path to the script
3. Select your optimizer from [muon, dist_muon, adam], defaults to dist_muon

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[BoxiangW]

Closing this PR in favor of https://github.com/NVIDIA/Megatron-LM/tree/dev branch.

[valentyn1boreiko]

Hi @BoxiangW, thank you for the implementation! Going through the code, it doesn't really support distributed Muon, does it? Also the distrib_optimizer.py is still mostly for Adam and is not touched in this PR.

[BoxiangW]

Hi, we have merged this in the new dev branch and added TP support as well https://github.com/NVIDIA/Megatron-LM/blob/dev/megatron/core/optimizer/layer_wise_optimizer.py. TP: https://github.com/NVIDIA/Megatron-LM/blob/dev/megatron/core/optimizer/muon.py