以DBRX 和 Mixtral 8*7B为例，介绍MOE模型

特里斯丹_专注前沿深度学习算法
365 阅读9分钟

模型简介

It uses a fine-grained mixture-of-experts (MoE) architecture with 132B total parameters of which 36B parameters are active on any input. It was pre-trained on 12T tokens of text and code data. Compared to other open MoE models like Mixtral-8x7B and Grok-1, DBRX is fine-grained, meaning it uses a larger number of smaller experts. DBRX has 16 experts and chooses 4, while Mixtral-8x7B and Grok-1 have 8 experts and choose 2. This provides 65x more possible combinations of experts and we found that this improves model quality. DBRX uses rotary position encodings (RoPE), gated linear units (GLU), and grouped query attention (GQA). It uses the GPT-4 tokenizer as provided in the tiktoken repository. We made these choices based on exhaustive evaluation and scaling experiments.

模型总参数量132B
激活参数量36B
专家数16
激活专家数4
其他grouped query attention (GQA)
rotary position encodings (RoPE)
gated linear units (GLU)
数据数据量12T token，自己收集、清洗的高质量数据集；绝大部分是英语，没有在其他语言上做过测试
tokenizertiktoken（速度更快）
词表大小100352
最大上下文长度32K（32768）

指标

在编程和数学任务上尤其强，超过了CodeLLaMA-70B这个领域模型，且速度快2倍；

在两个复合榜单上取得了第一名的成绩（第二名都是Mixtral 8*7B）：

  • the Hugging Face Open LLM Leaderboard (74.5% vs. 72.7% for the next highest model, Mixtral Instruct)
  • the Databricks Gauntlet (66.8% vs. 60.7% for the next highest model, Mixtral Instruct).

推理效率

整体跟参数量成正比，Mixtral 8*7B实际激活13B左右参数，DBRX实际激活36B参数；

上述结果是用半精度、TensorRT-LLM 优化后的模型评测的

模型训练/finetune

用3072张H100训练了3个月

框架

  • 训练框架： Composer

  • 数据流：Streaming

    • 用于从云端获取数据

  • MoE训练框架： Megablocks

    • Megablocks is a lightweight library for MoE training. Crucially, it supports “dropless MoE,” which avoids inefficient padding and is intended to provide deterministic outputs for a given sequence no matter what other sequences are in the batch.

  • 统一框架： LLM Foundry****

    • 把上述三个框架整合到一起，支持训练、测试等

    • LLM Foundry ties all of these libraries together to create a simple LLM pretraining, fine-tuning, and inference experience.

Trick

课程学习curriculum learning

  • 预训练过程中采用了curriculum learning（先简单后复杂， 逐步增加任务难度），动态调整数据混合比例；

  • 难度的评估维度：比如文本的长度、语言的复杂性、任务需解决的逻辑复杂度等。

MoE模型的负载均衡 - 辅助loss

MoE模型训练中存在专家倾斜（Expert Skew）或负载不平衡（Load Imbalance）的问题，DBRX采用了SwitchTransformer中的辅助loss，鼓励模型更均匀地分配输入到所有可用专家上。

理想中，这个额外的均衡损失会激励模型路由去平均地分配任务给每个专家，使得每个专家的负载都接近平均水平。 以下是实现均衡损失的一种方式：

  1. 负载损失（Load Loss）：度量每个专家实际处理的工作量和理想情况下的平均工作量之间的差异。例如，可以计算每个专家处理的tokens数量的标准差或方差，并将其作为损失项。

  2. 容量损失（Capacity Loss）：有时也被考虑进来，确保每个专家不会超出它的处理能力。这可以通过比较每个专家接收的任务量与其容量上限之间差异的损失来实现。

DBRX的实现代码：github.com/databricks/… (与Mixtral 8*7B逻辑一致）

  1. 计算每个专家的token百分比与路由概率： 使用专家掩码计算平均每个专家处理的tokens百分比（tokens_per_expert）和对应的路由概率（router_prob_per_expert）。

  2. 计算总损失： 计算 tokens_per_expert 和 router_prob_per_expert 的元素乘积和，乘以专家数量 num_experts，得到最终的负载平衡损失。

DBRX的router设计
  1. DBRX在router部分添加了扰动：

github.com/databricks/…

  1. DBRX的router还可以支持把token均匀地分配给各个专家（默认没有打开）

github.com/databricks/…

Mixtral 8x7B, a high-quality sparse mixture of experts model (SMoE)

概述

  1. 在大多数榜单都超过Llama 2 70B的效果，推理速度比llama 2 70B快6倍

  2. 支持商用，无需授权，最强的开源LLM

  3. 整个模型参数共46.7B，但由于是稀疏MoE，平均每个token只需要使用12.9B的参数量处理，因此速度也相当于12.9B的模型。

    1. Mixtral 用了 8个7B的experts，总参数量不是56B，而是46.7B，因为只有FFN层被设计为experts，其他层都是共享的

  4. 训练过程包括SFT 和 DPO，以支持更好的指令跟随任务。

Mixtral outperforms Llama 2 70B on most benchmarks with 6x faster inference. It is the strongest open-weight model with a permissive license and the best model overall regarding cost/performance trade-offs.

Concretely, Mixtral has 46.7B total parameters but only uses 12.9B parameters per token. It, therefore, processes input and generates output at the same speed and for the same cost as a 12.9B model.

This model has been optimised through supervised fine-tuning and direct preference optimisation (DPO) for careful instruction following.

核心代码

代码来自第三方复现：

github.com/nateraw/rep…

MoE

暂时无法在飞书文档外展示此内容

在MOE模型中，每个输入token会被分配给多个专家进行处理，而不同的专家可能具有不同的计算能力或特长。为了给每个专家提供相同数量的输入进行处理，需要将输入进行复制。

x.repeat_interleave(self.num_experts_per_tok, dim=0)的作用就是将x张量在维度0上重复扩充self.num_experts_per_tok次，生成相同数量的输入。

具体而言，它会将x的每一行按顺序重复扩充self.num_experts_per_tok次，保持其他维度不变。

通过这样的操作，就可以为每个专家提供相同数量的输入数据，以便进行并行计算和任务分发。

其中，稀疏性体现在：

在这段代码中，使用flat_expert_indices == i条件来筛选出对应于当前专家索引的输入。然而，如果某个专家没有任何输入被分配给它，那么在这段代码中，它的输出将会是一个空张量。

如果存在一些专家没有对应的输入进来，那么对应位置的输出将保持为空，也就是由torch.empty_like(x)生成的空张量。

MoE用在Transformer block中

暂时无法在飞书文档外展示此内容

MoE

MoE: Mixture of Experts

两个关键模块

  1. MoE layers

    1. 即experts，任意网络结构，一般为FFN，也可以是更复杂的结构，比如Transformer、MoE本身
    2. 一般有多个，比如8个（甚至更多，>1000)

  2. A gate network or router

    1. 用来控制哪个token送到哪个特定的expert中

    2. 可以是简单的全连接 + softmax激活函数，也可以是可学习的注意力模块

优缺点

优点

  • 每个专家都可以专门处理不同的任务或数据的不同部分

  • 增大了模型参数量，但不增加计算开销

缺点

  • 不好训练、容易过拟合

    • 流经到每个专家的batch_size不均衡、相比较全量的batch_size小很多
    • 每个batch里并不是每个专家都得到训练；而且有些专家会被强化（被选中的次数更多）

  • 虽然没有增加推理开销，但所有experts都需要加载到内存中

如何缓解上述缺点

首先是不好训练的问题。

先回顾一下门控机制的数学表示：

A learned gating network (G) decides which experts (E) to send a part of the input

大多数情况下，门控机制是简单的网络 + softmax函数

如果把部分权重设置为0，只保留top-k个，那就是个稀疏的门控机制。

还有一些门控机制，比如Noisy Top-K Gating：引入一些可调节的噪声、保留top-k个

稀疏门控专家混合层（Sparsely-Gated Mixture-of-Experts Layer）由许多专家网络组成，每个专家网络都是一个简单的前馈神经网络和一个可训练的门控网络。门控网络负责选择这些专家的稀疏组合来处理每个输入。

这里的重点是在门控功能中使用稀疏性，——这意味着对于每个输入实例，门控网络只选择少数专家进行处理，其余的专家保持非活动状态。

稀疏性的好处就是减少推理时的计算开销；但为啥不直接选top-k个专家，而是要添加噪声？

避免学的不均衡、都走某几个优势专家那了，劣势的那些专家可能从来不会被激活。

所以训练的时候添加一些可调节的噪声，可以更均衡地选到大多数专家，整个训练也就更好做了。

此外，还有很多方法，比如训练的时候为每个专家添加一个计数器，如果使用的次数过多，就降低该专家被选中的概率；

或者添加一个辅助loss，尽可能地让各个专家被均匀地选中。

tinkerd.net/blog/machin…

稀疏topk的实现

暂时无法在飞书文档外展示此内容

然后是参数量多了很多的问题。

有一些方向和思路

  • Distilling Mixtral into a dense model
  • Explore model merging techniques of the experts and their impact in inference time
  • Perform extreme quantization techniques of Mixtral

  • 蒸馏：把稀疏的多个专家，蒸馏成原始的密集模型

  • 考虑专家的合并

  • 量化

此外，MoE还有助于对稀疏矩阵进行高效计算；而MoE中的专家层可以并行计算，这样就有效利用了GPU的并行能力。

最后，MoE在帮助减少训练时间的同时，还可以做到有效地扩展模型，以更低的计算成本获得更好的结果。

不同的token送到不同的exper中，分别处理后，再聚合到一起

Mistral 相关资源

各种MoE模型

榜单

模型

DeepSeekMoE 16B 是一个具有 16.4B 参数的混合专家 （MoE） 语言模型。它采用了创新的 MoE 架构，该架构涉及两个主要策略：细粒度专家细分和共享专家隔离。它是在 2T tokens 从头开始训练的，并表现出与 DeekSeek 7B 和 LLaMA2 7B 相当的性能，只有大约 40% 的计算量。

论文

scale.com/blog/fine-t…

avatar
文章
阅读
粉丝