简介
论文《DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models》于 2024 年 1 月发布。它引入了细粒度专家分割和共享专家隔离策略,目标是让每位“专家”都能大展身手,而计算成本却丝毫不增,就像是一群精英各显神通却又默契配合的超级团队。
- 发布时间:2024 年 1 月
- 论文地址:arxiv.org/pdf/2401.06…
- Github地址:github.com/deepseek-ai…
摘要
传统 MoE 架构就像是一群“各怀绝技”的高手,但常常难以分工明确。DeepSeekMoE 的两大策略—— “细分” 与 “隔离” ——让每个专家都能专心致志地做自己的事。从小小的 2B 参数模型到后来的 16B 大腕,实验数据表明,DeepSeekMoE 无论在性能还是在专家专业化上,都远超传统架构,甚至还能与 LLaMA2 7B 相媲美,真是“少数精英,胜过一大群平庸”。
结论
实验表明,DeepSeekMoE 不仅在小规模上展现了惊人潜力,扩展到 16B 参数后仅用 40% 的计算量就达到了顶级表现。更别提对齐后的 MoE 聊天模型,各种对手只能望尘莫及。为了让更多人体验这款神器,团队还将 16B 版本的模型公开,让你在单个 40GB 显存的 GPU 上也能轻松部署!
同系列
论文详情如下:
摘要
在大语言模型时代,混合专家(Mixture-of-Experts,MoE)架构因能够在参数规模扩展时有效控制计算成本而被广泛关注。然而,传统的 MoE 架构(如 GShard,其采用从 N 个专家中激活前 K 个专家)在确保专家专门化方面存在挑战,即每个专家应学习不重叠且聚焦的知识。针对这一问题,我们提出了 DeepSeekMoE 架构,旨在实现终极的专家专门化。其主要包含两大策略:
- (1) 对专家进行精细划分(细分隔离),将原有专家细分为 mN 个,并激活其中 mK 个,从而实现激活专家组合的更高灵活性;
- (2) 将 Ks 个专家固定为共享专家,旨在捕捉共性知识并减少路由专家间的冗余。
从2B参数起步,我们证明 DeepSeekMoE 2B 在性能上可与拥有1.5倍专家参数及计算量的 GShard 2.9B 相媲美;同时,DeepSeekMoE 2B 几乎达到了同总参数下密集模型的上界。随后,我们将 DeepSeekMoE 扩展至 16B 参数,结果显示仅用约40%计算量即可实现与 LLaMA2 7B 相当的性能。此外,我们的初步尝试表明,当扩展至 145B 参数时,DeepSeekMoE 持续展现出对 GShard 架构的明显优势,并且其性能可与 DeepSeek 67B 媲美,仅需28.5%(甚至可能仅为18.2%)的计算量。
1. 引言
近年来的大量研究和实践表明,在充足训练数据支持下,通过增加参数量和计算预算扩展语言模型规模,往往能获得显著增强的模型能力(Brown 等,2020;Hoffmann 等,2022;OpenAI,2023;Touvron 等,2023a)。然而,模型极大规模扩展同时也伴随着极高的计算成本。为应对这一挑战,混合专家(MoE)架构(Jacobs 等,1991;Jordan 和 Jacobs,1994;Shazeer 等,2017)成为一种热门方案,它能够在扩展参数规模的同时保持计算成本在较低水平。近年来,MoE 架构在 Transformer 中的应用(Vaswani 等,2017)取得了不少成功案例(Du 等,2022;Fedus 等,2021;Lepikhin 等,2021;Zoph,2022),进一步证明了 MoE 模型在大规模模型训练中的潜力。
然而,现有的 MoE 架构可能面临知识混杂和知识冗余问题,从而限制了专家专门化,即每个专家学习到独特、聚焦的知识。常见做法是用 MoE 层替换 Transformer 中的前馈网络(FFN),每个 MoE 层包含多个与标准 FFN 结构相同的专家,每个 token 只分配给一位(Fedus 等,2021)或两位专家(Lepikhin 等,2021)。这种架构会引发以下两个问题:
- 知识混杂:由于专家数量有限(例如8或16个),分配到同一专家的 token 往往涉及多种知识,导致该专家参数中包含多种截然不同的知识,不易同时利用。
- 知识冗余:不同专家间的 token 可能需要共性知识,多个专家因此会在参数中学习相似的共性知识,造成冗余。
以上问题阻碍了现有 MoE 架构实现理论上 MoE 模型性能上界的专家专门化。
针对这些问题,我们提出了 DeepSeekMoE,这是一种全新设计的 MoE 架构,专注于实现专家专门化。该架构包含两大核心策略:
- 精细化专家划分:在保持专家参数总量不变的情况下,将每个专家 FFN 的中间隐层维度划分为 m 份,从而细化专家划分;同时,为保持计算成本恒定,将激活的专家数量提升至 mK,从而使激活专家组合更灵活多变。该策略使得多样化知识能更精细地分解到不同专家中,每个专家由此保留更高的专门化程度。
- 共享专家隔离:固定选择 Ks 个专家作为共享专家,不论路由模块如何,每个 token 均确定性地分配给这些共享专家,目的是捕捉和整合各上下文间的共性知识,从而减少其他路由专家中的冗余。该机制不仅提高了参数利用效率,还确保了各路由专家在专门化时聚焦于各自独特的知识。
【图1】展示了 DeepSeekMoE 16B 与其他开源模型在 Open LLM Leaderboard 上的对比。图中红色虚线为除 DeepSeekMoE 16B 外模型激活参数数量的数据线性拟合结果。从中可以看出,DeepSeekMoE 16B 在激活参数数量相近的模型中始终具有显著性能优势,并与拥有约2.5倍激活参数的 LLaMA2 7B 性能相当。
在 2B 参数规模上,我们通过12个零样本或少样本基准测试验证了 DeepSeekMoE 架构的优势。实验结果显示,DeepSeekMoE 2B 不仅大幅超越了 GShard 2B(Lepikhin 等,2021),甚至可与 GShard 2.9B(拥有1.5倍专家参数和计算量)相匹敌。更为关键的是,DeepSeekMoE 2B 的性能几乎达到了同总参数密集模型的上界。进一步的消融实验和专家专门化分析也验证了精细化专家划分和共享专家隔离策略的有效性。
基于此架构,我们将模型扩展至 16B 参数,并在2T token 大规模语料上训练。评测结果表明,DeepSeekMoE 16B 仅用约40%的计算量便达到了 DeepSeek 7B(DeepSeek-AI,2024)和 LLaMA2 7B 的性能水平。我们还将 DeepSeekMoE 16B 经过监督微调转化为聊天模型,其在聊天任务上同样展现了与 DeepSeek Chat 7B 和 LLaMA2 SFT 7B 相当的表现。初步尝试将 DeepSeekMoE 扩展至 145B 参数后,其在相同架构下持续展现出对 GShard 架构的优势,并在计算量仅为 DeepSeek 67B 的28.5%(甚至可能为18.2%)的情况下实现了性能匹敌。
本论文的主要贡献总结如下:
- 架构创新:我们提出了 DeepSeekMoE,一种旨在实现终极专家专门化的创新 MoE 架构,采用精细化专家划分和共享专家隔离两大策略。
- 实证验证:通过大量实验验证了 DeepSeekMoE 架构的有效性,实验结果表明 DeepSeekMoE 2B 已接近 MoE 模型性能上界。
- 可扩展性:我们将模型扩展至 16B,并证明其仅用约40%计算量便能达到与 DeepSeek 7B 和 LLaMA2 7B 相当的性能;同时,初步验证了 145B 规模下的持续优势。
- MoE 对齐:我们成功对 DeepSeekMoE 16B 进行监督微调,生成了对齐聊天模型,展现了其在不同任务下的适应性和通用性。
- 公开发布:本着开放研究精神,我们公开发布了 DeepSeekMoE 16B 模型检查点,值得注意的是,该模型可在40GB显存单卡上部署,无需量化处理。
2. MoE 在 Transformer 中的预备知识
我们首先介绍一种常用于 Transformer 语言模型中的通用 MoE 架构。标准 Transformer 语言模型由 L 层 Transformer 模块堆叠构成,每一模块可描述为:
其中 T 为序列长度,Self-Att(·) 表示自注意力模块,FFN(·) 表示前馈网络,uₗ^(1:T) ∈ ℝ^(T×d) 为第 l 层注意力模块后的所有 token 隐状态,hₗₜ ∈ ℝ^d 为第 l 层 Transformer 模块后第 t 个 token 的输出隐状态。为简洁起见,上式中省略了层归一化操作。
一种典型做法是用 MoE 层替换 Transformer 中的部分 FFN 层(Du 等,2022;Fedus 等,2021;Lepikhin 等,2021;Zoph,2022)。MoE 层由多个专家构成,每个专家均与标准 FFN 结构相同,随后每个 token 会被分配给一位(Fedus 等,2021)或两位专家(Lepikhin 等,2021)。设若将第 l 层的 FFN 替换为 MoE 层,则其输出 hₗₜ 的计算可表示为:
这里 N 表示专家总数,FFNᵢ(·) 为第 i 个专家,gᵢ,ₜ 为第 i 个专家对应的门控值,sᵢ,ₜ 为 token 与专家的亲和度,Topk(·, K) 表示从所有 N 个专家的亲和度中选出 K 个最大的值。注意,gᵢ,ₜ 是稀疏的,保证了每个 token 仅在 K 个专家上进行计算,从而实现了计算效率的提升。上式中同样为简洁起见省略了层归一化。
【图2】展示了三种架构:
(a) 传统 Top-2 路由策略的 MoE 层;
(b) 加入精细化专家划分后的结构;
(c) 集成共享专家隔离策略后的 DeepSeekMoE 架构。(一个 token 可以路由支更多专家)
3. DeepSeekMoE 架构
在第2节介绍的通用 MoE 架构基础上,我们提出了 DeepSeekMoE,以期充分发掘专家专门化的潜力。如【图2】所示,本架构采用了两大策略:精细化专家划分和共享专家隔离,二者共同提升了专家专门化程度。
3.1 精细化专家划分
在专家数量有限的情况下,分配到同一专家的 token 往往涉及多种知识,导致该专家参数难以同时有效学习所有知识。若能将每个 token 路由至更多专家,则不同知识可以分别被不同专家捕获,从而使每个专家保持较高的专门化。为实现此目标,我们在保持专家总参数和计算成本不变的前提下,将每个专家 FFN 划分为 m 个更小的专家,并相应增加激活专家的数量至 mK,从而获得更灵活的激活组合。这样,MoE 层的输出可表达为:
这种精细划分从组合角度极大提升了激活专家的灵活性。例如,当 N=16 时,传统 Top-2 路由可组合 16 选2=120 种方式;而若每个专家分成4个小专家,则总共 64 个专家激活 mK 个时,可能的组合数高达数亿数量级,从而有助于实现更精确、针对性的知识获取。
3.2 共享专家隔离
在传统路由策略下,不同专家间可能需要共性知识,导致多个专家同时学习相似知识而产生冗余。为缓解这一问题,我们引入共享专家隔离策略,即固定选定 Ks 个专家作为共享专家,每个 token 无论如何路由都将被分配到这些共享专家。为保证计算成本不变,其余路由专家激活数量相应减少 Ks 个。整合共享专家隔离后,MoE 层的表达式为:
共享专家隔离策略最早可追溯至 Rajbhandari 等(2022),不同之处在于我们从算法角度而非工程实现角度提出该策略。
3.3 负载均衡考虑
自动学习的路由策略可能面临负载不均问题,具体表现为:
- 路由崩溃,即模型总是选择少数专家,导致其他专家训练不足;
- 当专家分布在不同设备上时,负载不均可能加剧计算瓶颈。
为解决上述问题,我们引入了专家级和设备级的平衡损失。其中,专家级平衡损失定义为:
其中 fᵢ 为指示函数归一化的激活情况,Pᵢ 为该专家的平均门控值,α₁ 为超参数。设备级平衡损失则通过将所有路由专家分为 D 组,在设备上均衡计算量来设计,其表达式为:
其中 f′ᵢ 和 P′ᵢ 分别为每个设备上专家激活情况的均值和总和,α₂ 为设备级平衡因子。在实践中,我们通常设置较小的专家级平衡因子以防止路由崩溃,同时设定较大的设备级平衡因子以促进设备间均衡。
4. 验证实验
4.1 实验设置
4.1.1 训练数据与分词
我们的训练数据采自 DeepSeek-AI 构建的大规模多语言语料库,主要覆盖英文和中文,同时包括其他语言,数据来源涵盖网页文本、数学材料、代码、出版文献及其他各类文本。验证实验中,我们从语料库中采样 100B token,用 HuggingFace Tokenizer2 工具基于子集语料训练 BPE 分词器,构建词汇量为 8K 的分词器;在训练更大模型时词汇量将相应扩展。
4.1.2 计算平台
实验基于 HAI-LLM(High-Flyer,2023)训练框架,集成了张量并行、ZeRO 数据并行、PipeDream 流水线并行及专家并行等多种并行策略。为优化性能,我们开发了 CUDA 和 Triton 加速的 GPU 核,用于门控算法及跨线性层的计算融合。所有实验均在配备 NVIDIA A100 或 H800 GPU 的集群上进行。
4.1.3 超参数
【此处附超参数总览表,见附录 A】
验证实验中,Transformer 层数设为9,隐藏维度为1280,多头注意力共10个头(每头128维)。所有可学习参数初始化标准差为0.006。所有 FFN 均替换为 MoE 层,并确保专家总参数为标准 FFN 的16倍;同时激活专家参数(含共享专家和路由专家)约为标准 FFN 的2倍。在此配置下,每个 MoE 模型总参数约为2B,激活参数约为0.3B。
训练时采用 AdamW 优化器(β₁=0.9,β₂=0.95,weight_decay=0.1),学习率采用先热身后阶梯衰减策略,最大学习率为 1.08×10⁻³,梯度裁剪范数为1.0,批量大小为2000,最大序列长度为2000,每批约含4M token,总训练步数设为25,000。
4.1.4 评测基准
评测任务涵盖语言模型、语言理解推理、阅读理解、代码生成及闭卷问答等。具体基准包括:
- 语言模型:在 Pile 数据集(Gao 等,2020)上计算交叉熵损失。
- 语言理解与推理:HellaSwag、PIQA、ARC-easy、ARC-challenge,以准确率作为评测指标。
- 阅读理解:RACE-high 与 RACE-middle,评测指标为准确率。
- 代码生成:HumanEval 与 MBPP,指标为 Pass@1。
- 闭卷问答:TriviaQA 与 NaturalQuestions,指标为完全匹配率(EM)。
4.2 评测结果
我们将 DeepSeekMoE 与 Dense、Hash Layer、Switch Transformer、GShard 等模型进行对比。实验结果表明:
- 采用稀疏架构的 Hash Layer 和 Switch Transformer 均明显优于相同激活参数的密集基线;
- 与 Hash Layer 和 Switch Transformer 相比,GShard 由于激活参数更多表现略优;
- 在总参数及激活参数相同的情况下,DeepSeekMoE 明显优于 GShard,各项基准测试均有较大提升。
【表1】列出了验证实验的详细评测结果,具体指标及对比数据详见原文表格。
4.3 DeepSeekMoE 接近 MoE 模型上界
为更准确评估 DeepSeekMoE 性能,我们将其与更大规模的 GShard 和密集模型进行对比。结果显示,DeepSeekMoE 2B 的性能可与专家参数和计算量为其1.5倍的 GShard×1.5 模型相媲美,同时几乎达到密集模型 Dense×16 的上界。
【表2】展示了 DeepSeekMoE 与更大 GShard 模型及密集模型的详细对比。
4.4 消融实验
为了验证精细化专家划分和共享专家隔离策略的有效性,我们分别对两者进行了消融实验。结果表明:
- 添加共享专家隔离后,多数基准测试性能提升,验证了其在减少专家冗余方面的作用;
- 进一步细化专家划分(将专家划分为2或4个子专家)可持续提升模型整体性能;
- 关于共享专家与路由专家的比例,不同配置对性能影响不大,最终采用共享专家与激活路由专家比例为1:3 的配置获得了略优结果。
【图3】展示了消融实验结果,性能均归一化后进行比较。
4.5 专家专门化分析
我们通过禁用部分路由专家来检验专家冗余性。实验显示,DeepSeekMoE 对禁用部分激活专家更为敏感,表明其每个路由专家更为不可替代。进一步验证了共享专家在捕捉基础知识上的关键作用,一旦禁用共享专家,模型损失显著上升。
【图4、图5、图6】分别展示了不同禁用比例下的 Pile 损失变化和不同激活专家数量下的性能比较,进一步印证了 DeepSeekMoE 高效利用专家参数的优势。
5. 向 DeepSeekMoE 16B 扩展
利用 DeepSeekMoE 架构,我们将模型扩展至 16B 参数,并在 2T token 的大规模语料上训练。结果显示,DeepSeekMoE 16B 在仅用约40%计算量的前提下,可达到与 LLaMA2 7B 同等性能。
5.1 实验设置
5.1.1 训练数据与分词
训练数据同上,但采样量增至 2T token;分词器使用 HuggingFace Tokenizer,词汇量设置为 100K。
5.1.2 超参数
DeepSeekMoE 16B 模型设置为28层 Transformer,隐藏维度2048,多头注意力共16个头(每头128维)。除第一层外,所有 FFN 均替换为 MoE 层,每层包含2个共享专家和64个路由专家,每个专家大小为标准 FFN 的0.25倍。每个 token 路由至2个共享专家及6个路由专家。整体模型参数约为16.4B,激活参数约为2.8B。
训练采用 AdamW 优化器,学习率策略与验证实验类似,总训练步数约为106,449步,确保2T token 训练量。为避免设备间负载不均,各层专家部署在同一设备上,故不采用设备级平衡损失。
5.1.3 评测基准
在前述基准外,语言建模评测采用 Pile 数据集,并以每 byte 的比特数(BPB)作为指标;阅读理解新增 DROP 基准;数学推理新增 GSM8K 和 MATH 基准;多学科选择题采用 MMLU;歧义消解采用 WinoGrande;另外,由于预训练语料包含中文,额外评测了 CLUEWSC、CEval、CMMLU 与 CHID 等中文基准。
同时,基于 HuggingFace 支持的 Open LLM Leaderboard 对所有基准进行综合评测。
5.2 评测结果
5.2.1 与 DeepSeek 7B 内部对比
【表3】显示,DeepSeekMoE 16B 在仅用约40%计算量的前提下,与 DeepSeek 7B(密集模型,6.9B参数)在 Pile、HellaSwag、TriviaQA 等知识密集任务上均表现相当甚至略优,同时在数学推理、代码生成任务上更具优势;但在多选题任务上,由于注意力参数较少,略逊于 DeepSeek 7B。
5.2.2 与开源模型对比
【表4】展示了 DeepSeekMoE 16B 与 LLaMA2 7B 的对比。结果表明,DeepSeekMoE 16B 在计算量仅约40%的情况下,多数任务上均超过 LLaMA2 7B,尤其在数学推理、代码生成及中文任务上优势明显。
6. DeepSeekMoE 16B 的对齐(Alignment)
之前研究认为 MoE 模型通过微调难以获得明显提升(Artetxe 等,2022;Fedus 等,2021),但 Shen 等(2023)的研究表明,MoE 模型也能从指令调优中获益。为验证这一点,我们对 DeepSeekMoE 16B 进行了监督微调(SFT),生成聊天模型 DeepSeekMoE Chat 16B,其在语言理解、推理、阅读理解、数学和代码生成等任务上均表现出与 LLaMA2 SFT 7B 及 DeepSeek Chat 7B 相当甚至更优的性能。
6.1 实验设置
SFT 数据采自内部精心构建的 140 万样本数据集,涵盖数学、代码、写作、问答、推理、摘要等多领域,数据主要为中英文。微调时批量大小设为1024,每轮训练8个 epoch,最大序列长度为4000,学习率固定为 10⁻⁵,无额外调度。
评测基准与第5节类似,除去 Pile 与 CHID,并新增 BBH 以评估模型推理能力。
6.2 评测结果
【表5】展示了 LLaMA2 SFT 7B、DeepSeek Chat 7B 与 DeepSeekMoE Chat 16B 的详细对比。结果表明:
- DeepSeekMoE Chat 16B 在语言理解、阅读理解、数学及知识密集任务上均与两款7B模型表现相当或略优;
- 在代码生成任务上表现尤为突出;
- 尽管在多选题任务上仍略逊于 DeepSeek Chat 7B,但整体性能差距已被缩小;
- 预训练时融合中英文语料,使其在中文基准上大幅领先于纯英文 LLaMA2 SFT 7B。
7. DeepSeekMoE 145B 初步工作
基于 DeepSeekMoE 16B 的优秀表现,我们进一步尝试将模型扩展至 145B 参数。初步实验中,DeepSeekMoE 145B 在 245B token 上训练,已显示出对 GShard 架构的一致优势,并初步达到或超过 DeepSeek 67B(密集模型)的性能。最终完整版及全量训练完成后,我们计划公开发布该模型。
7.1 实验设置
训练数据与分词器同 DeepSeekMoE 16B,唯一区别在于训练 token 数扩展至 245B。模型设置为62层 Transformer,隐藏维度4096,多头注意力共32个头(每头128维)。除第一层外,其余 FFN 均替换为 MoE 层,每层包含4个共享专家和128个路由专家,每个专家大小为标准 FFN 的0.125倍;每个 token 路由至4个共享专家和12个路由专家。整体模型参数约为144.6B,激活参数约为22.2B。训练设置与前述类似。
同系列