Llama 3 Meets MoE: Efficient Upcycling

📄 arXiv: 2412.09952v1 📥 PDF

作者: Aditya Vavre, Ethan He, Dennis Liu, Zijie Yan, June Yang, Nima Tajbakhsh, Ashwath Aithal

分类: cs.LG

发布日期: 2024-12-13

💡 一句话要点

利用Llama 3高效训练MoE模型：低成本实现性能提升

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 混合专家模型 MoE 大型语言模型 Llama 3 高效训练 模型升级 预训练模型 模型微调

📋 核心要点

  1. 现有LLM扩展面临高昂计算成本，MoE模型虽能提升容量，但从头训练存在过拟合和路由问题。
  2. 该论文提出利用预训练的稠密模型检查点，以极低的计算成本训练MoE模型，实现高效升级。
  3. 实验结果表明，该方法在MMLU上的0-shot准确率提高了2%，模型FLOPs利用率达到46.8%。

📝 摘要（中文）

扩展大型语言模型(LLMs)能显著提升性能，但计算成本也随之增加。混合专家模型(MoE)提供了一种高效的替代方案，能在不按比例增加计算需求的情况下提升模型容量。然而，从头开始训练MoE模型面临过拟合和路由不稳定等挑战。我们提出了一种高效的训练方法，利用预训练的稠密模型检查点，以低于典型预训练计算量1%的成本，从Llama 3-8B训练了一个8专家Top-2 MoE模型。我们的方法增强了学术基准测试的下游性能，在MMLU上的0-shot准确率提高了2%，并且在使用我们的框架训练时，模型FLOPs利用率(MFU)达到了46.8%。我们还在NeMo中集成了在线升级，以无缝使用预训练权重，从而经济高效地开发高容量MoE模型。

🔬 方法详解

问题定义：论文旨在解决从头训练MoE模型计算成本高昂且易出现过拟合和路由不稳定的问题。现有方法要么需要大量的计算资源，要么难以保证MoE模型的训练稳定性和性能。

核心思路：论文的核心思路是利用预训练的稠密模型（如Llama 3）作为MoE模型的初始化，通过微调的方式将稠密模型转化为MoE模型。这种方法避免了从头训练MoE模型的困难，显著降低了计算成本，并继承了预训练模型的知识。

技术框架：整体框架包括以下几个阶段：1) 选择预训练的稠密模型作为基础；2) 在稠密模型中引入MoE层，通常是在Transformer的FFN层；3) 使用少量数据和计算资源对MoE模型进行微调，优化专家路由和模型参数；4) 在下游任务上评估微调后的MoE模型性能。

关键创新：最重要的技术创新点在于利用预训练的稠密模型进行MoE模型的“upcycling”，即高效升级。与从头训练相比，这种方法极大地降低了计算成本，并能更快地获得高性能的MoE模型。此外，论文还集成了在线升级到NeMo，方便使用预训练权重。

关键设计：论文使用Llama 3-8B作为预训练的稠密模型，构建了一个8专家Top-2 MoE模型。这意味着每个输入token会路由到两个最相关的专家。训练时，使用了特定的学习率策略和正则化方法，以防止过拟合和保证路由的稳定性。具体参数设置和损失函数细节未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，使用该方法训练的MoE模型在MMLU基准测试上的0-shot准确率提高了2%，达到了显著的性能提升。更重要的是，训练过程中的模型FLOPs利用率(MFU)达到了46.8%，表明该方法具有很高的计算效率。与从头训练MoE模型相比，该方法所需的计算资源大幅减少，成本降低。

🎯 应用场景

该研究成果可广泛应用于需要高容量和高性能语言模型的场景，例如智能客服、内容生成、代码生成、机器翻译等。通过利用预训练模型进行MoE升级，可以降低模型训练成本，加速模型迭代，并促进更大规模MoE模型的发展。该方法也适用于其他模态的模型，例如图像和语音。

📄 摘要（原文）

Scaling large language models (LLMs) significantly improves performance but comes with prohibitive computational costs. Mixture-of-Experts (MoE) models offer an efficient alternative, increasing capacity without a proportional rise in compute requirements. However, training MoE models from scratch poses challenges like overfitting and routing instability. We present an efficient training recipe leveraging pre-trained dense checkpoints, training an 8-Expert Top-2 MoE model from Llama 3-8B with less than $1\%$ of typical pre-training compute. Our approach enhances downstream performance on academic benchmarks, achieving a $\textbf{2%}$ improvement in 0-shot accuracy on MMLU, while reaching a Model FLOPs Utilization (MFU) of $\textbf{46.8%}$ during training using our framework. We also integrate online upcycling in NeMo for seamless use of pre-trained weights, enabling cost-effective development of high-capacity MoE models.