计算最优并非集群最优：稀疏专家混合的系统感知扩展

在大规模预训练中，算法、架构和系统决策通常是在断开的阶段中做出的。缩放法则阶段选择架构和训练配方，在计算约束下优化损失，然后一个单独的系统阶段优化硬件效率的实现。在这项工作中，我们开发了 MOSAIC，它将模型架构和系统协同设计制定为优化问题。 MOSAIC 将预测缩放定律与校准性能模型结合起来，该模型可估计模型 FLOP 利用率 (MFU)、通信成本、内存占用和最佳并行布局。我们实例化了稀疏专家混合 (MoE) 语言模型的框架，其中专家数量、路由稀疏性和其他 MoE 层维度都会影响损失和系统效率。我们在文本数据上训练的稀疏 MoE 模型上拟合了缩放法则，其缩放维度包括稀疏因子，即前向传递中每个标记不活动的模型参数的比例。缩放定律在我们的工作中涵盖了从 1.04 亿美元到 27 亿美元的活跃参数，总模型大小达到 79 亿美元的参数。我们表明，在校准的稀疏度范围内，与效率无关的模型 FLOPs 预算不承认内部最优稀疏度。拟合损失随着稀疏模型单调减少，并且计算最优值位于数据支持的上边界。相反，MoE 模型中的最佳稀疏性是在集群的系统约束下出现的，如 MOSAIC 所捕获的那样。我们的结果支持前沿语言 模型训练 向统一架构和系统协同设计的转变。