解读Qwen1.5 MoE：稀疏大模型的高效智能

显示全部楼层

引言

官方文档：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模型的结构和关键组件。

模型构造 (__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机制

classQwen2MoeSparseMoeBlock(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_inrange(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)

defforward(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_logitsbatch*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)
ifself.norm_topk_prob:
routing_weights/=routing_weights.sum(dim=-1,keepdim=True)
#wecastbacktotheinputdtype
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
)

#Onehotencodetheselectedexpertstocreateanexpertmask
#thiswillbeusedtoeasilyindexwhichexpertisgoingtobesollicitated
expert_mask=torch.nn.functional.one_hot(selected_experts,num_classes=self.num_experts).permute(2,1,0)

#Loopoverallavailableexpertsinthemodelandperformthecomputationoneachexpert
forexpert_idxinrange(self.num_experts):
expert_layer=self.experts[expert_idx]
idx,top_x=torch.where(expert_mask[expert_idx])

iftop_x.shape[0]==0:
continue

#intorchitisfastertoindexusingliststhantorchtensors
top_x_list=top_x.tolist()
idx_list=idx.tolist()

#Indexthecorrecthiddenstatesandcomputetheexperthiddenstatefor
#thecurrentexpert.Weneedtomakesuretomultiplytheoutputhidden
#statesby`routing_weights`onthecorrespondingtokens(top-1andtop-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_`onlysupporttorchtensorsforindexingsowe'lluse
#the`top_x`tensorhere.
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)
returnfinal_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模型在训练成本和推理时间上都取得了显著的降低。