Joomla外贸网站模板,网站文章来源seo,中国建筑校园招聘官网,网架加工价格1. 背景
根据本qiang~最新的趋势观察#xff0c;基于MoE架构的开源大模型越来越多#xff0c;比如马斯克的Grok-1(314B), Qwen1.5-MoE-A2.7B等#xff0c;因此想探究一下MoE里面的部分细节。
此文是本qiang~针对大语言模型的MoE的整理#xff0c;包括原理、流程及部分源码…1. 背景
根据本qiang~最新的趋势观察基于MoE架构的开源大模型越来越多比如马斯克的Grok-1(314B), Qwen1.5-MoE-A2.7B等因此想探究一下MoE里面的部分细节。
此文是本qiang~针对大语言模型的MoE的整理包括原理、流程及部分源码。
2. MoE原理
MoE的流行源于”欧洲的OpenAI” Mistral AI发布的论文及模型《Mixtral of Experts》评测集上的效果吊打众多开源模型如Llama 2 70B和GPT3.5。
《Mixtral of Experts》基础模型使用的是Mistral AI自研的Mistral 7B该模型的特点包括滑窗注意力(Sliding Window Aattention), 滚动缓冲区缓存(Rolling Buffer Cache)以及预填充-分块(Pre-fill and Chunking)具体细节可以查阅文末的论文地址。
本文以《Mixtral of Experts》为引子探究MoE的相关细节MoE的原理如下图所示 图2.1 MoE的原理
(1) Transformers架构中的每一层中的FFN网络均替换为了8个FFN(专家)且由一个网关路由(gate router)进行控制
(2) 针对每一个token每一层的网关路由仅选择其中的2个FFN(专家)来处理当前状态并进行加权输出
(3) 结果就是每一个token访问了47B参数但是在推理阶段仅仅使用了13B的激活参数(即只使用2个专家冻结其他6个专家)。
(4) 与Dropout机制对比Dropout让部分神经元失活而MoE是让部分专家失活。
3. 源码
本qiang~研读并尝试执行了Mistral官网的github推理代码该代码框架非常适合新手无他只因其几乎只是在torch上层做的封装很少引擎其他第三方库不像transformers功能强大但不适合新手研读代码…
为了普适性下面的代码截取了transformers框架中的代码。
首先看下通用Transformers中FFN中的代码模块代码位置在transformers.models.mistral.modeling_mistral, 主要流程是
(1) 先经过gate_proj和up_proj的2个[hidden_size, intermediate_size]的线性转换
(2) 使用激活函数对gate_proj进行激活
(3) 二者的内积再经过down_proj线性转换。
class MistralMLP(nn.Module):def __init__(self, config):super().__init__()self.config configself.hidden_size config.hidden_sizeself.intermediate_size config.intermediate_sizeself.gate_proj nn.Linear(self.hidden_size, self.intermediate_size, biasFalse)self.up_proj nn.Linear(self.hidden_size, self.intermediate_size, biasFalse)self.down_proj nn.Linear(self.intermediate_size, self.hidden_size, biasFalse)self.act_fn ACT2FN[config.hidden_act]def forward(self, x):return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))再来看下MoE中的专家模块代码位置在transformers.models.mixtral.modeling_mixtral主要流程是
(1) 首先经过网关路由self.gate
(2) 然后选择其中2个专家并归一化
(3) 之后遍历每个专家网络并按照expert_mask进行筛选
(4) 如果expert_mask有值则选择指定部分的隐藏层进行FFN操作且输出结果进行加权
(5) 最后原地增加先前初始化的最终结果变量final_hidden_states
class MixtralSparseMoeBlock(nn.Module):def __init__(self, config):super().__init__()self.hidden_dim config.hidden_sizeself.ffn_dim config.intermediate_sizeself.num_experts config.num_local_expertsself.top_k config.num_experts_per_tok# gatingself.gate nn.Linear(self.hidden_dim, self.num_experts, biasFalse)self.experts nn.ModuleList([MixtralBlockSparseTop2MLP(config) for _ in range(self.num_experts)])def forward(self, hidden_states: torch.Tensor) - torch.Tensor: batch_size, sequence_length, hidden_dim hidden_states.shapehidden_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, dim1, dtypetorch.float)routing_weights, selected_experts torch.topk(routing_weights, self.top_k, dim-1)routing_weights / routing_weights.sum(dim-1, keepdimTrue)# we cast back to the input dtyperouting_weights routing_weights.to(hidden_states.dtype)final_hidden_states torch.zeros((batch_size * sequence_length, hidden_dim), dtypehidden_states.dtype, devicehidden_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 sollicitatedexpert_mask torch.nn.functional.one_hot(selected_experts, num_classesself.num_experts).permute(2, 1, 0)# Loop over all available experts in the model and perform the computation on each expertfor 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 tensorstop_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 well use# the top_x tensor here.final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))final_hidden_states final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)return final_hidden_states, router_logits其中MixtralBlockSparseTop2MLP代码如下可以看到和传统MistralMLP内容完全一致。
class MixtralBlockSparseTop2MLP(nn.Module):def __init__(self, config: MixtralConfig):super().__init__()self.ffn_dim config.intermediate_sizeself.hidden_dim config.hidden_sizeself.w1 nn.Linear(self.hidden_dim, self.ffn_dim, biasFalse)self.w2 nn.Linear(self.ffn_dim, self.hidden_dim, biasFalse)self.w3 nn.Linear(self.hidden_dim, self.ffn_dim, biasFalse)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_states4. MoE微调
由于MoE只是将每一层的FFN改变为了每一层的gate网关路由8个FFN专家且gate网关路由和8个专家内部均为线性运算所以可以无缝地结合LoRA、QLoRA进行指令微调。
可以参考开源项目https://github.com/yangjianxin1/Firefly
5. 答疑解惑
(1) 问MoE 8*7B的模型是56B参数
答MoE 8*7B的参数量是47B而不是56B原因是每一层除了8个专家网络外其他层均是复用的。
(2) 问MoE的基础模型是Mistral 7B?
答不是MoE的模型架构与Mistral 7B相同但其中的FFN替换为了8个FFN且MoE是基于多语言数据集预训练而来的。
(3) MoE的稀疏性(sparse)体现在哪里
答在训练和推理时同时只有两个专家网络会被激活进行前向计算其它专家网络处于失活状态。
6. 总结
一句话足矣~
本文主要针对大语言模型的MoE包括原理及部分源码。
此外建议大家可以针对源码进行运行关于源码欢迎大家一块交流。
7. 参考
(1) Mistral 7Bhttps://arxiv.org/pdf/2310.06825v1.pdf
(2) MoE: https://arxiv.org/pdf/2401.04088v1.pdf
(3) MoE开源指令微调框架Firefly: https://github.com/yangjianxin1/Firefly
