微软：LLM在MoE架构的高效训练策略
微软：LLM在MoE架构的高效训练策略
微软最新研究提出了一种名为GRIN（GRadient IN formed MoE）的混合专家（MoE）模型训练策略，通过创新的梯度估计方法和模型并行技术，显著提升了LLM的训练效率和性能。
摘要
由于通过专家路由进行稀疏计算，选择性地仅激活一小部分专家模块，混合专家（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模型的训练效率和性能。