文章标题:DeepSeekMoE:迈向混合专家语言模型中的极致专家特化
作者/机构:Damai Dai∗1,2, Chengqi Deng1, Chenggang Zhao∗1,3, R.X. Xu1, Huazuo Gao1, Deli Chen1, Jiashi Li1, Wangding Zeng1, Xingkai Yu∗1,4, Y. Wu1, Zhenda Xie1, Y.K. Li1, Panpan Huang1, Fuli Luo1, Chong Ruan1, Zhifang Sui2, Wenfeng Liang1
1DeepSeek-AI 2北京大学多媒体信息处理全国重点实验室 3清华大学交叉信息研究院 4南京大学软件新技术全国重点实验室
在大语言模型时代,混合专家(MoE)架构是扩展模型参数同时控制计算成本的有前途的方案。然而,传统的MoE架构(如GShard)在激活N个专家中的前k个时,面临确保专家特化(即每个专家获得不重叠且专注的知识)的挑战。为应对这些挑战,本文提出了DeepSeekMoE架构,旨在实现极致的专家特化。
核心问题:
现有的MoE架构可能存在知识混合和知识冗余的问题,这限制了专家特化。
1. 知识混合(Knowledge Hybridity):现有的MoE实践通常使用有限数量的专家(如8或16个),导致分配给特定专家的token可能涵盖多样的知识。因此,指定的专家会试图在其参数中汇集截然不同类型的知识,而这些知识很难被同时利用。
2. 知识冗余(Knowledge Redundancy):分配给不同专家的token可能需要共同的知识。结果,多个专家可能会在其各自的参数中学习共享的知识,从而导致专家参数的冗余。
研究目标与创新点:
为了解决上述问题,本文引入了DeepSeekMoE,这是一种创新的MoE架构,旨在实现极致的专家特化。该架构包含两个主要策略:
1. 细粒度专家切分(Fine-Grained Expert Segmentation):在保持参数数量不变的情况下,通过缩小FFN中间隐藏维度,将专家切分成更细的粒度。相应地,在保持计算成本不变的情况下,激活更多细粒度的专家,以实现更灵活和自适应的激活专家组合。这使得多样化的知识能够被更精细地分解,并被更精确地学习到不同的专家中。
2. 共享专家隔离(Shared Expert Isolation):隔离一部分专家作为共享专家,这些专家总是被激活,旨在捕获和巩固跨不同上下文的通用知识。通过将通用知识压缩到这些共享专家中,可以减轻其他路由专家的冗余,从而提高参数效率,并确保每个路由专家专注于独特方面,保持特化。
主要贡献总结:
* 架构创新:引入了DeepSeekMoE,一种旨在实现极致专家特化的创新MoE架构,采用了细粒度专家切分和共享专家隔离两大策略。
* 实证验证:通过广泛的实验,验证了DeepSeekMoE架构的有效性。实验结果证实了DeepSeekMoE 2B具有高度的专家特化水平,并表明其性能可以接近MoE模型的性能上限。
* 可扩展性:将DeepSeekMoE扩展到16B模型,并证明其仅用约40%的计算量就达到了与DeepSeek 7B和LLaMA2 7B相当的性能。同时,初步尝试将DeepSeekMoE扩展到145B,验证了其相较于GShard架构的持续优势,并表现出与DeepSeek 67B相当的性能。
* MoE对齐:成功对DeepSeekMoE 16B进行了监督微调,创建了一个对齐的聊天模型,展示了其适应性和多功能性。
* 公开发布:本着开放研究的精神,公开发布了DeepSeekMoE 16B的模型检查点。该模型无需量化即可部署在单张40GB显存的GPU上。
通用Transformer架构。一个标准的Transformer语言模型由堆叠的$L$个标准Transformer块构成,每个块的计算过程可以表示如下:
$\mathbf{u}_{1:T}^{l} = \text{Self-Att}\left(\mathbf{h}_{1:T}^{l-1}\right) + \mathbf{h}_{1:T}^{l-1},$
其中,$T$表示序列长度,Self-Att(·)表示自注意力模块,FFN(·)表示前馈网络(FFN),$u_{1:T}^l \in \mathbb{R}^{T \times d}$是第$l$个注意力模块后所有token的隐藏状态,而$h_t^l \in \mathbb{R}^d$是第$t$个token经过第$l$个Transformer块后的输出隐藏状态。为简洁起见,上述公式中省略了层归一化。
MoE语言模型构建。构建MoE语言模型的典型做法通常是按指定间隔用MoE层替换Transformer中的FFN层【【14, Glam: Efficient scaling of language models with mixture-of-experts, 2022, ICML】; 【16, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021, CoRR】; 【33, Gshard: Scaling giant models with conditional computation and automatic sharding, 2021, ICLR】; 【70, Designing effective sparse expert models, 2022, IPDPS Workshops】】。一个MoE层由多个专家组成,每个专家的结构与标准FFN相同。然后,每个token将被分配给一个【16】或两个【33】专家。如果第$l$个FFN被替换为MoE层,其输出隐藏状态$h_t^l$的计算表示为:
架构设计。在第2节概述的通用MoE架构基础上,本文引入了DeepSeekMoE,它专门设计用于发掘专家特化的潜力。如图2所示,该架构融合了两大核心策略:细粒度专家切分和共享专家隔离。这两种策略都旨在提升专家特化的水平。
动机与实现。在专家数量有限的情况下,分配给特定专家的token更有可能涵盖多样化的知识类型。结果,指定的专家会倾向于在其参数中学习差异巨大的知识,而这些知识很难被同时利用。然而,如果每个token可以被路由到更多的专家,多样化的知识就有可能被分解并分别在不同的专家中学习。在这种情况下,每个专家仍然可以保持高水平的专家特化,有助于在专家之间形成更专注的知识分布。
具体方法。为了实现这一目标,在保持专家参数数量和计算成本一致的情况下,本文将专家切分得更细粒度。更细的专家切分使得激活的专家组合更加灵活和自适应。具体来说,在图2(a)所示的典型MoE架构之上,通过将FFN中间隐藏维度减小到原始大小的$1/m$倍,将每个专家FFN切分为$m$个更小的专家。由于每个专家变小了,相应地,也将激活的专家数量增加到$m$倍,以保持相同的计算成本,如图2(b)所示。通过细粒度专家切分,MoE层的输出可以表示为:
组合灵活性。从组合学的角度来看,细粒度专家切分策略显著增强了激活专家的组合灵活性。举例来说,考虑$N=16$的情况。典型的top-2路由策略可以产生 $\binom{16}{2}=120$ 种可能的组合。相比之下,如果每个专家被分成4个更小的专家($m=4$),细粒度的路由策略可以产生 $\binom{64}{8}=4,426,165,368$ 种潜在组合。组合灵活性的激增,增强了实现更精确和更有针对性的知识获取的潜力。
动机与实现。在使用传统路由策略时,分配给不同专家的token可能需要一些共同的知识或信息。结果,多个专家可能会在其各自的参数中学习共享的知识,从而导致专家参数的冗余。然而,如果有专门的共享专家来捕获和巩固跨不同上下文的通用知识,那么其他路由专家之间的参数冗余将会减轻。这种冗余的减轻将有助于构建一个参数效率更高、专家更特化的模型。
具体方法。为了实现这一目标,在细粒度专家切分策略的基础上,进一步隔离$K_s$个专家作为共享专家。无论路由模块如何决策,每个token都将被确定性地分配给这些共享专家。为了保持计算成本不变,其他路由专家中激活的专家数量将减少$K_s$,如图2(c)所示。整合了共享专家隔离策略后,完整的DeepSeekMoE架构中的一个MoE层公式如下:
与前期工作的区别。值得注意的是,共享专家隔离的原型可以追溯到Rajbhandari等人(2022)【42, Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation AI scale, 2022, ICML】。关键区别在于,他们是从工程角度得出这一策略,而本文是从算法角度来处理它。
负载不均衡问题。自动学习的路由策略可能会遇到负载不均衡的问题,这表现出两个显著的缺陷。首先,存在路由崩溃(routing collapse)的风险【58, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017, ICLR】,即模型总是只选择少数几个专家,导致其他专家得不到充分的训练。其次,如果专家分布在多个设备上,负载不均衡会加剧计算瓶颈。
专家级平衡损失(Expert-Level Balance Loss)。为了减轻路由崩溃的风险,本文也采用了一种专家级的平衡损失。该平衡损失的计算如下:
设备级平衡损失(Device-Level Balance Loss)。除了专家级平衡损失,本文引入了设备级平衡损失。当目标是缓解计算瓶颈时,不必在专家级别强制执行严格的平衡约束,因为对负载平衡的过度约束会损害模型性能。相反,主要目标是确保跨设备的计算平衡。如果将所有路由专家划分为$D$个组$\{E_1, E_2, ..., E_D\}$,并将每组部署在单个设备上,设备级平衡损失计算如下:
模型架构:
硬件配置:NVIDIA A100或H800 GPU集群。节点内通过NVLink/NVSwitch连接,节点间通过InfiniBand连接。
软件配置:
训练配置:
模型架构:
硬件配置:同验证实验。
实验结果:
分析结论:DeepSeekMoE实现了更高程度的专家特化,其路由专家冗余度低,共享专家功能独特,并且能以更高的效率(更少的激活专家)获取知识。
实验结果:
分析结论:DeepSeekMoE架构在大规模下依然高效,能以显著更低的计算成本达到甚至超越强大的密集模型的性能。
MoE技术的发展。混合专家(MoE)技术最早由Jacobs等人(1991)【27, Adaptive mixtures of local experts, 1991, Neural Computing】和Jordan与Jacobs(1994)【28, Hierarchical mixtures of experts and the EM algorithm, 1994, Neural Computing】提出,用于通过独立的专家模块处理不同的样本。Shazeer等人(2017)【58, Outrageously large neural networks: The sparsely-gated mixture-of-experts layer, 2017, ICLR】将MoE引入语言模型训练,并构建了基于LSTM的大规模MoE模型。随着Transformer成为NLP最流行的架构,许多工作尝试将Transformer中的FFN扩展为MoE层来构建MoE语言模型。
主流MoE架构。GShard【33, Gshard: Scaling giant models with conditional computation and automatic sharding, 2021, ICLR】和Switch Transformer【16, Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2021, CoRR】是先驱,它们采用可学习的top-2或top-1路由策略,将MoE语言模型扩展到极大。Hash Layer【44, Hash layers for large sparse models, 2021, CoRR】和StableMoE【12, Stablemoe: Stable routing strategy for mixture of experts, 2022, ACL】使用固定的路由策略以获得更稳定的路由和训练。Zhou等人(2022)【69, Mixture-of-experts with expert choice routing, 2022, NeurIPS】提出了一种专家选择路由策略,其中每个token可以被分配给不同数量的专家。Zoph(2022)【70, Designing effective sparse expert models, 2022, IPDPS Workshops】专注于MoE模型的训练不稳定和微调困难问题,并提出ST-MoE来克服这些挑战。
与本文工作的关系。除了对MoE架构和训练策略的研究外,近年来还涌现了许多基于现有MoE架构的大规模语言或多模态模型【【14, Glam: Efficient scaling of language models with mixture-of-experts, 2022, ICML】; 【35, M6: A chinese multimodal pretrainer, 2021, CoRR】; 【43, Pangu-Σ: Towards trillion parameter language model with sparse heterogeneous computing, 2023, CoRR】; 【67, Openmoe: Open mixture-of-experts language models, 2023, GitHub】】。总的来说,之前的大多数MoE模型都基于传统的top-1或top-2路由策略,在提升专家特化方面留下了很大的空间。作为回应,本文的DeepSeekMoE架构旨在最大程度地提高专家特化。
本文介绍了DeepSeekMoE架构,旨在实现MoE语言模型中的极致专家特化。通过细粒度专家切分和共享专家隔离,DeepSeekMoE相比于主流MoE架构,实现了显著更高的专家特化和性能。
* 验证与优势:从2B参数的小规模模型开始,验证了DeepSeekMoE的优势,证明了其有能力接近MoE模型的性能上限。此外,提供了实证证据表明DeepSeekMoE比GShard具有更高水平的专家特化。
* 规模化与应用:将模型扩展到16B参数,在2T token上训练的DeepSeekMoE 16B,仅用约40%的计算量就展现出与DeepSeek 7B和LLaMA2 7B相当的卓越性能。此外,通过监督微调成功构建了MoE聊天模型,进一步展示了其适应性和多功能性。
* 未来展望:初步探索将DeepSeekMoE扩展到145B参数,发现其依然保持对GShard架构的巨大优势,并以28.5%(甚至可能低至18.2%)的计算量实现了与DeepSeek 67B相当的性能。
* 开源贡献:为了研究目的,公开发布了DeepSeekMoE 16B的模型检查点,该模型可部署在单张40GB显存的GPU上。希望这项工作能为学术界和工业界提供有价值的见解,并为大规模语言模型的加速发展做出贡献。
下表7展示了不同规模DeepSeekMoE的超参数概览。
表8展示了DeepSeekMoE与GShard×1.2和GShard×1.5的比较。表9展示了DeepSeekMoE与Dense×4和Dense×16的比较。
在更大的13B总参数规模上,本文也比较了DeepSeekMoE与GShard×1.2和GShard×1.5,结果如表10所示。在更大规模上,DeepSeekMoE甚至明显优于GShard×1.5。
图7展示了DeepSeekMoE 16B和DeepSeek 7B(Dense)在训练过程中的基准曲线,以供参考。