研究人员提出了一种方法，能通过小模型实验，准确预测万亿 token 规模 MoE 模型的最优学习率，省去了大量计算成本。他们用 155B 参数的模型验证了方法的有效性。
该研究提出一种计算高效的两步超参数迁移框架，用于为大型 MoE 模型估算最优学习率。首先，通过 MuP 适应和 Muon 优化器，证明了最优学习率在宽度缩放模型间可一致迁移。其次，通过建立预测性缩放定律，将小模型在有限预算下的最优值线性回归，成功外推至万亿 token 规模（R²=0.95）。该方法在 155B 总参数的 MoE 基础模型预训练中得到验证。
Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts
Mixture-of-Experts (MoE) architectures significantly expand model capacity without a proportional increase in computational cost. However, optimizing their hyperparameters---particularly the learning rate---at extreme scales of both model size and token budget via sweeping remains computationally prohibitive. In this paper, we propose a compute-efficient, two-step hyperparameter transfer framework that estimates optimal learning rates for training large MoE models by transferring them across scaling model widths, and subsequently extrapolating to trillion-token horizons. First, we formulate a Maximal Update Parameterization ($μ$P) adaptation for MoE architectures utilizing Multi-head Latent Attention (MLA) and the Muon optimizer, demonstrating that optimal learning rates transfer consistently across width-scaled models. Second, we extend this transferability along the token dimension by establishing a predictive scaling law. By applying linear regression to the optimal values derived from small proxy models on limited budgets, we successfully extrapolate the ideal learning rate to massive training horizons (e.g., 10 trillion tokens) with high fidelity ($R^2=0.95$). Consequently, this indicates that proxy training on small models is sufficient to determine the optimal learning rate for the extensive training of large-scale MoEs. We apply the proposed methodology to pretrain our foundation model (155B total, 17B active parameters) from scratch, and the stable training and evaluation results validate that optimal configurations for full-scale target models can be accurately predicted with minimal ablation costs.