这是用户在 2024-6-16 23:01 为 https://zhuanlan.zhihu.com/p/701777558 保存的双语快照页面,由 沉浸式翻译 提供双语支持。了解如何保存?
从零实现一个MOE(专家混合模型)

从零实现一个MOE(专家混合模型)

ycz 等 121 人赞同了该文章
发布于 2024-06-05 16:13・IP 属地江苏
目录
收起
什么是混合模型(MOE)
从零实现一个MOE
1. 创建一个专家模型
2. 创建TopKrouter
3. 添加noisy噪声
4. 构建完整的稀疏MOE module
5. 将MOE与transformer结合
总结

什么是混合模型(MOE)

MOE主要由两个关键点组成:

一是将传统Transformer中的FFN(前馈网络层)替换为多个稀疏的专家层(Sparse MoE layers)。每个专家本身是一个独立的神经网络,实际应用中,这些专家通常是前馈网络 (FFN),但也可以是更复杂的网络结构。

二是门控网络或路由:此部分用来决定输入的token分发给哪一个专家。

可能有对FFN(前馈网络层)不太熟悉的小伙伴可以看一下下面的代码及图例,很简单就是一个我们平时常见的结构。

class FeedForward(nn.Module):
    def __init__(self, dim_vector, dim_hidden, dropout=0.1):
        super().__init__()
        self.feedforward = nn.Sequential(
            nn.Linear(dim_vector, dim_hidden),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(dim_hidden, dim_vector)
        )
        
    def forward(self, x):
        out = self.feedforward(x)
        return out

从零实现一个MOE

完整的从零实现MOE代码已集成至git代码训练框架项目,项目包括一个每个人都可以以此为基础构建自己的开源大模型训练框架流程、支持主流模型使用deepspeed进行Lora、Qlora等训练、主流模型的chat template模版、以及一些tricks的从零实现模块。欢迎大家star 共同学习!:

1. 创建一个专家模型

这一步也很简单了,其实就是一个多层感知机MLP。

class Expert(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        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)

2. 创建TopKrouter

即创建MOE的路由部分。

假设我们定义了4个专家,路由取前2名专家,即expert=4, top_k=2。接收注意力层的输出作为输入X,即将输入从(Batch size,Tokens,n_embed)的形状(2,4,32)投影到对应于(Batch size,Tokens,num_experts)的形状(2,4,4),其中num_experts即expert=4。其中返回的indices可以理解为对于每个token的4个专家来说,选的两个专家的序号索引。

代码如下:

# 这里我们假设定义n_embed为32, num_experts=4, top_k=2

class TopkRouter(nn.Module):
    def __init__(self, n_embed, num_experts, top_k):
        super(TopkRouter, self).__init__()
        self.top_k = top_k
        self.linear =nn.Linear(n_embed, num_experts)
    
    def forward(self, mh_output):
        logits = self.linear(mh_output)    # (2,4,32) ---> (2,4,4)
        # 获取前K大的值和索引,沿列。
        top_k_logits, indices = logits.topk(self.top_k, dim=-1)
        # 创建一个形状和logits相同全'-inf'矩阵,即(2,4,4)
        zeros = torch.full_like(logits, float('-inf'))
        # 按照索引和值填充上述zeros矩阵
        sparse_logits = zeros.scatter(-1, indices, top_k_logits)
        # 对其进行softmax,未被填充的位置会为0
        router_output = F.softmax(sparse_logits, dim=-1)
        return router_output, indices

看完代码之后配合整体流程图将会更清晰:

更清晰的图示如下,每个字代表一个token:

3. 添加noisy噪声

从本质上讲,我们不希望所有token都发送给同一组“受青睐”的expert。需要一个良好平衡,因此,将标准正态噪声添加到来自门控线性层的logits。

代码对比2中的代码只改动了几行,非常的简单。

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)
        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

4. 构建完整的稀疏MOE module

前面的操作主要是获取了router分发的结果,获取到这些结果后我们就可以将router乘给对应的token。这种选择性加权乘法最终构成了稀疏MOE。

代码部分如下所示:

class SparseMoE(nn.Module):
    def __init__(self, n_embed, num_experts, top_k):