MoE专家模块Demo

最新推荐文章于 2024-08-11 17:51:50 发布
最新推荐文章于 2024-08-11 17:51:50 发布
阅读量1k 收藏 27
点赞数 9
文章标签： MOE 混合专家模块
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。
本文链接：https://blog.csdn.net/weixin_38252409/article/details/140821366
版权

文章目录

前言

随着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)
tangjunjun-owen
关注
模块化MoE：视觉多任务学习的基础模型
DevWizard的博客
09-19 335
模块化MoE的核心思想是引入一个门控网络，用于动态地选择不同专家模块的输出。通过最小化整体模型输出与标签之间的差异，以及专家模块输出与任务标签之间的差异，来学习模型参数和门控网络的权重。门控网络根据输入数据动态地计算每个专家模块的权重，然后将权重与专家模块的输出进行加权组合。最终，得到整体模型的输出。模块化MoE是一种用于视觉多任务学习的基础模型，它通过将多个专家模块和门控网络组合在一起，实现对多个任务的综合处理。通过模块化MoE，我们可以在视觉多任务学习中取得更好的性能，同时提高模型的灵活性和可扩展性。
开源模型应用落地-Qwen1.5-MoE-1/3的激活参数量达到7B模型的性能
没有卑微的工作，只有卑微的心态，与其抱怨，不如埋头实干
05-20 6624
使用Qwen1.5-MoE-A2.7B，体验更快的推理速度
Fedformer-MOE模块
To be a better man
07-12 695
Fedformer-MOE模块
150亿参数，谷歌开源了史上最大视觉模型V-MoE的全部代码
机器学习社区
01-14 667
还记得谷歌大脑团队去年 6 月份发布的 43 页论文《Scaling Vision with Sparse Mixture of Experts》吗？他们推出了史上最大规模的视觉模型 V-MoE，实现了接近 SOTA 的 Top-1 准确率。如今，谷歌大脑开源了训练和微调模型的全部代码。在过去几十年里，深度学习的进步是由几个关键因素推动的：少量简单而灵活的机制、大型数据集、更专业的硬件配置，这些技术的进步使得神经网络在图像分类、机器翻译、蛋白质预测等任务中
Mixtral Moe代码解读
Swift's Blog
01-11 830
title: Mixtral Moe代码解读toc: true一直对稀疏专家网络好奇，有些专家没被选中，那么梯度是否为0，这一轮被选中有梯度，下一轮没被选中无梯度，模型可以训练收敛吗？more。
MoE结合TopK-router简单代码示例
To be a better man
06-29 511
请注意，这个示例是为了说明 MoE 模型的基本结构和工作原理，实际应用中可能需要更复杂的网络结构和训练策略。在这个例子中，我们使用 PyTorch 框架来定义模型。下面是一个简化的 MoE 模型的代码示例，其中包括一个。
Moe:一个简单的通用模板
06-08
项目说明 Includes:
美国加州大学圣克鲁兹分校王鑫教授招收NLP/CV方向全奖博士生
深度学习技术前沿
01-07 724
来源：AI求职加州大学圣克鲁兹分校UC Santa Cruzis one of America's Public Ivy universities and a member of the...
混合专家模型（MoE）入门
最新发布
weixin_44858814的博客
08-11 1447
模型规模是提升LLM大语言模型性能的关键因素，但也会增加计算成本。Mixture of Experts (MoE) 架构通过分布式专家层和动态门控机制，有效降低了计算资源，使模型能够在扩展参数规模的同时保持高效的运行。
【阅读笔记】多任务学习之MMoE(含代码实现)
简简单单的学习笔记，致力于帮助更多前进路上的朋友~
11-01 9353
本文作为自己阅读论文后的总结和思考，不涉及论文翻译和模型解读，适合大家阅读完论文后交流想法，文末含MMoE的Pytorch实现及Synthetic Data的生成代码。
jsteward.moe:jsteward.moe 的来源
07-24
jsteward.moe 的来源 住在这里： : 记得在本地测试之前获取主题： git submodule update --init --recursive
不同的图像patch由不同的专家模型来处理！南洋理工&Mila稀疏融合混合专家模型SF-MoE，具有超强泛化能力！代码已开源！...
我爱计算机视觉
07-15 677
关注公众号，发现CV技术之美本篇分享论文『Sparse Fusion Mixture-of-Experts are Domain Generalizable Learners』，不同的图像patch由不同的专家模型来处理！南洋理工&Mila稀疏融合混合专家模型SF-MoE，具有超强泛化能力！代码已开源！详细信息如下：论文地址：https://arxiv.org/a...
MOE原理解释及从零实现一个MOE（专家混合模型）
weixin_43013480的博客
06-05 6521
一个MOE主要由两个关键点一是将传统Transformer中的FFN（前馈网络层）替换为多个稀疏的专家层（Sparse MoE layers）。每个专家本身是一个独立的神经网络，实际应用中，这些专家通常是前馈网络 (FFN)，但也可以是更复杂的网络结构。二是门控网络或路由：此部分用来决定输入的token分发给哪一个专家。可能有对FFN（前馈网络层）不太熟悉的小伙伴可以看一下下面的代码及图例，很简单就是一个我们平时常见的结构。
使用PyTorch实现混合专家(MoE)模型
deephub
01-10 7562
我们测试的输出结果显示了混合专家模型的强大功能。该模型通过门控网络将各个专家模型的优势结合起来，取得了比单个专家模型更高的精度。门控网络有效地学习了如何根据输入数据权衡每个专家的贡献，以产生更准确的预测。混合专家利用了各个模型的不同专业知识，在测试数据集上提供了更好的性能。同时也说明我们可以在现有的任务上尝试使用MOE来进行测试，也可以得到更好的结果。作者：Shahriar Hossain。
一步一步理解大模型：混合专家模型（MoE）简介
热门推荐
chattyfish的博客
07-25 3万+
混合专家模型，英文叫Mixture of Experts (MoE) 是一种模型设计策略，它通过将多个模型（称为"专家"）的预测结果结合在一起，以获得更好的预测性能。
稀疏混合专家架构语言模型（MoE）~自实现
whaosoft143ai的博客
02-11 765
门控网络，也称为路由，确定哪个专家网络接收来自多头注意力的 token 的输出。该层将对应于（Batch size，Tokens，n_embed）的输入张量从（2，4，32）维度，投影到对应于（Batch size、Tokens，num_expert）的新形状：（2、4，4）。本文介绍了实现一个稀疏混合专家语言模型（MoE）的方法，详细解释了模型的实施过程，包括采用稀疏混合专家取代传统的前馈神经网络，实现 top-k 门控和带噪声的 top-k 门控，以及采用 Kaiming He 初始化技术。whao
Mixtral MOE代码理解
REfusing的博客
03-17 1199
我在看MOE的时候，虽然大概能够理解MOE的模型结构，但是看一些作者实现的代码（应该不是官方代码），虽然写的很好，但是始终理解无法彻底理解他代码的意思，于是，简单运行了一下，特此记录一下。
从Mistral 7B到MoE模型Mixtral 8x7B的全面解析：从原理分析到代码解读
结构之法 算法之道
12-23 1万+
23年12月9日，Mistral AI 在 X 平台甩出一条磁力链接，打开一看，发现是接近 87 GB 的种子看上去，Mistral 8x7B的架构此前传闻的GPT-4架构非常相似(很像传闻中GPT-4的同款方案)，但是「缩小版」：8 个专家总数，而不是 16 名(减少一半)每个专家为 7B 参数，而不是 166B(减少 24 倍)42B 总参数（估计）而不是 1.8T(减少 42 倍)与原始 GPT-4 相同的 32K 上下文。
MMOE核心代码块
suibianshen2012的专栏
03-28 721
class MMoE_Layer(tf.keras.layers.Layer): def __init__(self,expert_dim,n_expert,n_task): super(MMoE_Layer, self).__init__() self.n_task = n_task self.expert_layer = [Dense(expert_dim,activation = 'relu') for i in range(n_expert).
MOE 混合专家模型
04-10
MOE（Mixture of Experts）混合专家模型是一种机学习模型，用于解决复杂的非线性问题。通过将多个家模型组合起来，以提高整体模的性能和化能力。 在MOE模型中，每个专家模型都是一个独立的子模型，负责处理特定的输入数据子空间。每个专家模型都有自己的参数和权重，用于对输入数据进行建模和预测。而整个MOE模型则通过一个门控网络（Gate Network）来决定每个专家模型对于给定输入数据的权重。 门控网络接收输入数据，并输出一个权重向量，用于控制每个专家模型的贡献程度。这样，在给定输入数据时，MOE模型会根据门控网络的输出，将不同的专家模型的预测结果进行加权求和，得到最终的预测结果。 MOE模型的优势在于能够充分利用不同专家模型的特长，从而提高整体模型的性能。它可以应用于各种领域，如自然语言处理、计算机视觉等。
评论
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

当前余额3.43元 前往充值 >
需支付：10.00元
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

tangjunjun-owen

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付：¥1
获取中
扫码支付

您的余额不足，请更换扫码支付或充值

打赏作者

实付元
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明：

1.余额是钱包充值的虚拟货币，按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载，可以购买VIP、付费专栏及课程。

余额充值