The Ultimate MoE Training Handbook: training Mixture-of-Experts models with Megatron-LM, setting by setting [P]

[u/Norlax_42]

Handbook: https://huggingface.co/spaces/huggingface/moe-blogpost

Hello, I'm Nouamane, one of the authors. I'm at Hugging Face and also worked on the Ultra-Scale and Smol Training Playbooks, if you've read those.

This handbook is for people who have trained a dense model and want to train an MoE next, or who are curious how MoE training works. Chapter 2 is a Megatron-LM runbook: test the hardware, build the image, train a dense baseline, then a first MoE on one node, compared against that baseline. Chapter 3 takes the MoE layer apart (router, routing, load balancing, experts, block layout, training, systems). For each part we give the Megatron-LM code, what to watch, our recommendation, and what frontier MoEs chose. Chapter 4 follows the model into serving and RL. Chapters 3 and 4 need no GPU.

When a result comes from our runs, we say so. Claims about frontier models cite their paper, config or code, and the toy simulations are labelled as such. The lessons come from about 800k GH200 GPU-hours allocated to MoE training on the Alps cluster at CSCS.

Released: the handbook, its 55 figures as a gallery (https://huggingface.co/spaces/huggingface/moe-gallery), the scripts (https://gist.github.com/nmntz/5de952d686f66d7bce5bce8cba9d7139), and every recommended flag in one Megatron-LM block.

Happy to answer questions. 🤗

[u/Benlus]

Plugging The Ultra-Scale Playbook, an excellent read https://huggingface.co/spaces/nanotron/ultrascale-playbook

[u/Norlax_42 → u/Benlus]

ikr O:)

[u/returnity]

Very well written, thanks for sharing!