通过两步迁移,用小模型预测大模型最优学习率。
Let's Scale Step by Step: Compute-Efficient Hyperparameter Transfer for Large-Scale Mixture-of-Experts

- 用宽度缩放模型迁移学习率,实现高效参数配置。
- 在10万亿令牌训练中预测准确度达R²=0.95。
- 适合大规模MoE模型训练,节省大量调参成本。
Mixture-of-Experts(MoE)架构能在不显著增加计算成本的前提下大幅提升模型容量。然而,在模型规模和令牌预算均极高的情况下,通过网格搜索优化超参数(尤其是学习率)仍存在计算开销过大问题。本文提出一种计算高效的两步超参数迁移框架:首先利用多头隐式注意力(MLA)与Muon优化器,基于最大更新参数化(μP)在宽度缩放的模型间实现学习率稳定迁移;其次沿令牌维度拓展迁移能力,建立可预测的缩放规律。通过对小规模代理模型在有限预算下的最优学习率进行线性回归,成功将理想学习率外推至万亿级训练规模(如10万亿令牌),预测精度高达R²=0.95。结果表明,仅通过小模型代理训练即可准确预测大规模MoE全规模训练的最优配置。我们基于该方法从零开始预训练了一个基础模型(总计1550亿参数,活跃参数170亿),训练稳定且评估结果验证了配置预测的准确性,仅需极少的消融实验即可获得可靠设置。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。