群魔乱舞:MoE大模型详解
导读
大模型领域,近两年最火的概念之一,非MoE莫属。从GPT-4到Google的Switch Transformer,再到国内势头正猛的DeepSeek,似乎一夜之间,所有顶尖模型都跟“混合专家”这个词绑定在了一起。那么,这种架构到底有什么魔力?为什么它能用更低的成本训练出更大的模型?今天,我们就来系统性地梳理一下MoE大模型的来龙去脉、核心原理,以及那些让人拍案叫绝的优化技巧。
GPT-4是8个2200亿MoE模型
GPT-4的参数量远不止1万亿——这个说法在2023年6月被美国知名黑客George Hotz在一次采访中直接捅破:它由8个220B参数模型组成,简单一算,8 x 220B = 1.76万亿参数。这个信息甚至得到了PyTorch创始人Soumith Chintala的认同。下面这张“八头怪”示意图,也许就是GPT-4最形象的写照。
当然,MoE应用于大模型,GPT-4并不是第一个吃螃蟹的。早在2022年,Google就推出了MoE大模型
Switch Transformer
国内团队DeepSeek也紧随其后,
开源了国内首个MoE大模型
。
2B模型性能可接近2B Dense模型,但仅用了17.5%的计算量
这种压倒性的效率优势,让国内大模型开始集体向MoE方向冲刺。可以预见,2024年会有越来越多的模型选择MoE架构。那么,本质上,MoE究竟是什么?它凭什么能“花小钱办大事”?
什么是MoE大模型?
MoE,全称Mixed Expert Models,翻译过来就是混合专家模型。听起来很高端,但这个概念并不新——早在1991年的论文《Adaptive Mixture of Local Experts》中就已经被提出。
我们知道,模型规模是提升性能的关键因素之一。在有限的计算资源预算下,用更少的训练步数训练一个更大的模型,往往比用更多步数训练一个较小的模型效果更好。而MoE的核心优势就在于,它能在远少于Dense模型所需的计算资源下进行有效预训练。这意味着同样的计算预算,你可以极大地扩大模型或数据集的规模。
MoE基于Transformer架构,主要由两部分组成:
- 替代了传统Transformer模型中的前馈网络(FFN)层。每个MoE层包含若干“专家”(例如8个),每个专家本身是一个独立的神经网络(通常是FFN,但也可以是更复杂的结构)。
稀疏MoE层:
- 它的作用是决定每个token应该被发送到哪个专家。比如在图中,“More”这个token可能被分配给第二个专家,而“Parameters”则被分配给第一个专家。有些情况下,一个token甚至可以同时发送给多个专家。
门控网络或路由:
总结一下:在MoE中,我们将传统Transformer的每个FFN层替换为MoE层,而每个MoE层由两个核心部分组成——
一个路由器(或叫门控网络)和若干数量的专家
MoE大模型具备哪些优势?
最核心的优势还是效率。与Dense模型相比,在相同计算资源下,MoE的训练速度更快,且可以训练更大的模型。比如Google的Switch Transformer,模型大小是T5-XXL的15倍,但在固定计算预算下,达到固定困惑度PPL的速度比T5-XXL
快4倍
40%
总结下来,MoE的优点主要体现在以下四点:
训练速度更快,效果更好。
相同参数下,推理成本更低。
- ,可以在保持计算成本不变的情况下增加参数数量,轻松扩展到万亿参数规模。
扩展性好
- ,例如Switch Transformer在所有101种语言测试中都显示出性能提升。
多任务学习能力强
当然,硬币的另一面也存在。MoE的缺点同样不容忽视:
- MoE在训练过程中可能会遇到稳定性问题,需要特别注意。
训练稳定性:
- 在分布式训练环境中,专家路由机制会增加通信开销,尤其是在模型规模较大时。
通信成本:
- 设计相对复杂,可能需要更多的工程努力来实现和优化。
模型复杂性:
- 稀疏性导致Fine-tuning时容易出现过拟合。
下游任务性能:
接下来,我们深入技术细节,重点回答三个核心问题:
MoE为什么能实现更大的模型参数和更低的训练成本?如何解决训练稳定性问题?又如何应对Fine-tuning中的过拟合?
一、Adaptive mixtures of local experts
这篇1991年的论文是大多数MoE工作的开山鼻祖,作者中有两位大家熟知的大佬:Michael Jordan和Geoffrey Hinton。
论文的核心思想是:当用一个多层网络训练不同的子任务时,不同任务的学习过程会相互干扰,导致学习缓慢和泛化能力差。为了解决这个问题,论文提出使用多个模型(即专家,expert)来学习,并引入一个门控网络(gating network)来决定每个数据样本应由哪个模型负责训练,从而减轻不同类型样本之间的干扰。
对于一个样本,第i个expert的输出为o_i,期望输出为d,那么损失函数L的计算方式为:
L = -log(∑ g_i * exp(-1/2|d - o_i|^2))
其中g_i是门控网络分配给第i个expert的权重。
论文指出,这个损失函数可能导致expert网络之间
强烈耦合
鼓励专家网络之间的相互竞争而非合作
论文还提到,将损失函数进一步修改为以下形式效果更好:
L = -log(∑ g_i * exp(-1/2|d - o_i|^2))
这里先将每个expert的损失指数化后再加权求和,最后取log。这种技巧的好处在于:导数只与当前expert有关,不受其他expert与当前样本匹配程度的影响,从而进一步提升了训练的效率和稳定性。
二、Sparsely-Gated MoE
在2010到2015年间,两个独立的研究领域为MoE后续发展奠定了重要基础:
- 将MoE作为更深层网络的一个组件,嵌入到多层网络中的某一层,使模型既大又高效。
组件专家:
- 基于输入token动态激活或停用网络组件的方法。
条件计算:
2017年,Shazeer等人(团队包括Geoffrey Hinton和Jeff Dean)将这一概念应用于137B的LSTM。通过引入稀疏性,这项工作将模型规模提升了
1000多倍
与1991年的工作相比,这里的Sparsely-Gated MoE有两个关键区别:
- 只有极少数expert会被激活用于推理。这种稀疏性使得我们可以使用海量的experts来把模型容量做到极大。
Sparsely-Gated:
- 前面那篇论文是sample-level的(不同样本使用不同专家),而这里是token-level的,一个句子中不同的token可以分配给不同的专家。
Token-level:
在模型结构中,每个token都会经过一个MoE Layer。每个MoE layer包含一堆experts(每个expert都是一个小型FFN),还有一个Gating Network根据当前token选择少数几个expert进行计算。
2.1 门控网络(Gating Network)
门控网络是Sparsely-Gated MoE层的核心,负责为每个输入token选择一个稀疏的专家组合。其输出是所有experts的加权和:
y = ∑ G(x)_i * E_i(x)
如果门控网络输出G(x)为0,则对应的E(x)就无需计算,从而节省计算资源。典型的门控网络是一个带softmax的简单网络,学习将输入发送给哪个expert。
Shazeer等人还探索了“Noisy Top-K Gating”机制:先添加一些可调整的噪声,然后保留前K个值(K通常为2-4),再应用softmax。对于非TopK的部分,由于值是负无穷,经过softmax后变成0,因此不会被选中。噪声项则有助于不同expert的负载更加均衡。
2.2 平衡专家利用率(Balancing Expert Utilization)
论文发现,门控网络容易收敛到“总是为相同的几个专家产生大权重”的状态。这种不平衡是自我强化的——受到青睐的专家训练得更快,因此被更多地选择。这会导致某些专家可能从未被使用过。
为解决这个问题,论文提出了一种软约束方法:定义每个专家的重要性Importance(即该专家在一批样本中门控值的总和),然后添加一个额外的损失函数L_importance,其值为重要性值集合的CV(coefficient of variation)平方乘以一个缩放因子w_importance。这个额外的损失会鼓励所有专家的重要性相等,从而保证负载均衡。
三、GShard
前面两篇是MoE的基础工作,但都没有在大模型上得到广泛应用。GShard是Google在2021年提出的第一个将MoE思想拓展到Transformer上的工作。
具体的做法是:把Transformer的encoder和decoder中,每隔一个的FFN层替换成MoE层,使用Top-2门控网络。这种架构对大规模计算非常有效——当扩展到多个设备时,MoE层在不同设备间共享,而其他所有层则在每个设备上复制。
实现MoE跨设备分片的关键技术是模型并行化和数据并行化的结合。在模型并行化中,MoE层的专家网络被分配到不同设备上;在数据并行化中,输入数据被分割成多个部分,分配给不同设备。GShard模块提供了一套API和编译器扩展,允许用户简单注释关键张量,自动完成分片和通信。
由于专家被分配到不同设备并行计算,模型的计算效率大幅提升——这正是MoE能实现更大模型参数、更低训练成本的直接原因。
为了进一步保持负载平衡和训练效率,GShard还引入了一些关键变化:
- 在Top-2设置中,始终选择排名最高的专家,但第二个专家是根据其权重比例随机选择的。
随机路由:
- 设定一个阈值,定义每个专家能处理多少token。如果两个专家的容量都达到上限,token会溢出,通过残差连接传递到下一层,或在某些情况下被完全丢弃。
专家容量:
注意:
四、Switch Transformers
尽管MoE潜力巨大,但复杂性、通信成本以及训练和微调过程的不稳定性,仍制约着其广泛采用。2022年,Google提出的Switch Transformers一定程度缓解了这些问题。作者在Hugging Face上发布了一个1.6万亿参数的MoE模型(2048个专家),实现了与T5-XXL相比4倍的预训练速度提升。
Switch Transformers简化了MoE路由算法,降低了通信和计算成本,并首次证明了用较低精度(bfloat16)格式训练大型稀疏模型的可能性。
在和T5 Base、T5 Large的对比中,Switch Transformers在相同计算资源下获得了高达7倍的预训练速度提升。而且在多语言实验中,它在所有101种语言测试中都取得了提升。
实验数据显示:模型参数随着专家数量增加而增加,但FLOPs per token保持不变——这表明模型在保持计算效率的同时,能利用更多参数来提高性能。
4.1 Switch Transformer 主要优化
Swith Transformer的设计指导原则非常直观:
尽可能地把Transformer模型的参数量做大!
和之前所有MoE模型不同,
Switch Transformer的门控网络每次只路由到1个expert
- 减少路由计算,一个token每次只路由到一个专家
- 每个专家的batch size(专家容量)至少可以减半
- 简化路由实现,降低MoE中的通信成本
4.2 Switch Routing
这里要重点讲一下“专家容量”(Expert Capacity)的概念——它是指每个专家在模型中处理的token数量。为什么需要这个容量?因为在编译时,所有tensor的形状是静态确定的,但我们无法提前知道多少token会分配给每个专家。所以,需要将batch中的总token数平均分配给所有专家,并通过一个容量因子(capacity factor)来扩展每个专家的容量,以应对token分布不均的情况。
容量因子是一个大于1.0的数,为每个专家提供额外缓冲空间。当容量因子=1.0时,输入6个token,每个专家的容量等于2,如果有专家被分配了3个token,超出的token会“溢出”——模型会跳过计算,直接将token表示通过残差连接传递到下一层。当容量因子=1.5时,每个专家就能处理3个token,溢出问题就解决了。
但容量因子不能设得过高,否则会导致计算资源和内存的浪费。在论文中,Switch Transformers在低容量因子(例如1至1.25)下表现出色。下表展示了不同容量因子下的效果对比。
表4-1:Switch Transformer 和 MoE 的效果对比
| 模型 | 容量因子 | 训练100k steps后的负对数困惑度 | 到达-1.5所需时间(小时) | 训练速度(样本/秒) |
|---|---|---|---|---|
| T5-Base | - | -1.731 | 未达到 | 1600 |
| T5-Large | - | -1.550 | 131.1 | 470 |
| MoE-Base | 2.0 | -1.547 | 68.7 | 840 |
| Switch-Base | 2.0 | -1.554 | 72.8 | 860 |
| MoE-Base | 1.25 | -1.559 | 72.8 | 790 |
| Switch-Base | 1.25 | -1.553 | 65.0 | 910 |
| MoE-Base | 1.0 | -1.572 | 80.1 | 860 |
| Switch-Base | 1.0 | -1.561 | 62.8 | 1000 |
| Switch-Base+ | 1.0 | -1.534 | 67.6 | 780 |
Switch Transformer的作者还重新审视并简化了负载均衡损失,通过合理设置其系数,实现了良好的负载分布。
4.3 不同的负载均衡损失
在稀疏模型中,专家分布在多个设备上。理想情况下,每个专家应处理相同数量的数据。为此,论文引入了一种辅助损失函数:
L_aux = w_aux * N * ∑ (f_i * P_i)
其中f_i是batch中分配给专家i的token占比,P_i是所有输入token被路由到专家i的平均概率。通过最小化这个损失,可以鼓励均匀路由。最终loss乘以专家数量N,使得即使专家数量变化,loss也能保持恒定。
4.4 稀疏路由和负载均衡loss的合并效果
实验设置:在C4数据集上进行预训练,使用MLM作为预训练目标,对比Switch Transformer与MoE Transformer以及T5。所有模型都在相同硬件(TPUv3)上进行相同步数训练。
核心结论(参见表4-1):
- Switch Transformer在速度和效果上都优于MoE Transformer。
- 若将Switch Transformer的规模增加到匹配MoE Transformer的训练速度,它在每一步上都优于所有MoE模型。
- Switch Transformer在较低的容量因子(1.0, 1.25)下表现更好——较低的专家容量表明,大模型中内存稀缺,容量因子应尽可能小。
4.5 改进训练和Fine-Tuning技术
1. 精度选择
作者尝试了混合精度方法。最初,当专家和门控网络都使用bfloat16训练时,出现了不稳定现象。原因是路由计算涉及指数函数等对精度要求较高的操作。最终方案是:将路由器输入转换为float32,其他部分保持bfloat16。这种混合精度策略几乎达到了float32的稳定性,同时保持了bfloat16的训练速度。
表4-2:不同精度效果对比
| 模型精度选择 | 负对数困惑度 | 训练速度(样本/秒) |
|---|---|---|
| Switch-Base (float32) | -1.718 | 1160 |
| Switch-Base (bfloat16) | -3.780 | 1390 |
| Switch-Base (混合精度) | -1.716 | 1390 |
2. 更小的参数初始化
作者观察到,在Switch Transformer中,权重初始化特别重要。建议将默认的初始化超参数s从1.0减少10倍,即s=0.1。这种较小的初始化规模有助于提高模型效果和减少训练不稳定性。
表4-3:减小参数初始化规模可以提升训练稳定性
| 权重初始化规模 | 负对数困惑度 | 标准差 |
|---|---|---|
| 0.1倍初始化 | -2.72 | 0.01 |
| 1.0倍初始化 | -3.60 | 0.68 |
3. Fine-Tuning 过程正则化
针对Fine-tuning过程中的过拟合问题,作者提出了“expert dropout”策略——只在专家层增加dropout率。这种策略有效减少了过拟合风险,同时保持了模型在下游任务上的性能。
表4-4:Fine-Tuning过程中正则化效果
| 模型(dropout) | GLUE | CNNDM | SQuAD | SuperGLUE |
|---|---|---|---|---|
| T5-Base (d=0.1) | 82.9 | 19.6 | 83.5 | 72.4 |
| Switch-Base (d=0.1) | 84.7 | 19.1 | 83.7 | 73.0 |
| Switch-Base (d=0.2) | 84.4 | 19.2 | 83.9 | 73.2 |
| Switch-Base (d=0.3) | 83.9 | 19.6 | 83.4 | 70.7 |
| Switch-Base (d=0.1, expert d=0.4) | 85.2 | 19.6 | 83.7 | 73.0 |
4.6 高效训练:数据、模型、专家并行
任意增加专家数量会导致收益递减,但可以通过增加模型维度(如隐藏层大小d_model或前馈网络维度d_ff)来继续提升效果,这又需要并行技术支持。
这里补充一下各种并行方法的概念:
- 模型复制,数据并行处理。每个device上保存完整模型拷贝,前向计算后汇总梯度更新。
数据并行:
- 模型不同参数分配到不同device上,处理一个batch的数据。
模型并行:
- 将模型并行与数据并行结合,每个core处理一部分token和一部分权重。
模型和数据并行:
- 每个专家分配到一个core上,数据也切分成多份。
专家和数据并行:
- 三种并行策略合并使用,实现极致效率。
专家、模型和数据并行:
五、GLaM
除了Switch Transformer,Google还推出了另一个MoE模型:GLaM(Generalist Language Model)。它比GPT-3大三倍,但由于使用了Sparse MoE设计,训练成本仅为GPT-3的1/3,而且在29个NLP任务上全面超越了GPT-3。
表5-1:GLaM实验对比
| 模型 | 模型类型 | 参数量 | 激活的参数量 |
|---|---|---|---|
| BERT | Dense Encoder-only | 340M | 340M |
| T5 | Dense Encoder-decoder | 13B | 13B |
| GPT-3 | Dense Decoder-only | 175B | 175B |
| Jurassic-1 | Dense Decoder-only | 178B | 178B |
| Gopher | Dense Decoder-only | 280B | 280B |
| Megatron-530B | Dense Decoder-only | 530B | 530B |
| GShard-M4 | MoE Encoder-decoder | 600B | 1.5B |
| Switch-C | MoE Encoder-decoder | 1.5T | 1.5B |
| GLaM (64B/64E) | MoE Decoder-only | 1.2T | 96.6B |
GLaM虽然总参数量有1.2T,但实际推理中激活的参数量只有96B——这意味着在推断时,它比GPT-3这类Dense模型快得多。
六、ST-MoE
之前的负载均衡损失可能导致稳定性问题。虽然可以通过引入dropout等方法来稳定训练,但可能牺牲模型质量。ST-MoE论文提出了一种新的辅助损失——Router z-loss,专门针对稀疏专家模型中的路由器部分设计。
6.1 用Router z-loss稳定模型训练
Router z-loss的核心思想是鼓励路由器产生较小的logits值。因为较大的logits在softmax中会导致较大梯度,可能引起训练不稳定。其定义如下:
L_z = w_z * (1/N) * ∑ (log(∑ exp(x_i)))^2
其中N是batch中的token数量,x_i是路由器的logits。通过惩罚较大的logits值,这个损失有助于减少训练不稳定性,并可能提高模型泛化能力。
6.2 专家如何学习?
ST-MoE的研究者们发现:encoder中的不同专家倾向于专注于特定类型的token或浅层概念(如标点符号、专有名词等),而decoder中的专家通常专业化程度较低。此外,在多语言训练中,专家并不会按语言分派——尽管人们可能会预期每个专家处理一种特定语言,但由于token路由和负载均衡机制,没有任何专家被特定配置以专门处理某一语言。
6.3 专家的数量对预训练有何影响?
增加更多专家可以提升处理样本的效率和加速运算,但这些优势随专家数量增加而递减(尤其是专家数达到256或512之后)。同时,模型推理时需要更多显存来加载整个模型。但好消息是,Switch Transformers的研究表明,大规模模型中的特性在小规模模型下也同样适用——即便每层仅包含2、4或8个专家。
6.4 Fine-Tuning MoE 模型
稠密模型和稀疏模型在过拟合动态上存在显著差异。稀疏模型更易于过拟合,因此需要更强的内部正则化(如更高比例的dropout)。具体来说,可以为稠密层设定较低dropout率,而为稀疏层设置更高dropout率。
在Fine-tuning过程中是否使用辅助损失也是一个需要决策的问题。ST-MoE的作者尝试关闭辅助损失后发现,即使高达11%的token被丢弃,模型质量也未显著受影响——token丢弃可能本身就是一种正则化形式。
实验还观察到:在相同预训练PPL下,稀疏模型在下游任务中的表现不如对应稠密模型,尤其是在理解任务(如SuperGLUE)上;但在知识密集型任务(如TriviaQA)上,稀疏模型表现异常出色。另外,较少的专家数量有助于改善Fine-tuning性能。
一个可行的Fine-tuning策略是:尝试冻结所有非专家层的权重。实验显示,仅冻结MoE层的参数几乎与更新所有参数效果相当,这可以加速Fine-tuning并降低显存需求。Fine-tuning MoE时,还需要注意超参数的特殊设置——稀疏模型往往更适合使用较小的batch size和较高的学习率。
七、开源MoE模型
目前已经有一些开源的MoE大模型。国内的代表是DeepSeek团队开源的DeepSeekMoE,模型、代码、论文均已同步发布。
- 模型下载:https://huggingface.co/deepseek-ai
- 微调代码:https://github.com/deepseek-ai/DeepSeek-MoE
- 技术报告:https://github.com/deepseek-ai/DeepSeek-MoE/blob/main/DeepSeekMoE.pdf
国外也有一些开源的MoE模型和训练代码:
- Megablocks:https://github.com/stanford-futuredata/megablocks
- Fairseq:https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
- OpenMoE:https://github.com/XueFuzhao/OpenMoE
此外,还有开源了模型但未开源代码的项目:
- Switch Transformers (Google):基于T5,专家数从8到2048,最大模型1.6万亿参数。
- NLLB MoE (Meta):NLLB翻译模型的MoE变体。
- OpenMoE:社区基于Llama模型的MoE尝试。
- Mixtral 8x7B (Mistral):性能超越Llama 2 70B的高质量MoE,推理速度更快。
总结
本文系统性地介绍了混合专家模型(MoE)的核心原理、高效训练方法以及Fine-tuning技巧。现在,我们来回答开篇提出的三个问题。
第一个问题:MoE为什么能够实现在低成本下训练更大的模型?
主要原因在于稀疏路由——每个token只会选择top-k个专家进行计算,大幅降低了实际参与计算的参数量。同时,结合模型并行、专家并行和数据并行策略,可以进一步优化MoE的训练效率。而负载均衡损失则保证了每个设备的利用率。
第二个问题:MoE如何解决训练稳定性问题?
主要途径包括:采用混合精度训练(将路由器保持float32)、使用更小的参数初始化(将初始化规模降低10倍),以及引入Router z-loss来抑制异常的logits值。
第三个问题:MoE如何解决Fine-tuning过程中的过拟合问题?
主要策略是:对专家层使用更大的dropout率、采用更大的学习率以及更小的batch size。目前针对Fine-tuning的优化手段主要还是围绕这些常规思路进行。
参考
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
社区发布 | 深度求索开源国内首个 MoE 大模型,技术报告、模型权重同时发布
Adaptive mixtures of local experts
混合专家模型(MoE)详解
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
ST-MOE: Designing Stable and Transferable Sparse Expert Models
Mixtral of experts
-
- 关于宇宙的好的网名有哪些
- 角色扮演 | 1
- 网名