Mixtral-8x7B 背后的 MoE 技术

Mixtral-8x7B 背后的 MoE 技术

目录

  • 什么是专家混合(MoE)?
  • MoE 的历史
  • 什么是稀疏性?
  • MoE 和 Transformers
  • Switch Transformers
  • 微调 MoE
  • 加速 MoE 的运行
  • 服务技巧
  • 开源 MoE 项目

正文

什么是专家混合(MoE)?

模型的规模对于提升其质量至关重要。在有限的计算资源下,相较于用更多步骤训练一个小型模型,训练一个大型模型即便步骤更少效果通常更好。MoE 让模型以「远低于传统密集模型的计算成本」进行预训练,这意味着你可以在相同的计算预算下显著扩大模型或数据集的规模。特别是在预训练阶段,MoE 模型能比其同等规模的密集型模型更快地达到相同的性能水平。

那么,MoE 究竟是什么呢?在 Transformer 模型的背景下,MoE 主要由两个部分组成:

  • 稀疏 MoE 层代替了传统的密集前馈网络 (FFN) 层。MoE 层包含若干“专家”,例如 8 个,每个专家都是一个独立的神经网络。实际上,这些专家通常是 FFN,但它们也可以是更复杂的网络,甚至可以是 MoE 本身,形成一个层级结构的 MoE。
  • 门控网络或路由器,用于决定哪些 Token 分配给哪个专家。值得注意的是,一个 Token 可以被分配给多个专家。

如何高效地将 Token 分配给合适的专家,是使用 MoE 技术时需要考虑的关键问题之一。路由器由一系列可学习的参数构成,它与模型的其他部分一起进行预训练。

MoE 的历史

MoE 的概念最早出现在 1991 年的论文《Adaptive Mixture of Local Experts》中。这一理念与集成方法相似,目的是通过监督程序管理一个由不同网络构成的系统,每个网络处理训练样本的一部分。每个单独网络或“专家”,都在输入空间的不同区域有其特长。由单独的门控网络决定每个专家的权重,在训练过程中,同时对专家和门控网络进行训练。

在 2010 至 2015 年间,两个不同的研究领域推动了 MoE 的进一步发展:

  • 将专家作为组件:在传统的 MoE 结构中,系统由一个门控网络和多个专家组成。MoE 作为整体模型已在 SVM、高斯过程等方法中得到应用。Eigen 等人的研究将 MoE 作为更深层网络的一部分进行探索。这意味着 MoE 可以作为多层网络中的一层,使模型在大规模和高效率之间达到平衡。
  • 条件计算:传统网络会将所有输入数据通过每一层。在此期间,Yoshua Bengio 探索了一种基于输入 Token 动态激活或停用网络组件的方法。这些研究促进了在自然语言处理领域对混合专家模型的探索。具体来说,Shazeer 等人 (2017 年的研究,团队成员包括 Geoffrey Hinton 和 Jeff Dean) 将这一理念应用到了一个 137B 的 LSTM (当时的 NLP 主要架构) 上,通过引入稀疏性概念,即使在大规模应用中也能保持快速的推理速度。这项工作主要关注翻译领域,但也面临着高通信成本和训练不稳定等挑战。

什么是稀疏性?

稀疏性基于条件计算的概念。不同于密集模型中所有参数对所有输入都有效,稀疏性让我们能够只激活系统的部分区域。条件计算 (即网络的某些部分仅针对特定样本激活) 使得在不增加计算量的情况下扩大模型规模成为可能,从而在每层 MoE 中使用了数千名专家。

这种方法也带来了挑战。比如,虽然大批量处理通常能提高性能,但在 MoE 中,当数据通过活跃的专家时,实际的批量大小会减小。例如,如果我们的批量输入包含 10 个 Token,可能有 5 个 Token 由一个专家处理,另外 5 个 Token 分别由 5 个不同的专家处理,这导致批量大小不均匀,资源利用率低下。

那我们该如何解决这些问题呢?让我们深入探讨 Shazeer 在翻译领域对 MoE 的研究。通过一个学习型的门控网络 (G),决定将输入的哪些部分分配给哪些专家 (E):

$$y = \sum_{i=1}^{\text{n}} G(x)_i E_i(x)$$

在这种设置中,所有专家都参与处理所有输入——这是一种加权乘法过程。但如果 G 的值为 0 呢?这种情况下,就无需计算相应专家的操作,从而

节约了计算资源。

那么,典型的门控函数是什么样的呢?在传统设置中,我们通常使用一个简单的网络配合 softmax 函数。这个网络会学习如何选择最合适的专家处理输入。

$$G_\sigma(x) = \text{Softmax}(x \cdot W_g)$$

Shazeer 的研究还探索了其他类型的门控机制,如带噪声的 Top-K 门控。这种方法加入了一些可调节的噪声,然后只保留最高的 k 个值。具体来说:

$$添加噪音H(x)i = (x \cdot W_g)_i + \text{StandardNormal()} \cdot \text{Softplus}((x \cdot W{\text{noise}})_i)$$

$$仅保留前 k 个值\text{KeepTopK}(v,k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \text{ elements of } v, \ -\infty & \text{otherwise.} \end{cases}$$

$$应用softmax函数G(x) = \text{Softmax}(\text{KeepTopK}(H(x),k))$$

这种稀疏性带来了一些有趣的特性。如果使用较低的 k 值 (比如一到两个),我们可以比激活许多专家时更快地进行训练和推理。为什么不只选择最顶尖的专家呢?最初的假设是,为了让门控学习如何路由到不同的专家,需要路由到一个以上的专家,因此至少需要选择两个专家。

我们为什么要加入噪声?这是为了实现负载均衡!

MoE 的负载均衡

正如之前所讨论的,如果所有的 token 都被发送到少数几个受欢迎的专家,这将导致训练效率低下。在标准的多专家系统训练中,门控网络倾向于主要激活相同的几位专家。这会形成自我加强的循环,因为得到优先训练的专家会被更频繁地选择。为了减轻这种情况,引入了一种辅助损失来鼓励平等对待所有专家。这种损失确保所有专家获得大致相同数量的训练样本。

下文还将探讨「专家容量」的概念,这涉及到一个专家能处理的 token 数量上限。在 transformers 中,这种辅助损失可以通过 aux_loss 参数来调节。

MoE 和 Transformers

Transformers 模型展示了一个明显的趋势:「增加参数的数量可以显著提高性能」。Google 的 GShard 项目正是在这方面进行了深入探索,试图将 Transformers 模型扩展到超过 6000 亿个参数。在 GShard 中,编码器和解码器里的部分 FFN 层被 MoE 层替代,并采用了一种称为「top-2」的门控机制。下图显示了这种设计在编码器部分的应用。

这种设计对大规模计算尤其有利:当模型扩展到多个设备时,MoE 层在这些设备间共享,而其他层则在每个设备上独立存在。

为了在大规模应用中保持效率和均衡的负载,GShard 团队在设计上做了一些创新,除了引入了类似前一节提到的辅助损失机制外,还包括:

  • 随机路由机制:在 top-2 设计中,我们始终选择表现最优的专家,但第二选择的专家则根据其权重以一定概率被选中。
  • 专家处理能力限制:我们可以设定一个专家能处理的 Token 数量的上限。如果两个专家的处理能力都已达到上限,那么这个 Token 就会被认为是多余的,并通过残差连接传递到下一层,或在某些情况下被直接丢弃。这一概念在 MoE 的应用中非常关键。为什么这样做?因为在模型编译时所有的张量形状都是静态确定的,但我们无法预先知道每个专家将处理多少 Token,因此需要设定一个固定的处理能力上限。在模型推理过程中,只有部分专家会被激活。同时,一些计算过程如自注意力机制会被所有 Token 共享。因此,尽管一个拥有 8 个专家的 470 亿参数模型听起来庞大,但实际上它的计算需求相当于一个 120 亿参数的密集型模型。如果采用 top-2 机制,模型会涉及约 140 亿参数,但由于注意力等操作是共享的,实际上模型真正使用的参数量仍然是 120 亿。

Switch Transformers

尽管 MoE 充满潜力,但它们在训练和微调时面临稳定性挑战。Switch Transformers 这项研究深入剖析了这些问题,并发布了一个具有 2048 个专家和 1.6 万亿参数的 MoE 模型。相较于 T5-XXL,Switch Transformers 的预训练速度提高了四倍。

Switch Transformers 提出了一种处理两种

不同 token 的新型 Transformer 层,包含四个专家。不同于最初至少使用两个专家的设想,Switch Transformers 采用了更简洁的单专家策略。这种策略的影响包括:

  • 简化了路由计算
  • 每个专家处理的批量至少减少了一半
  • 减少了通信成本
  • 保持了模型质量

此外,Switch Transformers 还探讨了专家容量的概念。专家容量的计算公式是:

$$\text{Expert Capacity} = \left( \frac{\text{tokens per batch}}{\text{number of experts}} \right) \times \text{capacity factor}$$

每批 token 数量除以专家数量,再乘以容量因子。按此计算方式,可以均匀分配批次中的 Token 给每个专家。如果容量因子大于 1,可以为 Token 分配不均的情况提供缓冲。但容量增加会带来更高的设备间通信成本,这是一个需要权衡的问题。Switch Transformers 在较低的容量因子下表现优异。

Switch Transformer 的研究者还对上文提到的负载均衡损失进行了简化。在训练过程中,每个 Switch 层的辅助损失会加入到总模型损失中,这种做法促进了均匀的路由分配,并可以通过超参数进行调整。研究者们还尝试了一种选择性的精确度方法,例如在训练专家系统时使用 bfloat16 格式,而在其他计算过程中则采用全精度。降低精度能够显著减少处理器间的通信成本、计算成本以及存储数据的内存需求。但初期实验中,无论是专家系统还是门控网络都采用 bfloat16 进行训练,结果训练过程变得更加不稳定。特别是路由器计算部分,由于其涉及到指数函数,因此更高的精度能够显著改善准确性。

微调 MoE

在对 MoE 模型进行微调时,我们需要特别注意以下几点:

  • 选择合适的专家数量。专家数量过多会导致模型过拟合,而专家数量过少会导致模型欠拟合。
  • 选择合适的门控函数。门控函数决定了哪些 Token 分配给哪个专家。不同的门控函数会对模型的性能产生不同的影响。
  • 调整专家容量。专家容量是指每个专家能处理的 Token 数量。专家容量过大或过小都会影响模型的性能。
  • 使用辅助损失。辅助损失可以鼓励模型平等地对待所有专家,从而提高模型的性能。

加速 MoE 的运行

为了加速 MoE 模型的运行,我们可以采用以下方法:

  • 使用稀疏性。稀疏性可以减少模型的计算量,从而提高模型的运行速度。
  • 使用随机路由。随机路由可以减少模型的通信成本,从而提高模型的运行速度。
  • 使用专家处理能力限制。专家处理能力限制可以防止模型过拟合,从而提高模型的运行速度。
  • 使用选择性的精确度。选择性的精确度可以降低模型的计算成本,从而提高模型的运行速度。

服务技巧

在对 MoE 模型进行服务时,我们需要特别注意以下几点:

  • 选择合适的硬件。MoE 模型对硬件的要求很高,因此我们需要选择合适的硬件来运行模型。
  • 优化模型。我们可以对模型进行优化,以提高模型的性能。
  • 监控模型。我们需要监控模型的运行情况,以确保模型正常运行。

开源 MoE 项目

目前,有许多开源的 MoE 项目,例如:

这些项目提供了 MoE 模型的实现,我们可以直接使用这些项目来训练和微调 MoE 模型。

总结

MoE 是一种强大的技术,可以显著提高模型的性能。在本文中,我们介绍了 MoE 的基本原理、MoE 的历史、MoE 的应用以及 MoE 的加速方法。我们还提供了几个开源的 MoE 项目,供读者参考。

发表评论