一文带你详细了解：大模型MoE架构（含DeepSeek MoE详解）

MoE简史

混合专家模型（MoE）的理念最早源自 1991 年的论文《Adaptive Mixture of Local Experts》。它类似于集成学习方法，旨在通过一个监管机制管理由多个独立网络（称为“专家”）组成的系统。在这种系统中，每个专家处理训练数据的不同部分，专注于输入空间的特定区域。

那么，如何决定哪个专家来处理特定输入呢？这由门控网络负责，它根据输入分配权重给不同的专家。在训练过程中，专家和门控网络都会被同时优化，以提高整体决策能力。

在 2010 至 2015 年间，两个研究方向推动了 MoE 的发展：

• 组件专家 ：传统 MoE 由一个门控网络和多个专家组成，而一些研究探索了将 MoE 作为更深层神经网络的一部分。例如，《Learning Factored Representations in a Deep Mixture of Experts》提出了将 MoE 嵌入到多层网络中的某一层，使模型既庞大又高效。

  
  

• 条件计算 ：传统神经网络会处理所有输入数据，而这一时期的研究（如 Yoshua Bengio 等人的工作）开始探索根据输入动态激活或停用部分网络组件的方法。

  
  

这些研究促进了 MoE 在自然语言处理（NLP）中的应用。2017 年，Shazeer 等人将 MoE 用于一个 137B 规模的 LSTM（当时广泛用于 NLP 任务），通过引入稀疏性，实现了高效推理。然而，该方法在翻译任务中仍面临高通信成本和训练不稳定性等挑战。

混合专家模型 (MoE) 的引入使得训练具有数千亿甚至万亿参数的模型成为可能，如开源的 1.6 万亿参数的 Switch Transformers 等。这种技术不仅在自然语言处理 (NLP) 领域得到了广泛应用，也开始在计算机视觉领域进行探索。

MoE架构模型

在大模型时代，模型规模是提升性能的关键因素之一。在有限的计算资源下，与其用更多训练步数优化小模型，不如用更少的步数训练一个更大的模型，效果通常更佳。

混合专家模型（MoE）的优势在于，它能在远低于稠密模型的计算成本下进行高效预训练。这意味着，在相同预算下，可以训练更大规模的模型或数据集。特别是在预训练阶段，MoE 模型通常比稠密模型更快达到相同的质量水平。

什么是混合专家模型（MoE）？

MoE 是基于 Transformer 架构的模型，主要由两个核心部分组成：

• 稀疏 MoE 层 ：取代传统 Transformer 的前馈网络（FFN）层。MoE 层由多个“专家”（如 8 个）组成，每个专家是一个独立的神经网络，通常是 FFN，也可以是更复杂的结构，甚至是嵌套的 MoE 形成层级式结构。

  
  

• 门控网络或路由 ：决定哪些 Token 由哪个专家处理。例如，“More”可能被分配给第二个专家，而“Parameters”可能被分配给第一个。有时，一个 Token 甚至可以被多个专家处理。路由方式由可学习的参数控制，并与整个模型一同训练，是 MoE 关键机制之一。

上图来自 Switch Transformers 论文，展示了 MoE 层的结构。简单来说，MoE 通过将 Transformer 中的前馈网络（FFN）层替换为由门控网络和多个专家组成的 MoE 层。

MoE 具有高效预训练和比稠密模型更快的推理速度，但也面临一些挑战：

• 训练挑战
  ：MoE 在预训练时效率高，但在微调阶段容易泛化能力不足，导致过拟合。
• 推理挑战
  ：虽然 MoE 只激活部分参数进行推理，速度快于同等规模的稠密模型，但所有参数仍需加载到内存，导致较高的 VRAM 需求。例如，Mixtral 8×7B MoE 需要相当于 47B 参数稠密模型的内存，而非 8×7B=56B，因为 FFN 层是独立专家，其余参数共享。

此外，若每个 Token 仅使用两个专家，则推理 FLOPs 约等于 12B 模型，而非 14B，因为尽管进行了 2×7B 的矩阵运算，但部分层是共享的。

什么是稀疏性?

稀疏性的概念采用了条件计算的思想。在传统的稠密模型中，所有的参数都会对所有输入数据进行处理。相比之下， 稀疏性允许我们仅针对整个系统的某些特定部分执行计算 。这意味着并非所有参数都会在处理每个输入时被激活或使用，而是根据输入的特定特征或需求，只有部分参数集合被调用和运行。

让我们深入分析 Shazeer 对混合专家模型 (MoE) 在翻译应用中的贡献。条件计算的概念 (即仅在每个样本的基础上激活网络的不同部分) 使得在不增加额外计算负担的情况下扩展模型规模成为可能。这一策略在每个 MoE 层中实现了数以千计甚至更多的专家的有效利用。

这种稀疏性设置确实带来了一些挑战。例如，在混合专家模型 (MoE) 中，尽管较大的批量大小通常有利于提高性能，但当数据通过激活的专家时，实际的批量大小可能会减少。比如，假设我们的输入批量包含 10 个Token， 可能会有5个Token被路由到同一个专家，而剩下的5个Token分别被路由到不同的专家。这导致了批量大小的不均匀分配和资源利用效率不高的问题 。在接下来的部分中，将会介绍让MoE高效运行的其他挑战以及相应的解决方案。

那我们应该如何解决这个问题呢？一个可学习的门控网络 (G) 决定将输入的哪一部分发送给哪些专家 (E):

在这种设置下，虽然所有专家都会对所有输入进行运算，但通过门控网络的输出进行加权乘法操作。但是，如果 G (门控网络的输出) 为 0 会发生什么呢？如果是这种情况，就没有必要计算相应的专家操作，因此我们可以节省计算资源。那么一个典型的门控函数是什么呢？一个典型的门控函数通常是一个带有 softmax 函数的简单的网络。这个网络将学习将输入发送给哪个专家。

Shazeer 等人的工作还探索了其他的门控机制，其中包括带噪声的 Top-K 门控 (Noisy Top-K Gating)。这种门控方法引入了一些可调整的噪声，然后保留前 k 个值。具体来说:

1. 添加一些噪声
2. 选择保留前 K 个值
3. 应用 Softmax 函数

这种稀疏性引入了一些有趣的特性。通过使用较低的 k 值 (例如 1 或 2)，我们可以比激活多个专家时更快地进行训练和推理。为什么不仅选择最顶尖的专家呢？最初的假设是，需要将输入路由到不止一个专家，以便门控学会如何进行有效的路由选择，因此至少需要选择两个专家。

我们为什么要添加噪声呢？这是为了专家间的负载均衡！

MoE中的Token负载均衡

正如之前讨论的，如果所有的Token都被发送到只有少数几个受欢迎的专家，那么训练效率将会降低。在通常的混合专家模型 (MoE) 训练中，门控网络往往倾向于主要激活相同的几个专家。这种情况可能会自我加强，因为受欢迎的专家训练得更快，因此它们更容易被选择。

为了缓解这个问题，引入了一个 辅助损失 ，旨在鼓励给予所有专家相同的重要性。 这个损失确保所有专家接收到大致相等数量的训练样本，从而平衡了专家之间的选择 。接下来的部分还将探讨专家容量的概念，它引入了一个关于专家可以处理多少Token的阈值。在 transformers 库中，可以通过 aux_loss 参数来控制辅助损失。

MoE扩展Transformers

Transformer 类模型明确表明，增加参数数量可以提高性能，因此谷歌使用 GShard 尝试将 Transformer 模型的参数量扩展到超过 6000 亿并不令人惊讶。

GShard 将在编码器和解码器中的每个前馈网络 (FFN) 层中的替换为使用 Top-2 门控的混合专家模型 (MoE) 层，GShard 的工作对适用于 MoE 的并行计算模式也做出了重要贡献。下图展示了编码器部分的结构。这种架构对于大规模计算非常有效: 当扩展到多个设备时，MoE 层在不同设备间共享，而其他所有层则在每个设备上复制。 为了保持负载平衡和训练效率，GShard 的作者除了引入了上一节中讨论的类似辅助损失外，还引入了一些关键变化:

• 随机路由 : 在 Top-2 设置中，始终选择排名最高的专家，但第二个专家是根据其权重比例随机选择的。
• 专家容量 : 我们可以设定一个阈值，定义一个专家能处理多少Token。如果两个专家的容量都达到上限，Token就会溢出，并通过残差连接传递到下一层，或在某些情况下被完全丢弃。

专家容量是 MoE 中最重要的概念之一 。为什么需要专家容量呢？因为所有张量的形状在编译时是静态确定的，无法提前知道多少Token会分配给每个专家，因此需要一个固定的容量因子。

注意 : 在推理过程中，只有部分专家被激活。同时，有些计算过程是共享的，例如自注意力 (self-attention) 机制，它适用于所有Token。这就解释了为什么可以使用相当于 12B 稠密模型的计算资源来运行一个包含 8 个专家的 47B 模型。如果我们采用 Top-2 门控，模型会使用高达 14B 的参数。但是，由于自注意力操作 (专家间共享) 的存在，实际上模型运行时使用的参数数量是 12B。

MoE训练与微调

Switch Transformers

尽管混合专家模型 (MoE) 显示出了很大的潜力，但它们在训练和微调过程中存在稳定性问题。Switch Transformers 是一项非常激动人心的工作，它深入研究了这些话题。作者甚至在 Hugging Face 上发布了一个 1.6 万亿参数的 MoE，拥有 2048 个专家，实现了与 T5-XXL 相比 4 倍的预训练速度提升。 就像在 GShard 中一样，作者用混合专家模型 (MoE) 层替换了前馈网络 (FFN) 层。Switch Transformers 提出了一个 Switch Transformer 层，它接收两个输入 (两个不同的Token) 并拥有四个专家。与最初使用至少两个专家的想法相反，Switch Transformers 采用了简化的单专家策略。这种方法的效果包括:

• 减少门控网络 (路由) 计算负担
• 每个专家的批量大小至少可以减半
• 降低通信成本
• 保持模型质量

Switch Transformers 也对 专家容量 这个概念进行了研究。

上述建议的容量是将批次中的Token数量均匀分配到各个专家。如果我们使用大于 1 的容量因子，我们为Token分配不完全平衡时提供了一个缓冲。增加容量因子会导致更高的设备间通信成本，因此这是一个需要考虑的权衡。特别值得注意的是，Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表现出色。

Switch Transformer 的作者还重新审视并简化了前面章节中提到的负载均衡损失。在训练期间，对于每个 Switch 层的辅助损失被添加到总模型损失中。这种损失鼓励均匀路由，并可以使用超参数进行加权。

作者还尝试了混合精度的方法，例如用 bfloat16 精度训练专家，同时对其余计算使用全精度进行。较低的精度可以减少处理器间的通信成本、计算成本以及存储张量的内存。然而，在最初的实验中，当专家和门控网络都使用 bfloat16 精度训练时，出现了不稳定的训练现象。这种不稳定性特别是由路由计算引起的，因为路由涉及指数函数等操作，这些操作对精度要求较高。因此，为了保持计算的稳定性和精确性，保持更高的精度是重要的。 为了减轻不稳定性，路由过程也使用了全精度 。 使用混合精度不会降低模型质量并可实现更快的训练

Switch Transformers 采用了编码器 - 解码器的架构，实现了与 T5 类似的混合专家模型 (MoE) 版本。GLaM 这篇工作探索了如何使用仅为原来 1/3 的计算资源 (因为 MoE 模型在训练时需要的计算量较少，从而能够显著降低碳足迹) 来训练与 GPT-3 质量相匹配的模型来提高这些模型的规模。作者专注于仅解码器 (decoder-only) 的模型以及少样本和单样本评估，而不是微调。他们使用了 Top-2 路由和更大的容量因子。此外，他们探讨了将容量因子作为一个动态度量，根据训练和评估期间所使用的计算量进行调整。

稳定模型训练

之前讨论的平衡损失可能会导致稳定性问题。我们可以使用许多方法来稳定稀疏模型的训练，但这可能会牺牲模型质量。例如，引入 dropout 可以提高稳定性，但会导致模型质量下降。另一方面，增加更多的乘法分量可以提高质量，但会降低模型稳定性。

ST-MoE 引入的 Router z-loss 在保持了模型性能的同时显著提升了训练的稳定性。这种损失机制通过惩罚门控网络输入的较大 logits 来起作用，目的是促使数值的绝对大小保持较小，这样可以有效减少计算中的舍入误差。这一点对于那些依赖指数函数进行计算的门控网络尤其重要。为了深入了解这一机制，建议参考原始论文以获得更全面的细节。

专家学习特点

ST-MoE 的研究者们发现，编码器中不同的专家倾向于专注于特定类型的Token或浅层概念。例如，某些专家可能专门处理标点符号，而其他专家则专注于专有名词等。与此相反，解码器中的专家通常具有较低的专业化程度。

此外，研究者们还对这一模型进行了多语言训练。尽管人们可能会预期每个专家处理一种特定语言，但实际上并非如此。由于Token路由和负载均衡的机制，没有任何专家被特定配置以专门处理某一特定语言。

微调策略

1）增加更多专家可以提升处理样本的效率和加速模型的运算速度，但这些优势随着专家数量的增加而递减 (尤其是当专家数量达到 256 或 512 之后更为明显) 。同时，这也意味着在推理过程中，需要更多的显存来加载整个模型。值得注意的是，Switch Transformers 的研究表明，其在大规模模型中的特性在小规模模型下也同样适用，即便是每层仅包含 2、4 或 8 个专家。

2)稠密模型和稀疏模型在过拟合的动态表现上存在显著差异 。稀疏模型更易于出现过拟合现象，因此在处理这些模型时，尝试更强的内部正则化措施是有益的，比如使用更高比例的 dropout。例如，我们可以为稠密层设定一个较低的 dropout 率，而为稀疏层设置一个更高的 dropout 率，以此来优化模型性能。

3)在微调过程中是否使用辅助损失是一个需要决策的问题 。ST-MoE 的作者尝试关闭辅助损失，发现即使高达 11% 的Token被丢弃，模型的质量也没有显著受到影响。 Token丢弃可能是一种正则化形式，有助于防止过拟合 。

Switch Transformers 的作者观察到，在相同的预训练困惑度下，稀疏模型在下游任务中的表现不如对应的稠密模型，特别是在重理解任务 (如 SuperGLUE) 上。另一方面，对于知识密集型任务 (如 TriviaQA)，稀疏模型的表现异常出色。作者还观察到， 在微调过程中，较少的专家的数量有助于改善性能 。另一个关于泛化问题确认的发现是，模型在小型任务上表现较差，但在大型任务上表现良好。

在小任务 (左图) 中，我们可以看到明显的过拟合，因为稀疏模型在验证集中的表现要差得多。在较大的任务 (右图) 中，MoE 则表现良好。

4)一种可行的微调策略是尝试冻结所有非专家层的权重 。实践中，这会导致性能大幅下降，但这符合我们的预期，因为混合专家模型 (MoE) 层占据了网络的主要部分。我们可以尝试相反的方法: 仅冻结 MoE 层的参数。实验结果显示，这种方法几乎与更新所有参数的效果相当。这种做法可以加速微调过程，并降低显存需求 。 在微调稀疏混合专家模型 (MoE) 时需要考虑的最后一个问题是，它们有特别的微调超参数设置——例如， 稀疏模型往往更适合使用较小的批量大小和较高的学习率，这样可以获得更好的训练效果 。 最新MoE方法>> 专注大模型/AIGC、Agent、RAG等学术前沿分享！

MoE转折点

此时，您可能会对人们微调 MoE 中遇到的这些挑战而感到沮丧，但最近的一篇论文 “MoEs Meets Instruction Tuning” (2023 年 7 月) 带来了令人兴奋的发现。这篇论文进行了以下实验:

• 单任务微调
• 多任务指令微调
• 多任务指令微调后接单任务微调

当研究者们对 MoE 和对应性能相当的 T5 模型进行微调时，他们发现 T5 的对应模型表现更为出色。然而，当研究者们对 Flan T5 (一种 T5 的指令优化版本) 的 MoE 版本进行微调时，MoE 的性能显著提升。更值得注意的是，Flan-MoE 相比原始 MoE 的性能提升幅度超过了 Flan T5 相对于原始 T5 的提升，这意味着 MoE 模型可能从指令式微调中获益更多，甚至超过了稠密模型 。此外，MoE 在多任务学习中表现更佳。与之前关闭 辅助损失 函数的做法相反，实际上这种损失函数可以帮助防止过拟合。

Deepseek MOE 架构

和基础 MOE 结构的区别是：

1. 更精细地划分专家网络 ，提升每个专家的专业性，提高知识表达的准确度。
2. 引入部分共享专家 ，减少不同专家间的知识冗余，提升计算效率；所有 tokens 都会经过的共享专家，每个 token 会用计算的 Router 权重，来选择 topK 个专家，然后和共享的专家的输出一起加权求和。

DeepseekMOE 其实是有两类专家的：

• 共享专家（Shared Expert）：1 个共享专家，用于捕捉 通用 、全局的特征信息。
• 路由专家（Routed Experts）：每个 MoE 层都包含 256 个路由专家，负责精细化处理输入 tokens 的专业特征。

Gate 网络与 DeepseekMOE 计算流程

当一个 token 的向量传入 MoE 层时，首先会经过一个专门的 Gate 网络，该网络负责计算 token 与各个路由专家之间的匹配得分。具体流程如下：

1. 计算 tokens 和专家的匹配得分

• Gate 网络通过 线性变换 计算每个 token 与所有路由专家的兼容性得分。得分可以反映 token 和各专家“契合”的程度。

• 基于得分，Gate 网络为每个 token 选择 Top-K 个最合适的路由专家。在 DeepSeek‐V3 中，每个 token 通常选择 8 个路由专家（在一些实现中还可能对跨节点路由做限制，如最多路由到 4 个不同节点），从而只激活极少数专家进行计算。

• 被选中的专家各自对 token 进行独立处理（专家实际上是小型 FFN 模块，类似于 Transformer 中的 FFN 模块），并产生各自的输出。然后， 这些专家的输出会根据 Gate 网络给出的得分权重进行加权聚合 ，最后再和共享专家的输出进行融合，形成当前 MoE 层的最终输出表示。

DeepseekV2 模型的 MOE 参数如下：

{
  // 部分参数省略
"hidden_act": "silu",
"hidden_size": 5120,
"initializer_range": 0.02,
"intermediate_size": 12288,
"model_type": "deepseek_v2",
"moe_intermediate_size": 1536,
"moe_layer_freq": 1,
"n_group": 8,
"n_routed_experts": 160,
"n_shared_experts": 2,
"norm_topk_prob": false,
"num_experts_per_tok": 6,
"num_hidden_layers": 60,
"num_key_value_heads": 128,
"topk_group": 3,
"topk_method": "group_limited_greedy",
}

混合专家（MoE）参数说明：

Deepseek MOE结构代码

这里只考虑推理模式下的 DeepseekMOE 结构实现，且分步实现。

DeepseekV2MLP 实现

专家其实就是参数量更少的 FFN/MLP 结构，和 llama 中结构一样，只是参数量和计算量更少了，DeepseekV2MLP 代码如下所示。

classDeepseekV2MLP(nn.Module):
    def__init__(self, config, hidden_size=None, intermediate_size=None):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size if hidden_size isNoneelse hidden_size
        self.intermediate_size = (
            config.intermediate_size if intermediate_size isNoneelse intermediate_size
        )

        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act] # silu 激活函数

    defforward(self, x):
        mlp_out = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return mlp_out

门控/路由网络实现

门控网络的作用是，根据输入 tokens 动态的选择 Top-K 个专家，并为每个 Token 分配权重。关键流程如下：

1. 门控分数计算 ：通过线性层 + Softmax 生成专家选择概率分布。
2. Top-K 专家选择 ：支持两种模式（贪婪选择 vs 分组限制贪婪选择），贪婪模式直接使用 torch.topk 函数 选取分数张量中的前 k 个分数 。
3. 权重归一化 ：对 Top-K 权重进行归一化或缩放。

代码实现如下所示:

import torch
import torch.nn as nn
import torch.nn.functional as F
import math 
from dataclasses import dataclass

classMoEGate(nn.Module):
    def__init__(self, config):
        super().__init__()
        self.config = config
        self.top_k = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts
        self.routed_scaling_factor = config.routed_scaling_factor
        self.scoring_func = config.scoring_func
        self.topk_method = config.topk_method
        self.n_group = config.n_group
        self.topk_group = config.topk_group
        self.norm_topk_prob = config.norm_topk_prob
        
        # 静态化推理配置（假设配置固定）
        self.inference_norm = self.norm_topk_prob and (self.top_k > 1)
        self.use_group_limited = (self.topk_method == "group_limited_greedy")

        # 门控权重
        self.gating_dim = config.hidden_size
        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
        self.reset_parameters()

    defreset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    @torch.inference_mode()  # 禁用梯度与训练逻辑





    
    defforward(self, hidden_states):
        bsz, seq_len, h = hidden_states.shape
        hidden_states = hidden_states.reshape(-1, h)
        
        # 门控分数计算（保持原始数据类型）
        logits = F.linear(hidden_states, self.weight)  # [n_tokens, n_experts]
        scores = logits.softmax(dim=-1)  # 自动推断 dtype

        # Top-K 选择（静态分支）
        if self.use_group_limited:
            # 分组限制逻辑优化
            group_scores = scores.view(bsz * seq_len, self.n_group, -1).max(dim=-1).values
            group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
            group_mask = torch.zeros_like(group_scores).scatter_(1, group_idx, 1)
            score_mask = group_mask.unsqueeze(-1).expand(-1, -1, self.n_routed_experts // self.n_group).reshape(bsz * seq_len, -1)
            scores = scores.masked_fill(~score_mask.bool(), 0.0)
        
        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

        # 权重归一化（静态分支）
        if self.inference_norm:
            topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
        else:
            topk_weight = topk_weight * self.routed_scaling_factor

        return topk_idx, topk_weight, None# aux_loss 始终为 None

@dataclass
classDeepseekV2Config:
    # 1, Position Config
    max_position_embeddings: int = 163840
    vocab_size: int = 102400

    # 2, MLA Config
    # down_linear config
    q_lora_rank: int = 1536
    kv_lora_rank: int = 512

    # head_dim、heads and hidden_size config
    v_head_dim: int = 128
    qk_nope_head_dim: int = 128
    qk_rope_head_dim: int = 64
    hidden_size: int = 5120
    num_attention_heads: int = 128
    num_key_value_heads: int = 128
    
    attention_bias: bool = False

    attention_dropout: float = 0.1
    # rope config
    rope_theta: float = 10000

    # 3, MOE Config
    n_group: int = 8
    n_routed_experts: int = 160
    num_experts_per_tok: int = 6
    topk_group: int = 3
    routed_scaling_factor: float = 1.0
    scoring_func: str="softmax"
    topk_method: str="greedy"
    norm_topk_prob: bool = True

# 初始化配置
config = DeepseekV2Config()

# 模拟输入，CPU 电脑可直接跑，去除了 cuda 设备限制代码
device = torch.device("cuda"if torch.cuda.is_available() else"cpu")
hidden_states = torch.randn(32, 64, 5120, device=device)

# 创建模块
moe_gate = MoEGate(config)  # 半精度推理

# gate 网络推理
topk_idx, topk_weight, _ = moe_gate(hidden_states)

print("topk_idx shape ", topk_idx.shape) # 32 * 64 = 2048 个 tokens
print("topk_weight shape", topk_weight.shape)

"""
# 输出如下，表示每个 token 会激活 6 个专家参与计算
topk_idx shape  torch.Size([2048, 6]) 
topk_weight shape torch.Size([2048, 6])
"""

DeepseekMOE 实现

1. 门控计算

• 调用门控网络（self.gate），对输入 hidden_states 计算得到 top‑k 专家索引（topk_idx）、对应权重（topk_weight）以及辅助损失（aux_loss，推理时不参与梯度计算）。

• 将输入 hidden_states 展平为二维张量（形状 [B * T, d]），并将 topk_idx 也展平。
• 在推理模式下，通常不需要像训练时那样对每个 token 进行 repeat_interleave，因为每个 token 只会由对应专家处理一次。

• 根据展平后的 topk_idx ，依次对每个专家负责的 token 子集进行计算。
• 由于这里可能存在多个 token 被分配给不同专家，实际实现中需要将每个专家的输出按顺序记录下来。

• 将所有专家计算的输出进行合并。通过将输出重新整理（排序）回原始 token 顺序，并按照 topk_weight 对各个专家输出进行加权求和，从而获得最终输出。
• 整个过程保证最终输出形状与原始输入保持一致，即 [B, T, d]。

代码实现如下所示：

# 为了单元测试，模拟不使用分布式（ep_size默认为1）
classDeepseekV2MoE(nn.Module):
    """
    A mixed expert module containing shared experts.
    """
    def__init__(self, config):
        super().__init__()
        self.config = config
        self.num_experts_per_tok = config.num_experts_per_tok

        self.experts = nn.ModuleList(
            [
                DeepseekV2MLP(
                    config, intermediate_size=config.moe_intermediate_size
                )
                for i in range(config.n_routed_experts)
            ]
        )
        self.gate = MoEGate(config)
        if config.n_shared_experts isnotNone:
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
            self.shared_experts = DummyMLP(config=config, intermediate_size=intermediate_size)

    # 此处为简化实现，仅做推理示例，不涉及分布式通信
    @torch.no_grad()
    defmoe_infer(self, x, topk_ids, topk_weight):
        # x: [batch * seq_len, hidden_size]
        # 对每个 token 依然采用与训练类似的方式进行专家计算
        outputs = []
        flat_topk_ids = topk_ids.view(-1)
        for i, expert in enumerate(self.experts):
            mask = (flat_topk_ids == i)
            if mask.sum() == 0:
                continue
            outputs.append(expert(x[mask]))
        # 简单拼接，不做复杂排序和 all-to-all 操作
        outs = torch.cat(outputs, dim=0)
        new_x = torch.empty_like(outs)
        # 这里直接返回加权求和的结果（实际实现更复杂）
        final_out = (outs.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
        return final_out

稀疏 VS 稠密，如何选择?

稀疏混合专家模型 (MoE) 适用于拥有多台机器且要求高吞吐量的场景。在固定的预训练计算资源下，稀疏模型往往能够实现更优的效果。相反，在显存较少且吞吐量要求不高的场景，稠密模型则是更合适的选择。

注意 : 直接比较稀疏模型和稠密模型的参数数量是不恰当的，因为这两类模型基于的概念和参数量的计算方法完全不同。

MoE计算效率提升

最初的混合专家模型 (MoE) 设计采用了分支结构，这导致了计算效率低下。这种低效主要是因为 GPU 并不是为处理这种结构而设计的，而且由于设备间需要传递数据，网络带宽常常成为性能瓶颈。在接下来的讨论中，我们会讨论一些现有的研究成果，旨在使这些模型在预训练和推理阶段更加高效和实用。我们来看看如何优化 MoE 模型，让 MoE 起飞。

并行计算

让我们简要回顾一下并行计算的几种形式:

• 数据并行 : 相同的权重在所有节点上复制，数据在节点之间分割。
• 模型并行 : 模型在节点之间分割，相同的数据在所有节点上复制。
• 模型和数据并行 : 我们可以在节点之间同时分割模型和数据。注意，不同的节点处理不同批次的数据。
• 专家并行 : 专家被放置在不同的节点上。如果与数据并行结合，每个节点拥有不同的专家，数据在所有节点之间分割。

在专家并行中，专家被放置在不同的节点上，每个节点处理不同批次的训练样本。对于非 MoE 层，专家并行的行为与数据并行相同。对于 MoE 层，序列中的Token被发送到拥有所需专家的节点。Switch Transformers 论文中展示如何使用不同的并行技术在节点上分割数据和模型的插图，如下所示。 最新MoE方法>> 专注大模型/AIGC、Agent、RAG等学术前沿分享！

容量因子和通信开销

提高容量因子 (Capacity Factor, CF) 可以增强模型的性能，但这也意味着更高的通信成本和对保存激活值的显存的需求。在设备通信带宽有限的情况下，选择较小的容量因子可能是更佳的策略。

一个合理的初始设置是采用 Top-2 路由、1.25 的容量因子，同时每个节点配置一个专家 。在评估性能时，应根据需要调整容量因子，以在设备间的通信成本和计算成本之间找到一个平衡点。

部署技术

部署混合专家模型 (MoE) 的一个关键挑战是其庞大的参数规模。对于本地使用情况，我们可能希望使用更小的模型。为了使模型更适合部署，下面是几种有用的技术:

• 预先蒸馏实验 : Switch Transformers 的研究者们进行了预先蒸馏的实验。他们通过将 MoE 模型蒸馏回其对应的稠密模型，成功保留了 30-40%的由稀疏性带来的性能提升。预先蒸馏不仅加快了预训练速度，还使得在推理中使用更小型的模型成为可能。
• 任务级别路由 : 最新的方法中，路由器被修改为将整个句子或任务直接路由到一个专家。这样做可以提取出一个用于服务的子网络，有助于简化模型的结构。
• 专家网络聚合 : 这项技术通过合并各个专家的权重，在推理时减少了所需的参数数量。这样可以在不显著牺牲性能的情况下降低模型的复杂度。

高效训练

FasterMoE 深入分析了 MoE 在不同并行策略下的理论性能极限，并且探索了一系列创新技术 ，包括用于专家权重调整的方法、减少延迟的细粒度通信调度技术，以及一个基于最低延迟进行专家选择的拓扑感知门控机制。这些技术的结合使得 MoE 运行速度提升高达 17 倍。

Megablocks 则专注于通过开发新的 GPU kernel 来处理 MoE 模型中的动态性，以实现更高效的稀疏预训练 。其核心优势在于，它不会丢弃任何Token，并能高效地适应现代硬件架构 (支持块稀疏矩阵乘)，从而达到显著的加速效果。Megablocks 的创新之处在于，它不像传统 MoE 那样使用批量矩阵乘法 (这通常假设所有专家形状相同且处理相同数量的Token)，而是将 MoE 层表示为块稀疏操作，可以灵活适应不均衡的Token分配。下图为不同规模的专家和Token数量的块稀疏矩阵乘法。

开源混合专家模型

目前，Megablocks、Fairseq、OpenMoE等开源项目可以用于训练混合专家模型 (MoE)；对于开源的混合专家模型 (MoE)，你可以关注下面这些：Switch Transformers (Google)、NLLB MoE (Meta)、OpenMoE、Mixtral 8x7B (Mistral)、DeepSeek-R1(DeepSeek)。

研究方向

MoE 的 量化 也是一个有趣的研究领域。例如，QMoE通过将 MoE 量化到每个参数不到 1 位，将 1.6 万亿参数的 Switch Transformer 所需的存储从 3.2TB 压缩到仅 160GB。简而言之，一些值得探索的有趣领域包括:

• 将 Mixtral 蒸馏成一个稠密模型。
• 探索合并专家模型的技术及其对推理时间的影响。
• 尝试对 Mixtral 进行极端量化的实验。