混合专家模型 (MoE)笔记摘要

ref： https://huggingface.co/blog/zh/moe#%E4%BB%80%E4%B9%88%E6%98%AF%E6%B7%B7%E5%90%88%E4%B8%93%E5%AE%B6%E6%A8%A1%E5%9E%8B

简短总结
混合专家模型 (MoEs):

与稠密模型相比，预训练速度更快
与具有相同参数数量的模型相比，具有更快的推理速度
需要大量显存，因为所有专家系统都需要加载到内存中
在微调方面存在诸多挑战，但近期的研究表明，对混合专家模型进行指令调优具有很大的潜力。

混合专家模型 (MoE) 的一个显著优势是它们能够在远少于稠密模型所需的计算资源下进行有效的预训练。这意味着在相同的计算预算条件下，您可以显著扩大模型或数据集的规模。特别是在预训练阶段，与稠密模型相比，混合专家模型通常能够更快地达到相同的质量水平。

那么，究竟什么是一个混合专家模型 (MoE) 呢？作为一种基于 Transformer 架构的模型，混合专家模型主要由两个关键部分组成:

稀疏 MoE 层: 这些层代替了传统 Transformer 模型中的前馈网络 (FFN) 层。MoE 层包含若干“专家”(例如 8 个)，每个专家本身是一个独立的神经网络。在实际应用中，这些专家通常是前馈网络 (FFN)，但它们也可以是更复杂的网络结构，甚至可以是 MoE 层本身，从而形成层级式的 MoE 结构。
门控网络或路由: 这个部分用于决定哪些令牌 (token) 被发送到哪个专家。例如，在下图中，“More”这个令牌可能被发送到第二个专家，而“Parameters”这个令牌被发送到第一个专家。有时，一个令牌甚至可以被发送到多个专家。令牌的路由方式是 MoE 使用中的一个关键点，因为路由器由学习的参数组成，并且与网络的其他部分一同进行预训练。
在这里插入图片描述
总结来说，在混合专家模型 (MoE) 中，我们将传统 Transformer 模型中的每个前馈网络 (FFN) 层替换为 MoE 层，其中 MoE 层由两个核心部分组成: 一个门控网络和若干数量的专家。
尽管混合专家模型 (MoE) 提供了若干显著优势，例如更高效的预训练和与稠密模型相比更快的推理速度，但它们也伴随着一些挑战:

训练挑战: 虽然 MoE 能够实现更高效的计算预训练，但它们在微调阶段往往面临泛化能力不足的问题，长期以来易于引发过拟合现象。
推理挑战: MoE 模型虽然可能拥有大量参数，但在推理过程中只使用其中的一部分，这使得它们的推理速度快于具有相同数量参数的稠密模型。然而，这种模型需要将所有参数加载到内存中，因此对内存的需求非常高。

什么是稀疏性?
稀疏性的概念采用了条件计算的思想。在传统的稠密模型中，所有的参数都会对所有输入数据进行处理。相比之下，稀疏性允许我们仅针对整个系统的某些特定部分执行计算。这意味着并非所有参数都会在处理每个输入时被激活或使用，而是根据输入的特定特征或需求，只有部分参数集合被调用和运行。
稀疏性导致了批量大小的不均匀分配和资源利用效率不高的问题。在接下来的部分中，将会讨论让 MoE 高效运行的其他挑战以及相应的解决方案。

那我们应该如何解决这个问题呢？一个可学习的门控网络 (G) 决定将输入的哪一部分发送给哪些专家 (E):
在这里插入图片描述
最初的假设是，需要将输入路由到不止一个专家，以便门控学会如何进行有效的路由选择，因此至少需要选择两个专家。Switch Transformers 就这点进行了更多的研究。
为什么要添加噪声呢？这是为了专家间的负载均衡！
在这里插入图片描述

混合专家模型中令牌的负载均衡

正如之前讨论的，如果所有的令牌都被发送到只有少数几个受欢迎的专家，那么训练效率将会降低。在通常的混合专家模型 (MoE) 训练中，门控网络往往倾向于主要激活相同的几个专家。这种情况可能会自我加强，因为受欢迎的专家训练得更快，因此它们更容易被选择。为了缓解这个问题，引入了一个辅助损失，旨在鼓励给予所有专家相同的重要性。这个损失确保所有专家接收到大致相等数量的训练样本，从而平衡了专家之间的选择。接下来的部分还将探讨专家容量的概念，它引入了一个关于专家可以处理多少令牌的阈值。在 transformers 库中，可以通过 aux_loss 参数来控制辅助损失。’

补充概念：负载均衡：
大模型的负载均衡指的是在多个计算资源（如服务器、GPU等）之间合理分配计算任务，以确保系统的高效运行并避免单个计算节点的过载。在使用深度学习大模型时，由于模型的规模非常庞大，训练和推理过程中会涉及大量的计算资源和内存需求。因此，负载均衡的关键目的是优化计算资源的使用，使得所有计算节点能够均衡地分担工作，从而提高效率、减少延迟，并避免资源浪费。

负载均衡主要包括以下几个方面：

模型并行：将模型分割成不同的部分，分别放在多个计算节点上进行处理。例如，深度神经网络的不同层可以分配到不同的GPU上。

数据并行：将训练数据分割成多个小批次（mini-batch），每个计算节点处理数据的不同部分，计算结果在节点之间进行汇总（如梯度汇总）。

负载感知调度：根据每个计算节点的当前负载情况（例如，计算能力、内存使用等）动态调整任务分配，避免节点因负载过重而成为瓶颈。

跨节点通信优化：在分布式训练中，多个计算节点之间需要频繁交换信息（例如梯度、权重等），优化这些通信的效率对于负载均衡至关重要。

资源管理：合理管理计算资源（如内存、带宽等），避免单一资源的瓶颈。

通过这些策略，负载均衡有助于提高大模型训练和推理的效率，缩短训练时间，避免硬件资源过度利用或浪费。

补充概念：MOE的router

在 MOE（Mixture of Experts，专家混合模型）中，router（路由器）指的是一种机制，负责将输入数据动态地分配给不同的专家子模型。MOE模型的核心思想是，每个输入样本只由少数几个专家进行处理，而不是让所有专家都参与计算。这种设计能够显著提升计算效率，尤其是当专家数量很大时。

MOE中的Router如何工作？
在MOE模型中，通常会有一个 router，它负责决定对于每一个输入样本，哪个（哪些）专家会处理该样本。通常的做法是使用一个小的神经网络（通常称为 gating network 或 routing network）来计算这个路由决策。

Router的工作步骤：
输入数据通过router：每个输入样本首先通过一个router（通常是一个小的神经网络或多层感知机）进行处理。该router会为每个输入样本计算一个概率分布，表示每个专家（或子模型）被选择的概率。

计算路由决策：