deepseek技术解读(3)-MoE的演进之路

0. 引言

本篇讲讲deepseek在MoE（Mixture-of-Experts）上的演进过程。DeepSeek是MoE稀疏模型的忠实玩家。主版本模型从DeepSeekMoE（V1）到 DeepSeek V3，一直坚持走MoE的技术路线，并且持续做出一些创新。本文参考paper并结合源码阅读，理解MoE的演进过程和具体实现。

1.简述MoE的发展历程

首先我们简单回顾下MoE的发展历史，早在1991年一篇名为《Adaptive Mixtures of Local Experts 》的工作，最早提出了Mixture of Experts的原型框架，如图1，直至今日，MoE的框架依然保持这种形式。

图1、Adaptive Mixtures of Local Experts 框图

MoE(Mixture of Experts)是一种网络层结构，网络层主要包括三部分：

专家网络(Expert Network)：是一个前馈网络，逻辑上一个专家网络擅长处理一类专项的子任务，所有专家都接受相同的输入，来做特定计算处理，产出不同的输出
门控网络(Gating Network)：跟专家网络接收一样的输入，负责产出专家偏好的权重。来指示对于一个输入，不同专家的重要程度。
选择器(selector)：是一种根据专家权重来做专家选择的策略。可以选择权重最高的Top1专家或选择TopK专家来融合得到最终的结果。

随后一段时间，主要是Google在主导着MoE的发展。进入Transformer时代后，2020年Google发表了的《GShard》，也把模型训到了600B的规模。GShard刻画了在Transformer上做MoE的经典设计。主要包括Transformer MoE层设计和辅助负载均衡损失。

Transformer MoE层：MoE层替换Transformer的FFN层，计算逻辑：对于一个token 分别通过门控网络和专家网络计算门控值和专家输出，然后用门控值加权多个专家输出来产出最终结果。具体如下：

注：这里的专家是token级专家，而不是样本粒度，每个token都会做专家路由。此外专家是稀疏激活的，是根据门控值取topK个专家来融合计算最终的结果。GShard最多激活权重最高的2个专家。

负载均衡-辅助损失：引入负载均衡损失，目的是解决多专家token分布不均的问题。因为如果完全按门控权重选取topk专家，容易导致训练过程出现负载不均衡的问题。比如：大多数token被分配到少数几个专家，导致只有少数专家数据通信繁忙造成拥堵，从而减缓训练速度；也会导致其他专家得不到充分训练。为了解决这个问题，定义了一个辅助损失（aux_loss）来降低负载不均衡问题。

那么怎么定义负载均衡的辅助损失？

2.DeepSeek的工作

2.1. DeepSeek-moe(V1)

24年1月DeepSeek发布V1版MoE模型，作者指出当前方法存在两方面问题：

知识混合性：现有的MoE模型通常使用数量有限的专家（如8个或16个），由于token的知识是丰富多样的，将多样的知识分配给有限的专家，会导致特定专家的token很可能会涵盖多样化的知识，而使得专家变成一个杂糅多知识的专家，这样不能充分发挥专家的专业效果。
知识冗余性：分配给不同专家的token可能存在共同知识。因此，多个专家可能会在其各自的参数中学习到共享知识，从而导致专家参数存在冗余。这种问题也阻碍了现有MoE实践中专家的专业化，限制了MoE模型的理论上限性能。

针对上述问题，DeepSeek引入一种实现了专家专业化而设计的创新MoE架构。架构主要包含两方面优化：

细粒度专家分割（Fine-Grained Expert Segmentation）：在保持参数数量不变的情况下，作者通过分割FFN中间隐藏维度来将专家分割成更细的粒度。相应地，在保持计算成本不变的情况下，可激活更多细粒度的专家，以实现激活专家组合的更高灵活性。细粒度专家分割使得多样化的知识能够被更细致地分解，并更精确地学习到不同的专家中，每个专家将保持更高的专业化水平。
共享专家隔离（Shared Expert Isolation）：将某些专家隔离出来，作为始终激活的共享专家，旨在捕获不同上下文中的共同知识。通过将共同知识压缩到这些共享专家中，可以减轻其他路由专家之间的冗余，这可以提高参数效率，确保每个路由专家专注于不同方面而保持专业化。

如下图2所示。(b)是在(a)基础上，通过将隐层切分更细粒度，而形成细粒度专家，(c)又在(b)基础上隔离出来共享专家。DeepSeekMoE模型的演进过程，一直延续这两个创新的设置。

图2、DeepSeekMoE架构

DeepSeekMoE架构的公式形式：

图3、DeepSeekMoE计算公式

图4、Expert-Level Balance Loss 计算

				
					class MoEGate(nn.Module):
    def forward(self, hidden_states):
        bsz, seq_len, h = hidden_states.shape        
        ############################
        # 这里的hidden_states就是公式里的T，是一个Batch数据的全部token做计算，每个Batch会重新计算
        ############################
        hidden_states = hidden_states.view(-1, h)
        
        logits = F.linear(hidden_states, self.weight, None)
        scores_for_aux = logits.softmax(dim=-1)
        
        topk_weight, topk_idx = torch.topk(scores_for_aux, k=self.top_k, dim=-1, sorted=False)
        topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
        mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
        ce = mask_ce.float().mean(0)
        ############################
        # 计算Pi，fi 和 aux_loss。这里的计算并没有跨Batch累积，每个Batch单独计算
        ############################      
        Pi = scores_for_aux.mean(0)
        fi = ce * self.n_routed_experts
        aux_loss = (Pi * fi).sum() * self.alpha

2.2. DeepSeek V2 MoE升级

DeepSeek V2 相对于V1版，对MoE模块主要在负载均衡上做了三方面升级

1.设备受限的专家路由机制（Device-Limited Routing）

随着LLM的size越来越大，对MoE模型的训练，一般要采用专家并行(expert parallelism)来分布式加载模型，也就是对于网络的一个MoE层的多个专家，分配到多个设备上，来并行训练。由于DeepSeek的MoE是做了细粒度专家的设计，通常专家会很多（V2模型的路由专家数有160个，激活专家6个）。我们知道在MoE层多专家的输入是一样的，由当前层的Self-Attention输出的隐层激活值作为MoE层的输入。如果被激活的专家分布在多个机器上，那么要把输入传输到多机器，势必会带来成倍的通讯成本。

图6 、通信复杂均衡loss 公式

图7. DeepSeek V3 MoE网络层计算框架

从实现门控的效果上看，Softmax和Sigmoid都能做实现筛选TopK的功能，也能做概率分布的归一化处理。

但V3版的MoE为什么要做从Softmax -> Sigmoid的升级？

要解释这个问题，我们看看V3版相对于V2版的专家设置发生了哪些变化。

V2版：路由专家数： 160，激活专家数： 6个，模型总参数67B，激活参数21B
V3版：路由专家数： 256，激活专家数： 8个，模型总参数671B，激活参数37B

图8、Sequence-Wise Auxiliary Loss计算

相对于V1版的专家级辅助损失(Expert-Level Balance Loss)其实就是作用粒度不一样，Sequence-Wise的粒度是单条样本粒度的token做计算。Expert-Level Balance是一个Batch的多Sequence的token做计算。公式的计算形式并没有什么差异。

最后DeepSeekV3也强调通过上面的负载均衡的策略，能达到一个非常好的平衡效果，所以在V3版并没有Token被Drop掉。

3. 总结

我们再来回顾下DeepSeek在MoE方面创新和演进过程

V1版为了兼顾对通用知识和细粒度领域知识的建模，引入了共享专家（Shared Expert）和细粒度专家（Fine-Grained Expert）。同时为了平衡各个专家的计算负载，引入了专家级负载loss （Expert-Level Balance Loss）和设备级负载loss（Device-Level Balance Loss）。
V2版主要在通信负载上做了些优化，通过引入设备受限的专家路由机制和通信负载均衡loss确保设备输入、输出的通信负载均衡。
V3版考虑负载loss对主模型的优化会有影响，将辅助负载loss做了精简，通过在门控权重增加一个可调的bias来解决通信和计算的负载。也引入了一个更细粒度的sequence负载均衡loss。同时考虑随着路由专家增到256个，在门控权重计算上选择了值域更宽、打分差异更显著的sigmoid函数替换了原来的softmax函数。

整体演进过程，如下图所示：