Source-linked AI summary
OpenMoE: An Early Effort on Open Mixture-of-Experts Language Models
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, Yang You
TL;DR
MoE LLMs 能够降低计算压力，但其专家的行为与专业化仍缺乏充分理解。OpenMoE 训练并发布了可复现的开源 MoE 模型，同时分析其 routing，验证了其有效性，并发现三种 routing 现象：Context-Independent Specialization、Early Routing Learning 和 Drop-towards-the-End。
问题
LLMs 计算成本高昂，而 MoE 专家的专业化与 routing 行为仍缺乏充分理解。
方法
论文训练并发布了可复现的开源 decoder-only MoE LLMs，并对其 routing mechanisms 进行了深入分析。
结果
研究验证了基于 MoE 的 LLMs 的有效性，并识别出 Context-Independent Specialization、Early Routing Learning 和 Drop-towards-the-End。
要点与局限性
这些发现与提出的解决方案为未来的开源 MoE LLM 开发者提供了有价值的见解。
要点与局限性
当专家达到容量上限时，序列中较晚出现的 tokens 面临更高的丢弃风险；在 instruction-tuning 数据中，这一问题更为严重。
Abstract
from arXiv · showhide
To help the open-source community have a better understanding of Mixture-of-Experts (MoE) based large language models (LLMs), we train and release OpenMoE, a series of fully open-sourced and reproducible decoder-only MoE LLMs, ranging from 650M to 34B parameters and trained on up to over 1T tokens. Our investigation confirms that MoE-based LLMs can offer a more favorable cost-effectiveness trade-off than dense LLMs, highlighting the potential effectiveness for future LLM development. One more important contribution of this study is an in-depth analysis of the routing mechanisms within our OpenMoE models, leading to three significant findings: Context-Independent Specialization, Early Routing Learning, and Drop-towards-the-End. We discovered that routing decisions in MoE models are predominantly based on token IDs, with minimal context relevance. The token-to-expert assignments are determined early in the pre-training phase and remain largely unchanged. This imperfect routing can result in performance degradation, particularly in sequential tasks like multi-turn conversations, where tokens appearing later in a sequence are more likely to be dropped. Finally, we rethink our design based on the above-mentioned observations and analysis. To facilitate future MoE LLM development, we propose potential strategies for mitigating the issues we found and further improving off-the-shelf MoE LLM designs.
1 引言
OpenMoE 通过发布开源 MoE 模型来应对 LLM 的计算成本问题，同时研究路由机制与先进训练策略。该研究还分析了路由局限性，并据此为未来 MoE 发展提出方向。
- 发布 OpenMoE: OpenMoE 发布了涵盖调试、预训练和对话版本的开源 MoE LLM，包括一个每个 token 激活约 2B 参数、使用超过 1T tokens 训练的 8B 参数模型。发布的系列包括 OpenMoE-Base/16E、OpenMoE-8B/32E 和 OpenMoE-8B/32E-Chat。
- 探索先进训练策略: OpenMoE 通过在早期预训练中使用最高 52.25% 的代码数据，并研究 UL2 目标，探索超越传统实践的训练方式。这些策略突破了以文本为主的数据和下一 token 预测训练范式。
- 深入研究 MoE 路由: 该工作研究稀疏路由，因为 MoE 必须将每个 token 分配给少数专家，才能在扩展可训练参数的同时保持计算量不变。引言指出，稀疏专家选择对 MoE 的功能至关重要，并指出其行为仍缺乏充分研究。
- 深入研究 MoE 路由: 当专家达到容量上限时，序列后部的 token 面临更高的丢弃风险；在领域不同于预训练数据的指令微调数据上，这一问题会进一步加剧。早期预训练建立了平衡的 token 分配策略，但这些策略可能无法有效迁移到指令微调场景。
- 反思错误并提出潜在解决方案: 作者利用实证发现和路由可视化结果，重新审视项目中的次优决策，并为未来 MoE 发展提出潜在解决方案。除成功策略外，他们还将分享这一早期开源尝试中的错误与经验视为一项重要贡献。
2 OpenMoE 设计
OpenMoE 结合多语言分词、稀疏交错 MoE 架构、UL2 预训练和有限的有监督微调，构建可扩展且具备指令能力的模型。其设计强调稳定路由、专家负载均衡和降低路由开销。
- 分词: 模型采用 256K umT5 tokenizer，通过多语言词表和 byte fallback 支持低资源语言及词表外 token。相比词表更小的替代方案和不具备 byte fallback 的旧 tokenizer，该 tokenizer 更为合适。
- 架构与路由: OpenMoE 采用 token-choice top-2 routing，每个 expert 实现为 FFN，而不是完整的 Transformer。gating function 使用稀疏 softmax routing 从 E 个 expert 中选择 K=2 个，从而实现条件计算。
- 架构与路由: Residual MoE 为每个 token 保留一个始终激活的固定 FFN，同时每隔 4 或 6 个 Transformer layer 交错插入 MoE layer，以降低路由开销。OpenMoE-Base/16E 和 OpenMoE 34B/32E 每隔 4 个 layer 使用 MoE，OpenMoE-8B/32E 则每隔 6 个 layer 使用 MoE。
- 训练目标: 训练将负载均衡损失与 router z-loss 相结合，通过惩罚较大的 gating logit 来改善专家并行并稳定路由。z-loss 鼓励减小 logit 的绝对值，并减少 MoE layer 中的舍入误差。
- 训练目标: OpenMoE 探索 UL2，将 span corruption 与 prefix language modeling 相结合，并增加 PrefixLM，仅使用 20% 的低 mask ratio 训练。该配置优先采用 PrefixLM，以支持 zero-shot 和 in-context learning；同时，较低的 mask ratio 可避免因输出 token 更少而导致学习速度变慢。
- 有监督微调: OpenMoE 对 WildChat 中由 GPT-4 生成的 58K 段对话进行有监督微调，以增强指令遵循能力，并比较 SFT 前后的 MoE 行为。由于训练后期的计算资源受限，该项目将 SFT 数据限制为 GPT-4 生成的指令—响应对。
3 OpenMoE 训练
OpenMoE 的训练设计与 MoE 扩展性提升了其在多个基准上的成本效益和性能，但也暴露出 UL2 饱和、代码数据占比过高以及多轮 token 丢弃等局限。这些结果支持 MoE 的可扩展性，同时表明训练目标和数据构成需要谨慎调整。
- 消融实验: MoE 方法、UL2 目标、提高代码数据占比以及 LLaMA tokenizer 均有助于提升 zero-shot TriviaQA 性能。尽管词表更大可能带来不利影响，使用 LLaMA tokenizer 的模型仍优于使用 umT5 tokenizer 的模型。
- 消融实验: 代码领域训练更易于优化：模型能够达到更高的准确率和更低的损失，并且尽管模型规模较小，在 GitHub 上的 token 预测准确率超过 80%。论文认为，这部分源于 token 频率的长尾分布，其中换行和制表符 token 出现频率较高。
- 训练目标: UL2 能加速早期学习，但后期趋于饱和，因此 OpenMoE-8B/32E 在 390K steps 后切换至 CasualLM，而 OpenMoE-34B/32E 使用 UL2 训练 25K steps。在 780B tokens 之后出现的准确率下降，反映了难度更高的 CasualLM 目标和更难学习的代码数据，而由易到难的课程学习仍可能有用。
- 基准结果: 在 13 个 LM-Evaluation-Harness 任务上，OpenMoE 和 TinyLLaMA 的表现不及 OpenLLaMA，这可能是因为初期较高的代码数据采样比例损害了以文本为主的基准性能。尽管如此，作者认为 OpenMoE 的分数仍可接受，并将这一数据混合问题留待进一步讨论。
- 基准结果: 在单轮 MTBench 上，OpenMoE 的表现显著优于基线，尤其是在代码任务上；但第二轮的性能下降，导致多轮结果更差。论文认为，这种性能下降可能是由 token 丢弃造成的。
4 OpenMoE 分析
OpenMoE 的路由主要由 token 身份和位置决定，而非广泛的领域或编程语言语义；模型较早形成 specialization，且基本不依赖上下文。容量限制进而造成 Drop-towards-the-End 效应，而在更细粒度的输入中则可观察到语言和任务层面的 specialization。
- 领域与语言 specialization: OpenMoE 在不同领域或编程语言之间几乎没有明显的 expert specialization，包括低资源的 Assembly 和 Blitzmax，以及 Java 和 Python。各领域子集在 experts 之间大致均匀分布，仅表现出轻微偏好，例如 E21 偏向代码，E10 偏向书籍。
- 领域与语言 specialization: 自然语言路由表现出更清晰的 specialization：中文偏向 E5 和 E16，而日语和韩语偏向 E14。这些发现来自对 TED-Parallel-Corpus 中 12 种语言的分析。
- 任务与位置 specialization: MT-Bench 对话路由显示出类似的更细粒度 specialization，尤其是在数学数据上，这可能是因为数学任务包含更多 special tokens。Position IDs 也表现出 specialization，相邻位置倾向于选择相似的 experts，例如 E10 和 E19。
- Context-Independent Specialization: Token IDs 强烈决定路由：相同 tokens 即使出现在多样化的上下文中，也只偏向少数几个 experts，体现出 Context-Independent Specialization。Experts 还会将低层语义相似的 tokens 聚集在一起，例如 “can”、“will” 和 “would”。
- Early Routing Learning: 不同中间 checkpoints 之间的路由偏好几乎完全重叠，表明 token-to-expert assignments 在预训练的 very early 阶段就已学得并固定。Supervised fine-tuning 并未显著改变 Drop-towards-the-End 问题，这与路由早期固定的行为一致。
- Drop-towards-the-End: 当较早的 tokens 填满 expert capacity 后，较晚的 tokens 更容易被丢弃，从而在 decoder-only MoE models 中产生 Drop-towards-the-End 效应。该问题源于预先设定的 per-expert capacity 和自回归处理；除 OpenMoE 外，在 Mixtral 和 Deepseek-MoE 中也观察到了这一现象。
5 重新思考 OpenMoE
Section 5 总结了 OpenMoE 分析带来的设计启示，并提出未来改进方向，同时指出昂贵的消融实验以及对模型规模和数据质量的依赖所带来的不确定性。相关建议涵盖代码-数据构成、tokenizer 选择、高效 MoE 架构，以及尽早混合 instruction data，以改善路由平衡。
- 概述: 本节将这些建议作为作者从自身失误和分析洞见中得到的主要启示，旨在为未来实践者提供指导。作者明确承认，针对部分设计选择尚无精确答案，并毫无保留地报告了这些发现。
- 代码构成: 作者建议采用约 30% 的代码数据，正如 OpenMoE-34B/32E 所采用的那样，因为超过 50% 的代码数据可能损害文本任务能力。作者提醒，具体结论可能高度依赖模型规模和数据质量，而大规模消融实验成本高昂。
- Tokenizer 选择: 定量 tokenizer 分析发现，在多语言数据集上，umT5 远优于 LLaMA；尽管会增加输出层计算开销，这一结果仍支持更智能地选择 tokenizer。Table 10 使用激活词表量，以及不同数据集子集上的 umT5/LLaMA token 数量比，对这些 tokenizer 进行了比较。
- 更高效的 MoE 架构: 作者提出在 warmup 后移除可训练 router，并将 Transformer 与 FFN 计算并行化，同时将 attention 计算与 MoE 的 all-to-all 通信重叠。这些改动旨在提高硬件利用率；在扩展规模时，并行 Transformer 计算可以实现重叠，且不会造成性能下降 [9]。
- 负载均衡与 Drop-towards-the-End: 作者提出在预训练 warm-up 阶段混入 instruction-following 数据，以改善负载均衡并缓解 Drop-towards-the-End；这一点在多轮 MT-Bench 中尤为重要。其依据是，MoE 路由行为会在预训练早期被学习并固定下来。
6 结论 · 附录 · A 常见问题
本研究证明了开放 MoE LLM 的有效性与可复现性，同时识别出三种路由现象，并为未来开发者提出解决方案。研究旨在增进开源社区对 MoE 模型的理解，并针对潜在的常见问题提供逐点回答。
- 6 结论: 该研究验证了 MoE LLM 在 ChatGPT 之后阶段的有效性。
- 6 结论: 由于作者公开了全部细节以及开源代码和数据，所发布的模型完全可复现。
- 6 结论: 分析识别出 Context-independent Specialization、Early Routing Learning 和 Drop-towards-the-End 三种重要的 MoE 路由现象。
- 6 结论: 作者承认其 MoE 设计存在不足，并为未来开发者提出可能的解决方案。
- 6 结论: 本研究旨在帮助开源社区更好地理解 MoE 模型。
- A 常见问题: 附录列出了潜在的常见问题，并提供逐点回答。
A.1 为什么不展示预热阶段检查点的 token specialization？…… A.4 为什么不使用 AdamW optimizer 和 Cosine Learning Rate Schedule？
附录解释了 OpenMoE 中与检查点、比较、架构和优化相关的选择，强调可复现性、训练效率和稳定性约束。附录还指出，routing 的固定时间早于预期，而 AdamW 会频繁导致不稳定。
- A.1 为什么不展示预热阶段检查点的 token specialization？: Routing 的学习和固定时间早于预期，但由于存储空间有限，检查点仅每 200B tokens 保留一次。这种检查点保存频率使得无法展示预热阶段的 token specialization。
- A.2 为什么不与 Mixtral 和 DeepSeek-MoE 等先进的开源 MoE 模型进行比较？: OpenMoE 的发布公告和正式发布都比 Mistral 早 over four months，比 DeepSeek-MoE 早 more than four months。
- A.2 为什么不与 Mixtral 和 DeepSeek-MoE 等先进的开源 MoE 模型进行比较？: 不同于使用内部数据训练的模型，OpenMoE 完全透明，公开了全部细节和代码，能够从头训练出具有可比性的模型。
- A.3 为什么不使用 MoE upcycling？: 相比 inference，MoE 在 training 阶段更高效，因为较大的 batch size 能带来更好的并行性。
- A.3 为什么不使用 MoE upcycling？: 在 dense LLM 上构建 MoE 是更快的构建路径，但从长期看并不更高效；当性能损失较小时，将 MoE distilling 到 dense model [53] 可能有所帮助。
- A.4 为什么不使用 AdamW optimizer 和 Cosine Learning Rate Schedule？: 遵循 ST-MoE ，OpenMoE 使用 Adafactor 和 Inverse Square Root learning-rate schedule，而不是 AdamW 和 cosine decay。
- A.4 为什么不使用 AdamW optimizer 和 Cosine Learning Rate Schedule？: 之所以不采用 AdamW，是因为它经常产生不稳定的 NaN losses，而且在计算资源有限的情况下需要进行大量超参数搜索。
A.5 为什么不使用更好、更大的数据集？
项目启动时，可用的开源预训练数据集寥寥无几，限制了数据规模和质量。此后，更大且经过更仔细清洗的数据集相继出现，这表明未来改进训练数据可能显著提升 LLM 性能。
- 数据可用性: 2023年五月，项目启动时，只有少数开源预训练数据集可供使用。这限制了项目初期可供选择的方案。
- 数据可用性: Soldaini et al. 发布了 3T 个经过仔细清洗的 tokens，而 Computer 总共发布了 30T 个 tokens。这些例子体现了开源预训练数据在规模和质量上的后续提升。
- 影响: 作者预计，使用未来更好的数据进行训练，通常可以大幅提升 LLM 性能。这表明数据集质量和规模是未来改进的重要方向。
B 超参数 · C 相关工作 · C.1 OpenMoE 之前的工作
附录说明了 OpenMoE 的配置与训练选择，包括一种注重成本的优化器以及一个值得注意的 head dimension 限制，并将该模型置于 MoE 语言、视觉、路由和 scaling 研究工作的脉络中。
- B 超参数: OpenMoE 的模型配置定义了 hidden size、MoE-layer 布局、FFN size、attention dimensions、层数、总参数量和 activated parameters。“Every 4”表示每四个 Transformer blocks 配置一个 MoE layer；activated-parameter 的统计涵盖处理每个 token 时使用的参数。
- B 超参数: OpenMoE-8B/32E 使用 head dimension 128，不过对于其 2B activated Transformer parameters，dimension 64 可能带来更好的成本效益权衡。参数量统计通过纳入可训练的 attention- 和 FFN-layer 参数来简化计算，这些参数构成了 Transformer block 参数的大部分。
- B 超参数: OpenMoE 的训练超参数总结于一张专门的配置表中。所给段落指出了该表，但未提供其中的具体数值。
- B 超参数: Adafactor 取代 AdamW，因为其内存效率能够减少模型并行、增加数据并行，从而使报告配置下的训练成本更低。在训练步数相同的情况下，Adafactor 的表现略逊于 AdamW，而成本差距高度取决于硬件和模型规模。
- C.1 OpenMoE 之前的工作: 早期 MoE 语言模型工作从 Shazeer et al. [42] 的循环式 MoE layers，发展到 GShard [27] 和开源 Switch Transformer 中面向 Transformer-scale 的 expert parallelism。Switch Transformer 采用 encoder-decoder 架构，并在 C4 上训练。
- C.1 OpenMoE 之前的工作: 既有工作还通过更大的数据集和高效架构扩展 MoE 训练，包括 Artetxe et al. [3]、GaLM 和 Brainformer 。Artetxe et al. [3] 使用了 112B tokens，GaLM 使用了 1.6T tokens，而 Brainformer 搜索了 layer interleaving 和 capacity 等 MoE 属性。
- C.1 OpenMoE 之前的工作: 超越语言领域的 MoE 扩展包括用于可扩展视觉 Transformers 的 ViT-MoE 、提升参数效率的 WideNet ，以及用于稳定 soft routing 的 SoftMoE [35]。SoftMoE [35] 在保留效率的同时采用 soft token selection，并提升 routing-gradient 的稳定性。
C.2 OpenMoE 之后
OpenMoE 发布后，Mixtral [22] 和 Deepseek-MoE [12] 等部分开源的 MoE 模型取得了显著更好的最终结果。然而，由于它们使用内部训练数据，我们无法了解这些结果是如何获得的，这凸显了 OpenMoE 完全开源的特性。
D BigBench-Lite 结果
Table 14 报告了详细的 BigBench-Lite 结果，并明确 BIG-G-sparse 8B 是一个总参数量为 60B 的 MoE 模型。
- Table 14 提供了详细的 BigBench-Lite 结果。
- BIG-G-sparse 8B 是一个总参数量为 60B 的 MoE 模型。
E 路由决策标准差 · F 各专家选择的高频 Token
与 position IDs 相比，路由决策在 token IDs 间的差异更大，表明路由主要不依赖上下文。本节还汇总了各专家选择的高频 token。
- E 路由决策标准差: Figure 15 报告了不同 position IDs 间的路由决策标准差。该图考察了路由变异性如何随序列位置变化。
- E 路由决策标准差: Figure 16 报告了不同 token IDs 间的路由决策标准差，并将分析限制在 token 数量超过 128 的 IDs 上。由于资源极少的 token IDs 始终具有较大的标准差，因此将其排除。
- E 路由决策标准差: 从未出现过的 token IDs 也表现出路由决策方差。该观察结果被纳入 token-ID 分析。
- E 路由决策标准差: Token IDs 的路由决策标准差大于 position IDs，表明路由与 token identity 的关联更强。该比较结果来自 Figures 15 和 16。
- E 路由决策标准差: 大多数 token IDs 具有相对较大的路由决策标准差，支持对大多数 token 采用不依赖上下文的路由。该分析将 token IDs 间观察到的变异性解释为不依赖上下文的路由证据。
- F 各专家选择的高频 Token: Table 15 列出了各专家选择的高频 token。所给段落说明了该表的范围，但未提供具体 token 条目。