聊聊Mixtral 8x7B结构 & MoE模型

Mixtral 8x7B

Mixtral 8x7B的命名来源：8个Mixtral 7B大模型基于MoE(Mixture of Experts 混合专家模型)模型堆叠起来。

混合专家模型

混合专家模型（Mixture of Experts，简称MoE）是一种先进的深度学习技术，它通过将输入任务分配给多个专家模型，提高了模型的计算效率和性能。这种模型结构特别适用于处理多样化的数据集，能够让每个专家模型专注于其最擅长的任务，进而实现模型的稀疏性。

专家即某种模型

MoE 层由众多专家组成，既可以是小型的多层感知机（MLP），也可以是像 Mistral 7B 这样复杂的大型语言模型（LLM）；因此也需要某种机制将输入的token分配给对应的专家处理，即路由器。

与transformer模型的融合

MoE的核心特征在transformer模型中，表现为两个元素：稀疏MoE层和门控网络(路由器)。

稀疏MoE层

与transformer模型中的FFN层不同，MoE采用稀疏MoE层替代。在稀疏MoE层中，每层有多个”Expert”，每个”Expert”都是一个FFN层。再有一个门控网络(路由器)将输入token转发至对应的”Expert”中执行。

展开其MoE层，其结构如下：

对应的源码如下：

class MixtralSparseMoeBlock(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.hidden_dim = config.hidden_size
        self.ffn_dim = config.intermediate_size
        self.num_experts = config.num_local_experts
        self.top_k = config.num_experts_per_tok

        # MoE Block: gating + Expert
        # gating
        self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
        # Experts: 堆叠多个专家
        self.experts = nn.ModuleList([MixtralBLockSparseTop2MLP(config) for _ in range(self.num_experts)])

transformer模型中的FFN层，对应于模型源码(GLM、Llama2、Mixtral)的实现中，一般是MLP层。

而对于MLP层其实就很熟悉了，无非就是多个密集神经网络的堆叠。其源码如下

class MixtralBLockSparseTop2MLP(nn.Module):
    def __init__(self, config: MixtralConfig):
        super().__init__()
        self.ffn_dim = config.intermediate_size
        self.hidden_dim = config.hidden_size

        self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
        self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
        self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)

        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, hidden_states):
        current_hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states)
        current_hidden_states = self.w2(current_hidden_states)
        return current_hidden_states