混合专家模型（MoE）入门

模型规模是提升LLM大语言模型性能的关键因素，但也会增加计算成本。Mixture of Experts (MoE) 架构通过分布式专家层和动态门控机制，有效降低了计算资源，使模型能够在扩展参数规模的同时保持高效的运行。

Mixtral of Experts

代码地址：mistral-inference/src/mistral_inference at main · mistralai/mistral-inference

Mixtral 8x7B 模型配置：

MoE层细节：

混合专家层定义为为 { E 0 , E i , . . . , E n − 1 } \{E_0, E_i, ... , E_{n-1}\} {E0​,Ei​,...,En−1​}，路由层定义为 G G G，计算公式如下：
∑ i = 0 n − 1 G ( x ) i ⋅ E i ( x ) \sum^{n-1}_{i=0}G(x)_i \cdot E_i(x) i=0∑n−1​G(x)i​⋅Ei​(x)
在Mixtral中，每个专家层都是一个FFN。路由层提供不同专家的权重，与专家层的输出加权求和，得到MoE的输出。

路由层也是一个线性层，路由层的输出维度等于专家数量。定义 W g W_g Wg​定义为路由层权重，其形状为(dim, n_experts)
G ( x ) = S o f t m a x ( T o p K ( x ⋅ W g ) ) G(x)=Softmax(TopK(x\cdot W_g)) G(x)=Softmax(TopK(x⋅Wg​))
当top_k_experts=2时，选取得分排名前2的专家，沿着num_experts维度计算Softmax

Mixtral MoE 代码实现（Pytorch）：

class MoE(nn.Module):
    def __init__(self,num_experts: int,num_experts_per_tok: int,**kwargs,):
        super().__init__()
        # 初始化专家，例如8个专家分给2个GPU
        self.experts = nn.ModuleList([FeedForward(**kwargs).to(f"cuda:{i//4}") for i in range(num_experts)])
        # 门控线性层
        self.gate = nn.Linear(kwargs["dim"], num_experts, bias=False)
        # 路由的专家数量
        self.num_experts_per_tok = num_experts_per_tok

    def forward(self, x):
        orig_shape = x.shape
        # (b, n, d) -> (b*n, d)
        x = x.view(-1, x.shape[-1])
		
        # shape: (b*n, num_experts)
        scores = self.gate(x)
        # (b*n, k), 一个是权重，一个是索引
        expert_weights, expert_indices = torch.topk(scores, self.num_experts_per_tok, dim=-1)
        # 归一化，k个权重和为1
        expert_weights = expert_weights.softmax(dim=-1)
        # (b*n, k) -> (b*n*k,)
        flat_expert_indices = expert_indices.view(-1)
		
        # (b*n, d) -> (b*n*k, d)
        x = x.repeat_interleave(self.num_experts_per_tok, dim=0)
        y = torch.empty_like(x)
        for i, expert in enumerate(self.experts):
            # 根据索引进行路由，将每个token输入索引对应的专家
            y[flat_expert_indices == i] = expert(x[flat_expert_indices == i])
        # (b*n, k, d) * (b*n, k, 1) -> (b*n, d)
        y = (y.view(*expert_weights.shape, -1) * expert_weights.unsqueeze(-1)).sum(dim=1)
        return y.view(*orig_shape)  # (b*n, d) -> (b, n, d)

DeepSeekMoE

代码地址：modeling_deepseek.py · deepseek-ai/deepseek-moe-16b-base at main (huggingface.co)

针对两个问题：

提出解决方案：

DeepSeekMoE 模型配置：

Relative Expert Size是相对于标准FFN层的维度缩放，当专家数量增加时，有必要降低每个专家的参数。当专家数为64时，将每个FFN的参数量降低为原本的1/4，则可以保持和原本16个专家相当的参数。

DeepSeekMoE 实现细节：

负载均衡损失

举例说明，假设有6个token，4个专家，topK为2，且不考虑前面的常数项

负载均衡时，每个专家上的token数量近似

此时，负载均衡损失为 3 ∗ ( 0.24 + 0.26 + 0.25 + 0.25 ) = 3 3*(0.24+0.26+0.25+0.25)=3 3∗(0.24+0.26+0.25+0.25)=3

当负载不均衡时，tokens主要集中在专家1和2

此时，负载均衡损失为 6 ∗ 0.47 + 4 ∗ 0.32 + 1 ∗ 0.10 + 1 ∗ 0.11 = 4.57 6*0.47+4*0.32+1*0.10+1*0.11=4.57 6∗0.47+4∗0.32+1∗0.10+1∗0.11=4.57

deepseek-moe-16b-base代码阅读：

Gate network

class MoEGate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.top_k = config.num_experts_per_tok  # topK
        self.n_routed_experts = config.n_routed_experts  # num of experts

        self.scoring_func = config.scoring_func  # softmax
        self.alpha = config.aux_loss_alpha  # 负载均衡损失的权重

        # 门控神经网络
        self.gating_dim = config.hidden_size
        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))

    def forward(self, hidden_states):
        bsz, seq_len, h = hidden_states.shape        
        hidden_states = hidden_states.view(-1, h)
        # 计算每个专家的得分
        logits = F.linear(hidden_states, self.weight, None)
        if self.scoring_func == 'softmax':
            scores = logits.softmax(dim=-1)
        else:
            raise NotImplementedError(f'insupportable scoring function for MoE gating: {self.scoring_func}')
        
        # 选取topk专家，shape: (bsz * seq_len, k)
        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)
        
        # 计算负载均衡损失
        scores_for_aux = scores
        aux_topk = self.top_k
        topk_idx_for_aux_loss = topk_idx.view(bsz, -1)
        # (bsz * seq_len, n_routed_experts) -> (n_routed_experts, )
        Pi = scores_for_aux.mean(0)
        
        # 将分配的专家索引转为one-hot向量，shape: (bsz * seq_len * k, n_routed_experts)
        mask_ce = F.one_hot(topk_idx_for_aux_loss.view(-1), num_classes=self.n_routed_experts)
        ce = mask_ce.float().mean(0)
        fi = ce * self.n_routed_experts
        
        aux_loss = (Pi * fi).sum() * self.alpha
        return topk_idx, topk_weight, aux_loss

DeepSeekMoE layer

class DeepseekMoE(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_experts_per_tok = config.num_experts_per_tok
        self.experts = nn.ModuleList([DeepseekMLP(config, intermediate_size = config.moe_intermediate_size) for i in range(config.n_routed_experts)])
        self.gate = MoEGate(config)
        if config.n_shared_experts is not None:
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
            self.shared_experts = DeepseekMLP(config=config, intermediate_size = intermediate_size)
    
    def forward(self, hidden_states):
        identity = hidden_states
        orig_shape = hidden_states.shape
        # shape: (bsz * seq_len, k)
        topk_idx, topk_weight, aux_loss = self.gate(hidden_states)
        # (bsz, seq_len, d) -> (bsz*seq_len, d)
        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
        flat_topk_idx = topk_idx.view(-1)
        # 复制输入，方便使用for循环进行处理
        hidden_states = hidden_states.repeat_interleave(self.num_experts_per_tok, dim=0)
        y = torch.empty_like(hidden_states)
        # 计算路由专家的输出
        for i, expert in enumerate(self.experts):
            y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])
            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
            y =  y.view(*orig_shape)
            y = AddAuxiliaryLoss.apply(y, aux_loss)
        # 加上共享专家的输出
        if self.config.n_shared_experts is not None:
            y = y + self.shared_experts(identity)
        return y

从代码来看，不管是细粒度专家，还是共享专家，在专家本身的结构上都没有变化，变化的是专家负责的语义。一个是将语义细化，产生更多的组合去适应丰富的场景；另一个是将通用语义抽取出来，让其余的细粒度专家更专注于本身应该负责的语义。

DeepSeekMoE V2

代码未开源，这里放一个权重和配置文件地址：deepseek-ai/DeepSeek-V2 at main (huggingface.co)

论文的主要贡献是MLA (Multi-Head Latent Attention): Equipped with low-rank key-value joint compression, boosting inference efficiency.

为了提升推理速度，大型语言模型（LLM）采用Key-Value (KV) 缓存来存储中间结果，但这增加了存储负担，限制了批量大小和序列长度。

为降低KV缓存，Transformer中自注意力机制的发展历史：MHA，GQA，MQA，MLA

对比不同的注意力机制， n h n_h nh​分别表示头的数量， d h d_h dh​表示每个头的维度， l l l表示Transformer的层数， n k v n_{kv} nkv​表示GQA中kv头的数量

核心思想：将key和value从高维空间映射到低维空间，缓存投影矩阵和低维空间的矩阵，每次推理的时候重新计算key和value

q , k , v q, k, v q,k,v 经过线性投影层后，维度为 n h d h n_hd_h nh​dh​
q = W Q h ,   k = W K h ,   v = W V h q=W^Qh,\ k=W^Kh,\ v=W^Vh q=WQh, k=WKh, v=WVh
注意力机制计算公式：
o = S o f t m a x ( q T k d h ) v u = W O o o=Softmax(\frac{q^Tk}{\sqrt{d_h}})v \\ u=W^Oo o=Softmax(dh​ ​qTk​)vu=WOo
MLA使用低秩矩阵 c K V c^{KV} cKV进行压缩，维度为 d c < < d h d_c<<d_h dc​<<dh​，推理时将 c K V c^{KV} cKV映射回高维矩阵
c K V = W D K V h k C = W U K c K V ,   v C = W U V c K V c^{KV}=W^{DKV}h \\ k^C=W^{UK}c^{KV},\ v^C=W^{UV}c^{KV} cKV=WDKVhkC=WUKcKV, vC=WUVcKV
算子融合：上映射这一步的算子 W U K W_{UK} WUK​可以与 W Q W_Q WQ​融合， W U V W_UV WU​V可以与 W O W_O WO​融合，从而不需要计算key和value，直接输入 c K V c^{KV} cKV输出 u u u

此外，为了降低存储资源，论文甚至在训练过程中对query使用了低秩矩阵，即使略微增加了计算量

然而，算子融合也带了问题，算与旋转位置编码(ROPE)不兼容，因为ROPE必须直接作用在query和key上

解决方法：如下图，将query和key分成两组，分别用MLA和MQA，只在MQA这一组使用ROPE

修正一下前面的KV 缓存，MLA真实的KV缓存应该为 ( n h d c + d R ) l (n_hd_c+d_R)l (nh​dc​+dR​)l，其中 d R d_R dR​为MQA中共享k的维度

举例说明，假设有8个头，每个头的维度为64，GAQ中 n k v = 4 n_{kv}=4 nkv​=4

使用for循环对专家进行路由是非常低效的，如果做高效推理，通常会在并行策略（数据并行、模型并行、专家并行）中重新实现token的重排和分发

参考资料：

混合专家模型（MoE）详解 (huggingface.co)

Mixtral of Experts

DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models

DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

总结