最近春招和训练已开启了。
不同以往的是，当前职场环境已不再是那个双向奔赴时代了。求职者在变多，HC 在变少，岗位要求还更高了。
最近，我们又陆续整理了很多大厂的口试题，资助一些球友解惑答疑，分享技术口试中的那些弯弯绕绕。
喜好本文记得收藏、关注、点赞。
混合专家模子 (Mixed Expert Models，简称 MoEs) ，最早是随着 Mixtral 8x7B 的推出而逐渐引起人们的广泛关注。最近随着DeepSeek的爆火，MoE又进一步引起大家的关注。
我将带你了解 MoEs的发展史、核心组件、训练方法，推理中各因素考量和DeepSeek MoE详解。干系头脑导图如下所示：
如需文中头脑导图、论文等，复兴：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 个值。具体来说:
- 添加一些噪声
- 选择保留前 K 个值
- 应用 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转折点
此时，您可能会对人们微调 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 结构的区别是：
- 更精致地分别专家网络，提拔每个专家的专业性，提高知识表达的准确度。
- 引入部分共享专家，减少不同专家间的知识冗余，提拔计算服从；所有 tokens 都会颠末的共享专家，每个 token 会用计算的 Router 权重，来选择 topK 个专家，然后和共享的专家的输出一起加权求和。
DeepseekMOE 其实是有两类专家的：
- 共享专家（Shared Expert）：1 个共享专家，用于捕捉通用、全局的特征信息。
- 路由专家（Routed Experts）：每个 MoE 层都包罗 256 个路由专家，负责精致化处理输入 tokens 的专业特征。
Gate 网络与 DeepseekMOE 计算流程
当一个 token 的向量传入 MoE 层时，起首会颠末一个专门的 Gate 网络，该网络负责计算 token 与各个路由专家之间的匹配得分。具体流程如下：
- 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",
- }
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 分配权重。关键流程如下：
- 门控分数计算：通过线性层 + Softmax 生成专家选择概率分布。
- Top-K 专家选择：支持两种模式（贪心选择 vs 分组限定贪心选择），贪心模式直接使用 torch.topk 函数选取分数张量中的前 k 个分数。
- 权重归一化：对 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])
- """
- 调用门控网络（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
稀疏混合专家模子 (MoE) 实用于拥有多台呆板且要求高吞吐量的场景。在固定的预训练计算资源下，稀疏模子通常可以或许实现更优的效果。相反，在显存较少且吞吐量要求不高的场景，稠密模子则是更合适的选择。
注意: 直接比力稀疏模子和稠密模子的参数数量是不恰当的，因为这两类模子基于的概念和参数量的计算方法完全不同。
MoE计算服从提拔
最初的混合专家模子 (MoE) 设计采用了分支结构，这导致了计算服从低下。这种低效主要是因为 GPU 并不是为处理这种结构而设计的，而且由于设备间必要传递数据，网络带宽常常成为性能瓶颈。在接下来的讨论中，我们会讨论一些现有的研究结果，旨在使这些模子在预训练和推理阶段更加高效和实用。我们来看看如何优化 MoE 模子，让 MoE 腾飞。
并行计算
让我们扼要回首一下并行计算的几种形式:
- 数据并行: 相同的权重在所有节点上复制，数据在节点之间分割。
- 模子并行: 模子在节点之间分割，相同的数据在所有节点上复制。
- 模子和数据并行: 我们可以在节点之间同时分割模子和数据。注意，不同的节点处理不同批次的数据。
- 专家并行: 专家被放置在不同的节点上。假如与数据并行结合，每个节点拥有不同的专家，数据在所有节点之间分割。
在专家并行中，专家被放置在不同的节点上，每个节点处理不同批次的训练样本。对于非 MoE 层，专家并行的行为与数据并行相同。对于 MoE 层，序列中的Token被发送到拥有所需专家的节点。Switch Transformers 论文中展示如何使用不同的并行技术在节点上分割数据和模子的插图，如下所示。最新MoE方法>>
容量因子和通信开销
提高容量因子 (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 进行极端量化的实验。
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。 |
