引言

官方文档：Qwen1.5-MoE: 1/3的激活参数量达到7B模型的性能 | Qwen

3月28日，阿里首次宣布开源MoE技术大模型Qwen1.5-MoE-A2.7B，这个模型以现有的Qwen-1.8B模型为基础。Qwen1.5-MoE-A2.7B激活参数为2.7亿，但在一系列基准评估中可以达到7B模型的性能。此外，与7B模型相比，它在训练成本和推理速度上具有显著优势。据官方评测显示，Qwen1.5-MoE-A2.7B在与最佳的7B模型相比取得了非常接近的性能。本文将根据官方博客内容与开放代码，针对Qwen1.5-MoE-A2.7B进行解读。

文章结构如下：

参数量

Qwen1.5-MoE-A2.7B。它仅拥有27亿个激活参数，但其性能却能与当前最先进的70亿参数模型，如Mistral 7B和Qwen1.5-7B相媲美。相较于包含65亿个Non-Embedding参数的Qwen1.5-7B，Qwen1.5-MoE-A2.7B只有20亿个Non-Embedding参数，约为原模型大小的三分之一。此外，相比Qwen1.5-7B，Qwen1.5-MoE-A2.7B的训练成本降低了75%，推理速度则提升至1.74倍。

Qwen1.5-MoE-A2.7B在与最佳的7B模型相比取得了非常接近的性能。

模型结构

Qwen1.5-MoE模型中采用了特别设计的MoE架构。通常情况下，如Mixtral等方法所示，每个transformer block中的MoE层会配备8个expert，并采用top-2门控策略进行routing。这种配置还存在很大的优化空间。Qwen1.5-MoE对这一架构进行了多项改进：

• Finegrained experts
• 初始化
• 新的routing机制

DeepSeek-MoE和DBRX已经证明了finegrained experts的有效性。从FFN层过渡到MoE层时，我们一般只是简单地复制多次FFN来实现多个expert。而finegrained experts的目标是在不增加参数数量的前提下生成更多expert。为了实现这一点，我们将单个FFN分割成几个部分，每个部分作为一个独立的expert。我们设计了具有总共64个expert的的MoE，对比其他配置，我们认为这个实现能达到效果和效率的最优。

模型初始化阶段至关重要。初步实验表明，从零开始训练MoE模型可能效率低下，且难以提升至预期的最优性能水平。因此，我们首先利用已有的Qwen-1.8B，将其改造为Qwen1.5-MoE-A2.7B。此外，在初始化阶段引入随机性可以显著加快收敛速度，并在整个预训练过程中带来更好的整体性能表现。

目前，一个明显的趋势是在MoE中实现共享expert与routing expert。从更宏观的角度看，这是一种广义的routing方法，因为在没有共享expert的情况下，实际上就退化为传统的MoE路由设置。对于Qwen1.5-MoE-A2.7B模型，我们在其中整合了4个总是被激活的共享expert和每次只激活其中4个的60个routing expert。这种方式非常灵活，同时在我们实验中效率最佳。

• “类似”8*1.8B TOP2激活 MOE：这里的"类似"并不是真的8*1.8B，而是采用Finegrained experts，总共64个expert，激活8个“Finegrained expert”

• Finegrained Experts：参照DeepSeek-MOE,DBRX，将单个FFN拆分成几部分，总共64个expert，激活8个“Finegrained expert”

• 初始化：基于QWen-1.8B初始化，并引入随机性提高收敛速度，但是随机性引入在哪里未体现在博客与代码中

• Routing机制：共享Expert和Routing Expert，整合了4个总是被激活的共享expert和每次只激活其中4个的60个routing expert

训练成本与推理效率

MoE模型的训练成本与dense模型存在显著差异。尽管MoE模型通常拥有更多的参数，但由于其稀疏性，训练开销可以显著降低。首先先对比各个模型的三个关键参数，分别是总参数数量、激活参数数量和Non-embedding参数：

• 尽管MoE模型的总参数量较大，但Non-embedding激活参数量远小于7B模型。

• 在实践中，使用Qwen1.5-MoE-A2.7B相比于Qwen1.5-7B，训练成本显著降低了75%。

• 由于初始化方法，不需要训练同样数量的token即可达到很好的模型效果，这也显著了降低了训练成本

代码解读

代码地址：https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_moe/modeling_qwen2_moe.py

config解读

config配置参数：

{
  "architectures": [
    "Qwen2MoeForCausalLM"
  ],
  "attention_dropout": 0.0,
  "bos_token_id": 151643,
  "eos_token_id": 151643,
  "hidden_act": "silu",
  "hidden_size": 2048,
  "initializer_range": 0.02,
  "intermediate_size": 5632,
  "max_position_embeddings": 8192,
  "max_window_layers": 21,
  "model_type": "qwen2_moe",
  "num_attention_heads": 16,
  "num_hidden_layers": 24,
  "num_key_value_heads": 16,
  "rms_norm_eps": 1e-06,
  "rope_theta": 1000000.0,
  "sliding_window": 32768,
  "tie_word_embeddings": false,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.39.0.dev0",
  "use_cache": true,
  "use_sliding_window": false,
  "vocab_size": 151936,
  "decoder_sparse_step": 1,
  "moe_intermediate_size": 1408,
  "shared_expert_intermediate_size": 5632,
  "num_experts_per_tok": 4,
  "num_experts": 60,
  "norm_topk_prob": false,
  "output_router_logits": false,
  "router_aux_loss_coef": 0.001
}

Config文件中可以窥见许多设计模型的细节：

• hidden_act：激活函数采用silu

• 支持的上下文长度：8K

• rope_theta：100W

• tie_word_embeddings：非参数共享

• moe_intermediate_size，shared_expert_intermediate_size：这里可以对应上blog中4个共享Expert

• num_experts_per_tok，num_experts：对应blog中60个expert选4个expert

• router_aux_loss_coef：处理专家之间均衡的系数

模型主结构介绍

Qwen2MoeForCausalLM模型是基于“混合专家(Mixture of Experts, MoE)”架构的变体，其中使用了一种稀疏激活策略，即每个输入序列只激活少数专家（expert）。下面我们会逐步理解整个Qwen2MoeForCausalLM模型的结构和关键组件。

1. 模型构造 (__init__方法):

   模型的初始化方法首先调用了它的父类Qwen2MoePreTrainedModel的初始化，然后创建了以下关键结构：

• self.model: 这是Qwen2MoeModel的一个实例，它本身是由多个Qwen2MoeDecoderLayer建立的解码器模型。每个解码层(Qwen2MoeDecoderLayer)包含注意力层(Qwen2MoeAttention 或其变种)和前馈网络(在本例中可能是Qwen2MoeMLP或Qwen2MoeSparseMoeBlock)。

• self.lm_head: 这是一个线性层，它将解码器的输出映射为词汇表长度维度的logits。换句话说，它将最后一层隐藏状态转换为预测每个词汇索引的未归一化的分数。

• self.router_aux_loss_coef: 这是用于计算路由器辅助损失的系数。该损失用于平衡专家之间的负荷，防止某些专家过载而其他专家则鲜少被选中。

• 前向传播 (forward方法):

  forward方法指导数据如何通过模型。以下是其主要步骤：

• 调用self.model执行解码器模型的前向传播，这将依序通过每个解码层处理输入数据。

• 解码器输出的隐藏状态被送入self.lm_head，产生对词汇表每个单词概率的logits。

• 如果提供了labels，则计算损失。对于因果语言模型任务（Causal Language Modeling，CLM），这通常是交叉熵损失。

• 如果forward方法收到output_router_logits的指示，它将计算和返回用于损失函数中的路由器logits，或者说这是路由每个token到不同专家的概率分布。

这里是一些模型的特色部分：

• Mixture of Experts (MoE): Qwen2MoeSparseMoeBlock豪华表示了MoE的实现，它根据动态确定的路由决策将计算路由至不同的专家。

• 共享专家: 代码中通过self.shared_expert定义了一个共享的全连接层，这个层会在所有输入上都应用其变换（即总是被激活）。

• 稀疏激活: Qwen2MoeSparseMoeBlock的forward方法中使用routing_weights和selected_experts实现了专家的稀疏激活，这意味着对于每个输入，只有少数权重最高的专家被选择和激活。

• 路由器辅助损失: 通过调用load_balancing_loss_func，模型计算了一个附加损失，以促进在所有专家之间均匀分配计算负载。这也帮助模型避免依赖于少数专家，从而利用分布在多个专家上的知识。

• 参数初始化: _init_weights方法确保模型的权重以一种提高收敛速度的方式被初始化。这往往涉及到给定一定范围的随机性，来避免参数被初始化为相同的值，这可以通过使用正态分布来实现。

  在整个Qwen2MoeForCausalLM模型中，由于解码器每一层都可能包含MoE结构，因此这个模型可以通过动态路由机制大大增加模型的容量，而不会线性增加计算成本。模型的稀疏激活策略保证了在每个时间步只有少量的专家被调用，这极大地提升了参数利用率和效率。

  总结来说，Qwen2MoeForCausalLM通过以下特点体现其创新和功能：

  在Qwen2MoeForCausalLM模型的实际应用中，可能还会包括不同的解码器层选择和训练策略等高级功能，以解决特定的NLP任务或提升模型性能。

• **混合专家架构 (MoE)**：多个专家网络可以让模型更好地扩展和专门化。

• 稀疏激活：每个输入只激活少量的专家，这降低了运行时的计算复杂性。

• 共享与路由专家的结合：共享专家始终激活并处理每个token，增加了模型的普适性，而路由专家提供了选择性的知识处理。

• 辅助损失：提供了一种激励机制，确保负载在专家之间平衡分布，这有助于模型的广义性和稳定性。

• 权重初始化：采用特定的随机初始化策略，以提高收敛速度。

Routing机制

class Qwen2MoeSparseMoeBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.num_experts = config.num_experts
        self.top_k = config.num_experts_per_tok
        self.norm_topk_prob = config.norm_topk_prob

        # gating
        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
        self.experts = nn.ModuleList(
            [Qwen2MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
        )

        self.shared_expert = Qwen2MoeMLP(config, intermediate_size=config.shared_expert_intermediate_size)
        self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """ """
        batch_size, sequence_length, hidden_dim = hidden_states.shape
        hidden_states = hidden_states.view(-1, hidden_dim)
        # router_logits: (batch * sequence_length, n_experts)
        router_logits = self.gate(hidden_states)

        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
        if self.norm_topk_prob:
            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
        # we cast back to the input dtype
        routing_weights = routing_weights.to(hidden_states.dtype)

        final_hidden_states = torch.zeros(
            (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
        )

        # One hot encode the selected experts to create an expert mask
        # this will be used to easily index which expert is going to be sollicitated
        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)

        # Loop over all available experts in the model and perform the computation on each expert
        for expert_idx in range(self.num_experts):
            expert_layer = self.experts[expert_idx]
            idx, top_x = torch.where(expert_mask[expert_idx])

            if top_x.shape[0] == 0:
                continue

            # in torch it is faster to index using lists than torch tensors
            top_x_list = top_x.tolist()
            idx_list = idx.tolist()

            # Index the correct hidden states and compute the expert hidden state for
            # the current expert. We need to make sure to multiply the output hidden
            # states by `routing_weights` on the corresponding tokens (top-1 and top-2)
            current_state = hidden_states[None, top_x_list].reshape(-1, hidden_dim)
            current_hidden_states = expert_layer(current_state) * routing_weights[top_x_list, idx_list, None]

            # However `index_add_` only support torch tensors for indexing so we'll use
            # the `top_x` tensor here.
            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))

        shared_expert_output = self.shared_expert(hidden_states)
        shared_expert_output = F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output

        final_hidden_states = final_hidden_states + shared_expert_output

        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
        return final_hidden_states, router_logits

这段代码定义了一个称为Qwen2MoeSparseMoeBlock的模块，该模块是Qwen2Moe架构中实现混合专家（Mixture of Experts，MoE）机制的关键部分。MoE结构允许网络将不同部分的计算任务动态地分配给专家子模型（expert sub-models），其中每个专家专注于模型的一小部分。

下面是代码的逐行解释：

__init__ 方法（初始化）

• num_experts: 配置参数，指定要使用的专家数量。
• top_k: 为每个token选择的专家数量。
• norm_topk_prob: 一个标志，指是否对选择的top_k专家的权重进行归一化处理。
• gate: 一个线性层，负责根据输入的hidden states生成每个专家的logits，用于后续的专家选择。
• experts: 一个模块列表，包含所有按配置创建的专家网络。每个专家是一个Qwen2MoeMLP对象。
• shared_expert: 另一个Qwen2MoeMLP对象，表示共享的专家，它总是被激活的，不参与动态路由。
• shared_expert_gate: 线性层，用于控制共享专家在最终输出中的权重。

forward 方法（前向传播）

• 首先，将进入到此模块的隐藏状态张量的形状改为 [batch_size * sequence_length, hidden_dim]。
• router_logits 是由上面定义的 gate 线性层计算出来的，它决定输入的hidden states将如何分配给不同的专家。
• 接着通过softmax函数得到每个专家的routing权重。
• selected_experts 选取了top_k权重最大的专家。
• 如果norm_topk_prob为真，将通过每一行（每个token）的权重和进行归一化。
• 创建一个全零的final_hidden_states张量用来累积各个专家的计算结果。

接下来进入到处理每个专家的循环:

• 对每个专家，它首先检查是否有token被分配给该专家（如果没有，则跳过）。
• expert_mask 是一个one-hot编码，用来标识哪个token被分配到了哪个专家。
• 之后代码将各个专家应用到其对应的token上，并根据routing_weights缩放结果。
• final_hidden_states.index_add_ 通过加法累积更新对应的token预测。

共享专家的处理:

• shared_expert 处理所有的隐藏状态，并通过 sigmoid 激活函数和 shared_expert_gate 的输出来缩放其贡献。
• 将共享专家的输出添加到 final_hidden_states。

最后:

• 将 final_hidden_states 形状变回原来的 [batch_size, sequence_length, hidden_dim]。
• 返回 final_hidden_states 和 router_logits.

整个流程实现了混合专家机制，在某一层将隐藏状态注入各个专家和一个共享专家，并汇总输出作为下一层输入的机制。这样可以显著增加模型表示能力的同时保持效率。

重要的是注意，forward 方法返回了两个结果: final_hidden_states 表示了由专家处理过的隐藏状态，router_logits 表示了用于专家选择的原始logits。后者可以用来计算辅助损失（如上述代码所提）、平衡专家之间的工作负载。

重点TIPS:

整合了4个总是被激活的共享expert和每次只激活其中4个的60个routing expert:

• routing 60选4是很显而易见的，self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)，通过gate实现topk的筛选

• 4个共享expert但是上述代码里面只有一个shared_expert_output = self.shared_expert(hidden_states)，这是因为在config中shared_expert隐藏层的大小已经设置为4倍的普通expert大小

总结

阿里首个MoE模型Qwen1.5-MoE-A2.7B，它对比当前最好的7B参数模型取得了相当的性能。此外，与传统的7B模型相比，MoE模型在训练成本和推理时间上都取得了显著的降低。