Most hyperparameter searches for massive MoE pretraining are computationally prohibitive because optimal settings (especially learning rate) shift with both model width and token budget. The core insight of this paper is that width and token dimensions can be decoupled: μP-style parameterization yields learning-rate transferability across width, and a simple log-log linear law lets you extrapolate the optimal rate across token budgets from small proxy runs.
Key Findings
- μP adaptation for MoE (with Multi-head Latent Attention and the Muon optimizer) makes optimal learning rates robustly transferable across width-scaled proxy models, eliminating the need for large-scale width sweeps.
- Optimal learning rates across token budgets follow a predictable trend in log-log space; linear regression on proxy-derived optima achieves R² ≈ 0.95, enabling reliable extrapolation to very long training horizons.
- Applying the two-step method, the authors predict an optimal learning rate of 3.85×10⁻⁴ for a 155B-total (17B active) MoE model trained to 10 trillion tokens and validate stable training and competitive downstream results.
- The proxy search compute cost is small relative to full pretraining (proxy runs used ≈98× less total compute than the full pretraining run used for validation), making the approach practical for resource-constrained labs.
Who it's for and tradeoffs
Great fit if you need a principled, low-cost way to set the learning rate for large MoE pretraining: researchers and engineers planning wide MoE models or very long token budgets who can run small proxy experiments. Look elsewhere or plan extra experiments if you need guaranteed global optimality (the method targets learning-rate transfer and extrapolation rather than full multi-dimensional HP sweeps), if your architecture or optimizer differs substantially from the paper's setup, or if routing/load-balancing hyperparameters require joint tuning. The extrapolation assumes stable training dynamics after initial budget/scale-induced transients, so early-phase instability or radically different data curricula may reduce prediction fidelity.