Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

A two-step framework for selecting the learning rate of very large Mixture-of-Experts models without running prohibitively expensive full-scale sweeps. It first adapts Maximal Update Parameterization (μP) to MoEs using Multi-head Latent Attention and the Muon optimizer, transferring the optimal learning rate across model widths and expert counts. It then fits a log-linear scaling law over token budgets to extrapolate from short proxy runs to trillion-token training horizons.

The authors predict the learning rate for a 10T-token run with R² = 0.95; held-out tests near 500B tokens differ from the measured optimum by about 4.4% on average. They validate the recipe by stably pretraining an unnamed 155B-total / 17B-active MLA MoE from scratch on 10T tokens using NVIDIA H200 GPUs and an internal Megatron-LM fork. The work was led at Kakao; co-author Hojin Lee is affiliated with Upstage AI and notes that his contribution was completed at Kakao.