前言
随着MoE越来越火热,MoE本质就是将Transformer中的FFN层替换成了MoE-layer,其中每个MoE-Layer由一个gate和若干个experts组成。这里gate和每个expert都可以理解成是nn.linear形式的神经网络。既然如此,本篇文章将结合transformers结构构建一个MoE的demo供大家学习。该源码可直接使用。
一、MoE原理与设计原则
expert:术业有专攻。假设我的输入数据是“我爱吃炸鸡”,在原始的Transformer中,我们把这5个token送去一个FFN层做处理。但是现在我们发现这句话从结构上可以拆成“主语-我”,“谓语-爱吃”,“宾语-炸鸡”,秉持着术业有专攻的原则,我把原来的1个FFN拆分成若干个expert,分别用来单独解析“主语”,“谓语”,“宾语”,这样可能会达到更好的效果。
gate:那么我怎么知道要把哪个token送去哪个expert呢?很简单,我再训练一个gate神经网络,让它帮我判断就好了。
当然,这里并不是说expert就是用来解析主谓宾,只是举一个例子说明:不同token代表的含义不一样,因此我们可以用不同expert来对它们做解析。除了训练上也许能达到更好的效果外,MoE还能帮助我们在扩大模型规模的同时保证计算量是非线性增加的(因为每个token只用过topK个expert,不用过全量expert),这也是我们说MoE-layer是稀疏层的原因。
最后需要注意的是,在之前的表述中,我们说expert是从FFN层转变而来的,这很容易让人错理解成expert就是对FFN的平均切分,实际上你可以任意指定每个expert的大小,每个expert甚至可以>=原来单个FFN层,这并不会改变MoE的核心思想:token只发去部分expert时的计算量会小于它发去所有expert的计算量。
引用:https://mp.weixin.qq.com/s/76a-7fDJumv6iB08L2BUKg
二、构建完整transformers与MoE集成模块
这个模块定义了一个名为Block的PyTorch模块,代表了一个混合专家Transformer块,包括多头自注意力和计算MoE部分(SparseMoE)。其源码如下:
class Block(nn.Module):
"""Mixture of Experts Transformer block: communication followed by computation (multi-head self attention + SparseMoE) """
def __init__(self, n_embed, n_head, num_experts, top_k):
super().__init__()
self.sa = nn.MultiheadAttention(n_embed, n_head) # 我直接调用官网的attention方法
self.smoe = SparseMoE(n_embed, num_experts, top_k)
self.ln1 = nn.LayerNorm(n_embed)
self.ln2 = nn.LayerNorm(n_embed)
def forward(self, x):
qkv = self.ln1(x)
x = x + self.sa(qkv,qkv,qkv)[0] # 并不包含FFN结构
x = x + self.smoe(self.ln2(x))
return x
这段代码的解读:
__init__方法:在初始化方法中,定义了模块的结构。模块包含了一个nn.MultiheadAttention实例sa用于多头自注意力计算,我直接使用pytorch官网的方法模块,一个SparseMoE实例smoe用于稀疏专家计算,以及两个nn.LayerNorm实例ln1和ln2用于层归一化。
forward方法:在前向传播方法中,首先对输入张量x进行层归一化处理,然后通过多头自注意力模块对处理后的张量进行注意力计算,并将注意力计算结果与原始输入张量相加。接着,将相加后的张量通过稀疏专家模块进行计算,再次与原始输入张量相加。最后返回处理后的张量。
总体来说,这段代码实现了一个混合专家Transformer块,结合了多头自注意力和稀疏专家计算。在前向传播过程中,通过多头自注意力和稀疏专家计算两部分对输入张量进行处理,并保持了张量的维度一致。
注: x = x + self.smoe(self.ln2(x))使用了类似残差方法
。
三、专家模块定义
专家模块定义了一个名为Expert的PyTorch模块,代表了一个专家模块,用于对输入进行线性变换和非线性变换。
class Expert(nn.Module):
def __init__(self, n_embd):
super().__init__()
dropout=0.1
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
四、路由门控模块
门控网络,也称为路由,确定哪个专家网络接收来自多头注意力的 token 的输出。举个例子解释路由的机制,假设有 4 个专家,token 需要被路由到前 2 个专家中。首先需要通过线性层将 token 输入到门控网络中。该层将对应于(Batch size,Tokens,n_embed)的输入张量从(2,4,32)维度,投影到对应于(Batch size、Tokens,num_expert)的新形状:(2、4,4)。其中 n_embed 是输入的通道维度,num_experts 是专家网络的计数。
class NoisyTopkRouter(nn.Module):
def __init__(self, n_embed, num_experts, top_k):
super(NoisyTopkRouter, self).__init__()
self.top_k = top_k
self.topkroute_linear = nn.Linear(n_embed, num_experts)
# add noise
self.noise_linear = nn.Linear(n_embed, num_experts)
def forward(self, mh_output):
# mh_ouput is the output tensor from multihead self attention block
logits = self.topkroute_linear(mh_output)
# Noise logits
noise_logits = self.noise_linear(mh_output)
# Adding scaled unit gaussian noise to the logits
noise = torch.randn_like(logits) * F.softplus(noise_logits)
noisy_logits = logits + noise
top_k_logits, indices = noisy_logits.topk(self.top_k, dim=-1) # 在4个专家中选择最好的专家
zeros = torch.full_like(noisy_logits, float('-inf'))
sparse_logits = zeros.scatter(-1, indices, top_k_logits)
router_output = F.softmax(sparse_logits, dim=-1)
return router_output, indices
五、稀疏MoE集成模块
通过路由获得indices在对每个专家循环获得对应mask,挑选有用的信息不断叠加而获得最终结果,我已改成伪代码如下说明:
# 1. 输入进入router得到两个输出
gating_output, indices = router(x)
# 2.初始化全零矩阵,后续叠加为最终结果
final_output = zeros_like(x)
# 3.展平,即把每个batch拼接到一起,这里对输入x和router后的结果都进行了展平
flat_x = flatten(x)
flat_gating_output = flatten(gating_output)
# 4. 对每个专家进行操作
for each expert in experts:
# 5. 查看当前专家对哪些tokens在前top_k
expert_mask = check_top_k_indices(indices, expert)
# 6. 获取当前专家作用的token输入
expert_input = select_expert_input(flat_x, expert_mask)
# 7. 将token输入经过专家处理得到输出
expert_output = expert(expert_input)
# 8. 计算当前专家对作用的token的权重分数
gating_scores = calculate_gating_scores(flat_gating_output, expert_mask)
# 9. 将expert输出乘上权重分数
weighted_output = expert_output * gating_scores
# 10. 将结果叠加到最终输出中
final_output += weighted_output
return final_output
总结解释:
这段代码实现了一个稀疏专家(SparseMoE)模块,结合了一个路由器(router)和多个专家(experts)。
输入首先通过路由器得到两个输出:一个是门控输出(gating_output),一个是索引(indices)。
然后初始化一个全零矩阵作为最终输出。
对每个专家进行操作,根据索引找出每个专家对哪些token起作用,然后将这些token输入到对应的专家中进行处理,根据门控输出的权重分配,将专家处理后的输出加权叠加到最终输出中。
最终返回加权叠加后的最终输出。
其源码如下:
class SparseMoE(nn.Module):
def __init__(self, n_embed, num_experts, top_k):
super(SparseMoE, self).__init__()
self.router = NoisyTopkRouter(n_embed, num_experts, top_k)
self.experts = nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])
self.top_k = top_k
def forward(self, x):
# 1. 输入进入router得到两个输出
gating_output, indices = self.router(x) # x.shape=[2,6,64] gating_output.shape=[2,6,4]
# 2.初始化全零矩阵,后续叠加为最终结果
final_output = torch.zeros_like(x) # final_output.shape=[2,6,64]
# 3.展平,即把每个batch拼接到一起,这里对输入x和router后的结果都进行了展平
flat_x = x.view(-1, x.size(-1)) # flat_x.shape=[12,64]
flat_gating_output = gating_output.view(-1, gating_output.size(-1)) # flat_gating_output.shape=[12,4]
# 以每个专家为单位进行操作,即把当前专家处理的所有token都进行加权
for i, expert in enumerate(self.experts):
# 4. 对当前的专家(例如专家0)来说,查看其对所有tokens中哪些在前top2
expert_mask = (indices == i).any(dim=-1) # expert_mask.shape=[2,6]
# 5. 展平操作
flat_mask = expert_mask.view(-1) # flat_mask=12
# 如果当前专家是任意一个token的前top2
if flat_mask.any():
# 6. 得到该专家对哪几个token起作用后,选取token的维度表示
expert_input = flat_x[flat_mask] # 假设第0个专家选择了2个batch的7个维度为[7,64]
# 7. 将token输入expert得到输出
expert_output = expert(expert_input) # 这个专家选择通道走了mlp结构,[7,64]
# 8. 计算当前专家对于有作用的token的权重分数
gating_scores = flat_gating_output[flat_mask, i].unsqueeze(1) # 使用flat_mask选择第i个专家的权重
# 9. 将expert输出乘上权重分数
weighted_output = expert_output * gating_scores
# 10. 循环进行做种的结果叠加,也就是越重要被专家选择越多就叠加越多
final_output[expert_mask] += weighted_output.squeeze(1) # 变成[7,64]
return final_output
六、完整MoE的Demo
最后,给出即插即用的MoE完整Demo代码如下:
import torch
import torch.nn as nn
import torch.nn.functional as F
class Expert(nn.Module):
def __init__(self, n_embd):
super().__init__()
dropout=0.1
self.net = nn.Sequential(
nn.Linear(n_embd, 4 * n_embd),
nn.ReLU(),
nn.Linear(4 * n_embd, n_embd),
nn.Dropout(dropout),
)
def forward(self, x):
return self.net(x)
class NoisyTopkRouter(nn.Module):
def __init__(self, n_embed, num_experts, top_k):
super(NoisyTopkRouter, self).__init__()
self.top_k = top_k
self.topkroute_linear = nn.Linear(n_embed, num_experts)
# add noise
self.noise_linear = nn.Linear(n_embed, num_experts)
def forward(self, mh_output):
# mh_ouput is the output tensor from multihead self attention block
logits = self.topkroute_linear(mh_output)
# Noise logits
noise_logits = self.noise_linear(mh_output)
# Adding scaled unit gaussian noise to the logits
noise = torch.randn_like(logits) * F.softplus(noise_logits)
noisy_logits = logits + noise
top_k_logits, indices = noisy_logits.topk(self.top_k, dim=-1) # 在4个专家中选择最好的专家
zeros = torch.full_like(noisy_logits, float('-inf'))
sparse_logits = zeros.scatter(-1, indices, top_k_logits)
router_output = F.softmax(sparse_logits, dim=-1)
return router_output, indices
class SparseMoE(nn.Module):
def __init__(self, n_embed, num_experts, top_k):
super(SparseMoE, self).__init__()
self.router = NoisyTopkRouter(n_embed, num_experts, top_k)
self.experts = nn.ModuleList([Expert(n_embed) for _ in range(num_experts)])
self.top_k = top_k
def forward(self, x):
# 1. 输入进入router得到两个输出
gating_output, indices = self.router(x) # x.shape=[2,6,64] gating_output.shape=[2,6,4]
# 2.初始化全零矩阵,后续叠加为最终结果
final_output = torch.zeros_like(x) # final_output.shape=[2,6,64]
# 3.展平,即把每个batch拼接到一起,这里对输入x和router后的结果都进行了展平
flat_x = x.view(-1, x.size(-1)) # flat_x.shape=[12,64]
flat_gating_output = gating_output.view(-1, gating_output.size(-1)) # flat_gating_output.shape=[12,4]
# 以每个专家为单位进行操作,即把当前专家处理的所有token都进行加权
for i, expert in enumerate(self.experts):
# 4. 对当前的专家(例如专家0)来说,查看其对所有tokens中哪些在前top2
expert_mask = (indices == i).any(dim=-1) # expert_mask.shape=[2,6]
# 5. 展平操作
flat_mask = expert_mask.view(-1) # flat_mask=12
# 如果当前专家是任意一个token的前top2
if flat_mask.any():
# 6. 得到该专家对哪几个token起作用后,选取token的维度表示
expert_input = flat_x[flat_mask] # 假设第0个专家选择了2个batch的7个维度为[7,64]
# 7. 将token输入expert得到输出
expert_output = expert(expert_input) # 这个专家选择通道走了mlp结构,[7,64]
# 8. 计算当前专家对于有作用的token的权重分数
gating_scores = flat_gating_output[flat_mask, i].unsqueeze(1) # 使用flat_mask选择第i个专家的权重
# 9. 将expert输出乘上权重分数
weighted_output = expert_output * gating_scores
# 10. 循环进行做种的结果叠加,也就是越重要被专家选择越多就叠加越多
final_output[expert_mask] += weighted_output.squeeze(1) # 变成[7,64]
return final_output
class Block(nn.Module):
"""Mixture of Experts Transformer block: communication followed by computation (multi-head self attention + SparseMoE) """
def __init__(self, n_embed, n_head, num_experts, top_k):
super().__init__()
self.sa = nn.MultiheadAttention(n_embed, n_head) # 我直接调用官网的attention方法
self.smoe = SparseMoE(n_embed, num_experts, top_k)
self.ln1 = nn.LayerNorm(n_embed)
self.ln2 = nn.LayerNorm(n_embed)
def forward(self, x):
qkv = self.ln1(x)
x = x + self.sa(qkv,qkv,qkv)[0] # 并不包含FFN结构
x = x + self.smoe(self.ln2(x))
return x
if __name__ == '__main__':
# 假设分词64个embed与4个头
n_embed, n_head = 64, 4
dropout = 0.1
# 假设4个专家与2个top_k
num_experts, top_k = 4, 2
# 假设数据是2个batch、6个分词,后面代码注释都已这些假设具体化
inputs = torch.rand(2, 6, n_embed)
M = Block(n_embed, n_head, num_experts, top_k)
y = M(inputs)
print(y.shape)