OLMoE

论文：OLMoE: Open Mixture-of-Experts Language Models
地址：https://arxiv.org/pdf/2409.02060
权重：https://hf.co/allenai/OLMoE-1B-7B-0924
数据：https://hf.co/datasets/allenai/OLMoE-mix-0924
代码：https://github.com/allenai/OLMoE
日志：https://wandb.ai/ai2-llm/olmoe/reports/OLMoE-1B-7B-0924--Vmlldzo4OTcyMjU3

一、简介

OLMoE 是一种完全开放且最先进的语言模型，利用稀疏的 MoE（Mixture-of-Experts）。OLMOE-1B-7B 模型在相似活跃参数数量的情况下，表现优于所有可用的模型，甚至超过了更大的模型如 Llama2-13B-Chat 和 DeepSeekMoE-16B。(如果有和 sota 结果相比实验更有说服力了)
贡献：现有的 MoE 模型大多是闭源的，缺乏训练数据、代码和训练方法的开放资源，这限制了研究的进展和开源 MoE 模型的发展。OLMoE 开放了论文、代码、数据以及实验过程和结果，保证了广大研究者们能复现的可能。

二、模型细节

1、架构

OLMOE 是一个 decoderonly 的模型，与传统 MoE 差不多，将 dense 块中前馈网络（FFN）用 MoE 模块替代，MoE 模块由多个较小的专家模块(小型 FFN，如果是 lstm 之类的怎么样？)组成，每个输入标记只激活其中的一个子集。

2、关键模块设计

专家：使用 64 个小专家，每层激活 8 个。
路由：采用无 dropout 的标记选择路由。
损失：使用负载平衡损失（LLB）和路由器 z 损失（LRZ）来优化模型的训练。

三、数据与实验

1、数据

预训练使用 DCLM 和 Dolma 1.7 的数据集，其中包括 Common Crawl、StarCoder、arXiv 等等。

2、训练过程

从头开始训练，总共训练 5.133T 标记。在训练的退火阶段，先重新打乱整个数据集，然后线性衰减学习率至 0。在预训练过程中，使用多种下游任务进行评估，包括 MMLU、ARC-C、BoolQ 等。适应训练过程中，使用指令调优（SFT）和偏好调优（DPO）来提升模型性能。

3、实验参数

预训练时，优化器为 AdamW，混合精度，学习率为 5.0E-4
SFT 时，使用 BF16,batch_size 为 128，DPO 时使用 RMSProp，batch_size 为 32。

四、实验结果