📖标题:GRIN: GRadient-IN formed MoE
🌐来源:arXiv, 2409.12136
摘要
🔸由于通过专家路由进行稀疏计算,选择性地仅激活一小部分专家模块,混合专家(MoE)模型比密集模型更有效地扩展。然而,稀疏计算挑战了传统的训练实践,因为离散专家路由阻碍了标准的反向传播,从而阻碍了基于梯度的优化,而这是深度学习的基石。
🔸为了更好地追求MoE的缩放能力,我们引入了GRIN(GRadient IN形成的MoE训练),它结合了稀疏梯度估计用于专家路由,并配置了模型并行性以避免令牌丢弃。将GRIN应用于自回归语言建模,我们开发了一个top2 16×3.8B MoE模型。我们的模型只有6.6B的激活参数,其性能优于7B密集模型,并与在相同数据上训练的14B密集模型的性能相匹配。
🔸对不同任务的广泛评估表明,GRIN具有显著提高MoE疗效的潜力,在MMLU上达到79.4,在HellaSwag上达到83.7,在HumanEval上达到74.4,在MATH上达到58.9。
🛎️文章简介
🔸研究问题:在大规模预训练模型中,如何通过混合专家(MoE)模型提高模型可扩展性和训练效率?
🔸主要贡献:论文提出了一种新的MoE模型GRIN,并开发了相应的训练技术,显著提升了模型的训练效率和性能。
📝重点思路
🔺相关工作
🔸大规模预训练的成功凸显了模型可扩展性的重要性,MoE已成为一种有前途的方法。
🔸通过专家路由选择性地仅激活特定输入上的一小部分模块,从而将模型可扩展性提高几个数量级。
🔸MoE的稀疏激活机制给模型训练带来了一些挑战,如反向传播不能直接应用于专家路由的梯度计算。
🔺论文方案
🔸模型设计:设计了名为GRIN的MoE模型,通过混合专家架构来提高模型的可扩展性。
🔸梯度估计:提出了SparseMixer-v2方法,通过离散变量采样和Heun’s第三阶方法来估计专家路由的梯度,解决了传统MoE模型中梯度不可微的问题。
🔸模型并行:采用了数据并行、流水线并行和激活检查点技术,避免了专家并行和容量因子使用,提高了训练效率。
🔸实验设计:进行了控制实验和半控制实验,比较了GRIN MoE与传统GShard方法以及不同规模的密集模型在多个任务上的性能。
🔎分析总结
🔸SparseMixer-v2在自回归语言模型训练中表现出色,尤其是在训练后期,性能优于传统的GShard方法。
🔸GRIN MoE模型在代码和数学任务中表现尤为突出,例如在MMLU、HumanEval、MBPP和MATH任务中取得了优异的成绩。
🔸GRIN MoE在训练吞吐量方面也表现出显著的优势,相对于具有相同活动参数的密集模型,训练效率提高了80%以上。
🔸尽管GRIN MoE在自然语言任务上的表现相对较弱,但在推理、编码和数学任务上的表现优于许多开源模型。
💡个人观点
论文的核心是提出了一种新的梯度估计方法SparseMixer-v2,并结合新的模型并行策略,显著提升了MoE模型的训练效率和性能。
附录