【LLM MoE Pretrain】论文分享:DeepSeekMoE

341 阅读9分钟

论文名称:DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

论文链接:arxiv.org/abs/2401.06…

机构:DeepSeek AI + 北大 + 清华 + 南大

简介

DeepSeekMoE这篇论文,早在2024年1月就发布出来了,当时毕竟是国内首个开源MoE大模型,而且效果确实不错。我们团队内部很快就阅读了技术报告,并进行了一些讨论,现在过去快一年半了,市面上确实也出现越来越多的MoE架构模型,证明了这种架构的潜力。

现在回头再看这篇论文,在记录关键内容的同时,也谈谈自己的认识。

相关常识

MoE 是 DeepSeek 提出的吗?

MoE的全称是 Mixture-of-Experts,这个架构不是DeepSeek提出来的,早在1991年就有相关论文提出该概念了,附上论文里面贴的三篇论文:

①【1991】Adaptive Mixtures of Local Experts

②【1994】Hierarchical mixtures of experts and the EM algorithm

③【2017】Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

MoE 是 LLM 的专属架构吗?

不是,严格来说,是属于神经网络架构的一种。目前为止,LLM的基础架构是基于Transformer不断改进的,Transformer也是一种神经网络的结构。

DeepSeek 第一个应用了 MoE 架构?

不是,在其之前已经有多个知名模型采用了MoE架构了,比如Mistral AI的Mixtral 8x7B、Google的GShard与Switch Transformer等,OpenAI的GPT-4疑似也采用了MoE架构。

DeepSeek MoE 是第一个开源的 MoE 大模型?

国际上首个开源的MoE大模型是Mixtral 8x7B(2023年12月发布),DeepSeek MoE是国内首个开源的MoE大模型,是否为闭源第一,不确定。

DeepSeek 全系列都是 MoE 大模型?

不是,MoE(DeepSeek MoE 16B) 和 Dense(DeepSeek 67B) 都有。

通用 MoE 架构介绍

标准Transofmer语言模型

标准Transformer语言模型是通过堆叠L个标准Transformer模块层来构建的,其中每个模块可以表示如下:

主要是表达了两个逻辑:

  • L层注意力模块之后所有token的隐藏状态 = L-1层Transformer块之后的输出隐藏状态 + L-1层Transformer块之

后的输出隐藏状态的自注意力块状态

  • L层Transformer块之后的输出隐藏状态 = L层注意力模块之后所有token的隐藏状态 + L层注意力模块之后所有

token的隐藏状态的前馈神经网络输出状态

注:类似下图的描述,思想来源于残差网络的设计,能使网络层数堆叠的更深。

MoE 的常见做法

构建MoE语言模型的一种常见做法是在Transformer的指定间隔处用MoE层替换FFNs

MoE层由多个专家组成,每个专家在结构上与标准FFN相同。然后,每个Token将被分配给一个或两个专家。如果第L个FFN被MoE层替换,其输出隐藏状态的计算表示为:

其中s表示Token到专家的亲和力(token-to-expert affinity),每个Token都被分配到一个或多个专家,而亲和力表示了某个Token与各个专家之间的关联程度或选择概率。

DeepSeek MoE 架构介绍

整体框架

图2是DeepSeekMoE的示意图,展示了三个子图,分别说明了不同的MoE层架构。

  • 子图(a):展示了具有传统Top-2路由策略的MoE层。在这种情况下,每个Token被分配给两个专家中的一个,即Top-2。这是传统的MoE路由策略,其中每个Token只与两个专家相关。
  • 子图(b):说明了Fine-grained expert segmentation策略。相比于传统的Top-2路由,DeepSeekMoE采用了更细粒度的专家划分,将专家进一步分为多个子专家。这样,每个Token可以与更多的专家相关,实现更灵活的激活专家的组合。
  • 子图(c):展示了Shared expert isolation策略。在这种情况下,一些专家被标记为共享专家,旨在捕捉共同知识并减轻激活专家之间的冗余。这种共享专家隔离策略有助于提高模型的性能和效率。

Fine-grained expert segmentation 策略

如何保持计算成本不变?通过减少每个专家的大小,来提高专家的数量

Shared expert isolation策略

如何保持计算成本不变?假设每次选择Top-K个专家,其中Ks个专家作为共享专家,剩下K-Ks个作为激活专家

Load Balance Consideration Loss

自动学习的路由策略可能会遇到负载不平衡的问题,这表现为两个值得注意的缺陷:

  • 路由崩溃的风险:模型总是只选择少数专家,从而阻止其他专家进行充分训练。
  • 计算瓶颈:如果专家分布在多个设备上,负载不平衡会加剧计算瓶颈。

所以,提出了两个Loss来解决这个问题:

① Expert-Level Balance Loss

专家级别的平衡损失有助于防止路由崩溃。

② Device-Level Balance Loss

设备级别的平衡损失用于确保在设备之间的平衡计算。

实验部分

数据处理

训练数据是从DeepSeek-Al创建的大规模多语言语料库中采样的。语料库主要侧重于英语和中文,但也包括其他语言。它来源于多种来源,包括网络文本、数学材料、编码脚本、已发表的文献和各种其他文本材料。出于验证实验的目的,作者从语料库中抽取一个包含100B个标记的子集来训练模型

对于Tokenization,作者使用HuggingFace Tokenizer,在训练语料库的较小子集上训练byte pair encoding tokenizer。在验证实验中,作者准备了一个词汇量为8K的tokenizer,在训练更大的模型时,词汇量会放大。

训练框架

实验基于HAI-LLM框架进行,这是一个高效且轻量级的训练框架,整合了多种并行策略,包括张量并行、ZeRO数据并行、PipeDream pipeline并行以及专家并行通过结合数据和张量并行。作者为了优化性能,使用CUDA和Triton为门控算法开发了GPU内核,并在不同专家的线性层之间进行计算融合

所有实验在配备A100或H800的集群上进行。A100集群中的每个节点包含8个GPU,通过NVLink成对连接。H800集群还具有每个节点8个GPU,节点内部使用NVLink和NVSwitch相互连接。对于A100和H800集群,InfiniBand互连用于节点间的通信。

超参设置

作者在附录A中提供了不同大小的DeepSeekMoE超参数概述表。

实验结果

① 【表1】在验证实验中,对比Dense、Hash Layer、Switch Transformer、GShard等模型,具有相同总参数和激活参数的DeepSeekMoE在各项任务指标上表现优异,展现出相对于其他MoE架构的显著性能优势。

② 【表2】对比了DeepSeekMoE与更大规模的GShard×1.5和Dense×16模型,结果显示DeepSeekMoE能与专家参数和计算量1.5倍于自身的GShard×1.5达到相当的性能,且接近参数为其16倍的Dense×16模型的性能,体现了DeepSeekMoE架构的优势以及其性能接近MoE模型理论上限。

③【图3】通过对DeepSeekMoE进行消融实验,在保证对比模型总参数和激活参数相同的情况下,证明了细粒度专家分割和共享专家隔离策略均有助于提升模型整体性能,且在总专家数和激活专家数不变时,共享专家与路由专家的比例对性能影响不显著

④ 【图4】通过对比DeepSeekMoE和GShard×1.5在禁用不同比例顶级路由专家后的Pile损失,表明DeepSeekMoE对禁用顶级路由专家更敏感,其路由专家间冗余度更低,且共享专家不可被路由专家替代,对模型性能至关重要。

⑤【图5】在 DeepSeekMoE 中,激活 4 个路由专家时就能达到与 GShard 相当的 Pile 损失,表明其能以更少激活专家更准确高效地获取知识。

⑥ 【图6】在总专家参数相同且激活专家参数仅为一半的情况下,DeepSeekMoE 的性能仍优于 GShard,体现其对专家参数的利用更高效。

论文里面还有更多实验分析,不多赘述了,感兴趣的可以详细阅读原文。

总结

通过DeepSeekMoE论文,结合这一年多来的经历,有以下一些感受与思考。

MoE 架构是否值得长期坚持?

值得。

① 首先拿DeepSeek-V3举例,其671B总参数中仅激活37B即可实现性能对标传统Dense模型,算力需求却能降低至1/20,这种“参数规模与计算成本解耦”的特性在算力稀缺的2025年更加有实际意义。

② 很多企业已经能将 MoE 模型应用于金融风控、数字人直播等高并发场景,推理吞吐量能够提升3.2倍以上,验证了 MoE 架构在商业化场景中的可行性。

③ 2025年出现的专家链(CoE)通过专家间串行通信进一步优化推理效率,也能说明 MoE 仍处于创新活跃期。

DeepSeeK MoE 对后续研究的影响?

比如 Hybrid-MoE 就借鉴了其专家划分思路,通过残差MoE模块组合实现17倍计算节省。

再后面的阿里云QwQ-32B等模型采用“共享+任务专家”双层结构,减少了通用知识冗余的问题。

有无比MoE更好的架构去替代它?

目前看尚无颠覆性的架构出现,但混合架构可能有一定的潜力,比如Hybrid-MoE、Hybrid-Mamba等。

MoE 架构有没有什么缺点?

最大的挑战依然在端侧的部署,毕竟也需要加载全部专家的参数。

其次就是门控网络的波动,可能也会导致相同输入激活不同专家,对于一些稳定性要求较高的场景(如医疗、法律)会有点吃力。

剩下的就是,负载均衡不好做,而且在一些数据稀缺的领域(如小语种翻译),专家如果特化不足就会导致性能下降。