论文

让我们逐步扩展:大规模专家混合的计算高效的超参数传输

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

模型训练预训练

摘要

专家混合 (MoE) 架构显着扩展了模型容量,而计算成本却没有成比例增加。然而,通过扫描在模型大小和 词元预算 的极端规模下优化其超参数(尤其是学习率)在计算上仍然令人望而却步。在本文中,我们提出了一种计算效率高的两步超参数传输框架,该框架通过跨缩放模型宽度传输模型来估计训练大型 MoE 模型的最佳学习率,然后外推到万亿词元的范围。首先,我们利用多头潜在注意力(MLA)和 Muon 优化器为 MoE 架构制定最大更新参数化($μ$P)适应,证明最佳学习率在宽度模型上一致地转移。其次,我们通过建立预测缩放定律来沿词元维度扩展这种可转移性。通过将线性回归应用于有限预算下的小型代理模型得出的最佳值,我们成功地将理想学习率外推到高保真度的大规模训练范围(例如 10 万亿个词元)($R^2=0.95$)。因此,这表明小模型上的代理训练足以确定大规模 MoE 的广泛训练的最佳学习率。我们应用所提出的方法从头开始预训练我们的基础模型(总共 155B 个,17B 个活动参数),稳定的训练和评估结果验证了可以以最小的消融成本准确预测全尺寸目标模型的最佳配置。