OLMoE：开源混合专家语言模型
OLMoE: Open Mixture-of-Experts Language Models
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Walsh, Pete, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi
arXiv (Cornell University) · 2024-09-03
摘要
我们推出了OLMoE，这是一种完全开源、处于前沿水平的语言模型，采用了稀疏混合专家架构。OLMoE-1B-7B拥有70亿个参数，但每个输入词元仅调用10亿个参数。我们使用5万亿个词元对其进行了预训练，随后进一步微调得到OLMoE-1B-7B-Instruct。在活跃参数规模相近的模型中，该系列模型的表现均优于其他同类模型，甚至超过了Llama2-13B-Chat和DeepSeekMoE-16B等参数规模更大的模型。我们开展了多项关于混合专家架构训练的实验，分析了该模型的路由机制，发现其具备高度专业化特征；同时我们开源了研究的全部成果：模型权重、训练数据、代码及日志。
本摘要由 AI（Hy-MT2-Pro）翻译生成，仅供参考，请以原文为准。
AI 综述
基于全文生成- 主要观点
- 论文提出OLMoE，一个完全开放的稀疏Mixture-of-Experts语言模型系列，核心对象为OLMoE-1B-7B：总参数6.9B，每个输入token仅激活1.3B，经5.1T token预训练后通过指令与偏好调优得到OLMoE-1B-7B-Instruct。作者主张在相近激活参数预算下OLMoE达到最优，并超过Llama2-13B-Chat、DeepSeekMoE-16B等更大模型。模型设计采用每层64个小专家、激活8个，dropless token-choice路由，不使用共享专家与稀疏upcycling，并加入权重0.01的load balancing loss和0.001的router z-loss。路由分析显示路由在预训练早期饱和、专家很少共激活，并呈现领域与词汇专门化。论文同时开源模型权重、训练数据、代码与日志，强调开放MoE资源对理解复杂设计问题的重要性。
- 研究方法
- 方法上采用受控预训练实验与下游评测结合：在128张H100上以相同1.3B激活参数对比MoE与dense，MoE训练吞吐23,600 tokens/s/GPU、dense为37,500 tokens/s/GPU。围绕MoE设计逐一改变专家粒度、共享专家、expert-choice与token-choice、稀疏upcycling、负载均衡损失与router z-loss，并做数据集、初始化、RMSNorm、embedding参数衰减、QK-Norm、AdamW epsilon等通用设置实验。预训练数据为OLMoE-Mix，混合DCLM-Baseline、StarCoder、peS2o、arXiv、OpenWebMath、Algebraic Stack、Wikipedia/Wikibooks，共训练5.133T tokens；适应数据包括Tulu 2 SFT Mix、No Robots、CodeFeedback-Filtered-Instruction、MetaMathQA、Daring Anteater子集及UltraFeedback。评测覆盖MMLU、HellaSwag、ARC-Challenge、PIQA、WinoGrande、COPA、GSM8k、BBH、HumanEval、AlpacaEval、XSTest、IFEval等，并以router saturation、expert co-activation、domain specialization、vocabulary specialization分析路由行为。
- 学术贡献
- 贡献主要在开放资源与MoE设计实证两方面：将MoE的模型权重、训练数据、训练代码、日志和每5000步中间checkpoint一并开放，使MoE研究可在完整配方上复现与审计。通过控制实验给出若干可检验的设计结论：细粒度小专家优于少量大专家，dropless token-based routing优于expert-based routing，共享专家未带来收益，稀疏upcycling仅在较小计算预算下可能有益。性能上，OLMoE-1B-7B在少于2B激活参数模型中表现最佳，并在适应后平均分超过Qwen1.5-3B-14B chat等更大模型。路由分析贡献了关于专家专门化的经验刻画：路由早期饱和、专家共激活少、存在领域与词汇专门化，对开放MoE的可解释性与路由机制研究有价值。
- 存在局限
- 局限首先体现在配置选择：尽管按3e22计算预算的预测最优专家数为256，作者因收益递减与预测不适用而固定为64专家，未在该规模上验证更高粒度。部分结论如共享专家无效、稀疏upcycling收益有限，是在特定1.3B active/6.9B total、5T token与单变量控制设置下得到，材料承认超参众多，结果在不同配置下可能变化，且无法保证每个超参选择的正确性。实验之间因基础模型随成功发现更新而不可横向比较，削弱了跨实验结论的累积性。MoE相对dense仅实现约2倍训练加速，而非按token计算的约3倍，原因包含7B总参数带来的内存开销，材料未给出系统级优化后的加速上限。提供材料未展开附录H中限制与未来工作的具体内容，也未报告多语言、长上下文、推理延迟与能耗的系统测评。
- 进一步研究方向
- 后续可在相同1.3B active/6.9B total与5T token预算下，扫描8、32、64、256专家及相应激活数，测量MMLU、HellaSwag与路由饱和、共激活、领域与词汇专门化的变化，检验更高粒度是否仍只有1–2%增益。可针对MoE训练吞吐瓶颈做系统优化，比较不同内存布局、并行策略与kernel实现，验证能否把约2倍加速推进到约3倍token等价加速。可利用每5000步checkpoint做纵向路由追踪，定位专门化形成的时间点，并用专家掩蔽或替换做因果消融，检验领域与词汇专门化是否直接决定下游任务表现。可进行适应数据消融，改变Tulu 2 SFT Mix、CodeFeedback、MetaMathQA等代码与数学样本比例，验证GSM8k与HumanEval增益来源。还可在更大计算预算下重做共享专家与稀疏upcycling实验，检验当前否定性结论是否随预算改变。
本综述由 AI（deepseek-flash）阅读开放获取全文自动生成，仅供参考，不代表原文作者观点，请以原文为准。
参考文献
相关论文
话题：Topic Modeling