MoE 稀疏性的超参数缩放定律

专家混合 (MoE) 模型在训练计算量不成比例增加的情况下扩展了模型容量，但稀疏性的增加使可靠的超参数传输变得具有挑战性。在这项工作中，我们表明传统的超参数缩放定律对于超稀疏 MoE 来说是不够的：最佳学习率和批量大小随激活率而变化，并且这些变化不能仅用总参数数或激活参数数来解释。为了表征这种依赖性，我们进行了 1,800 次 预训练 运行，涵盖六个激活参数尺度和模型，总非嵌入参数高达 6B，处理大约 20 万亿个词元，成本为 200,000 个等效 H800 GPU 小时。我们的结果通过揭示两种缩放机制来调和先前工作中相互矛盾的发现。在固定稀疏度下，最佳批量大小遵循与训练标记 $D$ 的幂律关系，而最佳学习率随训练计算 $C$ 缩放，并且对模型大小和数据之间的分配保持鲁棒性。在稀疏级别上，激活比率 $A$ 作为附加的乘法幂律因子进入这两种关系。这些观察结果导致了统一的超参数缩放定律，可以在 MoE 稀疏度级别上传递。大规模评估表明，缩放形式优于替代函数形式。在具有 12B 总参数且仅激活 1/64 专家的超稀疏 MoE 上，预测的超参数仍然接近观察到的最优值，支持跨模型规模和稀疏性的联合外推。进一步的实验证明了跨专家粒度的传输，并将激活率的影响与专家总数的影响分开。