Skip to content
Bud's Wiki

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

TL;DR

Section titled “TL;DR”

Nouamane Tazi (Hugging Face) publishes a practical engineering guide to training Mixture-of-Experts language models with Megatron-LM, setting by setting. Written as a direct follow-on to The Smol Training Playbook (which stopped at dense), the handbook walks a reader from hardware tests through a dense baseline, up to a single-node MoE training run, then into the design choices inside an MoE layer — routers, load balancing, expert block layouts — with recommended settings and cross-references to how recent frontier MoEs (DeepSeek-V3, Kimi-K2, Qwen3-Next, GPT-OSS) set them. Final sections cover post-training serving and RL on the trained MoE. Framed for engineers who already trained dense LLMs and now need to cross to sparse.

Key claims

Section titled “Key claims”
  • Megatron-LM is the right framework to center the handbook on because it is battle-tested and powers production MoE runs including Kimi-K2 and Nemotron, and pioneered 3D parallelism [§Picking a training framework]. Fits the broader cluster the wiki already tracks: Fireworks’ production stack, PithTrain’s compact Megatron-compatible rewrite, and ByteDance’s veScale-FSDP all target the same Megatron-LM reference point.
  • The intended reader path is: (1) hardware characterization and GPU microbenchmarks, (2) a dense baseline on a single node, (3) a single-node MoE run with Megatron-LM, then only after those steps the inside-the-layer MoE design choices [§Structure].
  • Chapter 3 enumerates the inside-the-MoE-layer design axes as router + load balancing + expert block layout, with recommended settings grounded in recent MoE releases [§Chapter 3]. Matches the vocabulary the wiki’s moe-routing-design concept has been building: router topology (TopK vs sigmoid vs loss-free bias), auxiliary-loss scope, expert pool size vs active count, shared-expert isolation, fine-grained-vs-coarse experts.
  • Chapter 4 extends the scope past training into serving the trained MoE and running RL post-training on it — placing the handbook in the same “end-to-end recipe” slot as Fireworks’ trainer stack (Scaling and Optimizing Frontier Model Training (Fireworks AI)) rather than the “one slice of the pipeline” slot of a tech report [§Chapter 4].
  • Ships alongside a companion MoE Memory Estimator for Megatron-LM HF Space that predicts GPU memory from a training command or model-parameter form [§Companion tooling].

Method

Section titled “Method”

The handbook is structured as a four-chapter engineering sequence rather than a research survey. Chapter 1 frames the gap (Smol Training Playbook covered dense; MoE needs its own treatment) and picks Megatron-LM. Chapter 2 is a Megatron-LM setup walkthrough: hardware microbenchmarks → dense baseline training run → extending that baseline to a single-node MoE run. Chapter 3 is the architectural core — one sub-section per MoE-layer component (router, load balancing, expert block), each presenting the design axis, the recommended setting, and a short citation list of how recent MoEs set it. Chapter 4 extends past the training loop: serving the trained MoE and doing RL on top of it.

The accompanying MoE Memory Estimator for Megatron-LM HF Space lets a reader paste a training command (or fill in model parameters: hidden size, num_layers, FFN dim, num_experts, top_k, parallelism config, etc.) and get a predicted per-GPU memory footprint before launching.

Results

Section titled “Results”

Not applicable as written — the handbook is a guide rather than a research artifact with evaluation numbers. The underlying recommended settings reference measured results from the cited MoE releases (DeepSeek-V3 1-shared + 256 routed, top-8; Kimi K3 16-of-896; Qwen3-Next 512-expert; etc.) rather than running new ablations.

Why it’s interesting

Section titled “Why it’s interesting”

Fills a specific gap the wiki’s MoE coverage has been circling: the recipe-side companion to the design-space debates tracked in MoE Routing Design and the systems-side debates tracked in Distributed training parallelism. The wiki has load-bearing papers on individual routing design choices (Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models on LBL scope, Routing-Free Mixture-of-Experts on removing the router, Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method on hyperparameter-free balancing, Mitigate Silent Expert Death In Ultra-Sparse MoE on silent expert death) but the end-to-end “if you’ve trained dense before, do these things” is exactly what practitioners have been missing. Sits at the opposite end of the same curve as PithTrain: A Compact and Agent-Native MoE Training System (compact Python-native reimplementation) and Scaling and Optimizing Frontier Model Training (Fireworks AI) (production-scale 4D-parallelism stack): same Megatron-LM reference, different audience. Complements The MoE 101 Guide: From Theory to Production (Daria Soboleva, Cerebras)‘s theory-to-production guide from Cerebras with a Megatron-LM-specific implementation path.

See also

Section titled “See also”