【元来如此】第三章——Mixtral 8x7B深入挖掘，改变游戏规则的AI模型！

正文共：2285字 8图

预计阅读时间：6分钟

作者：思成

软件生态中心·应用平台部·模型应用组

Mixtral 8x7B是由Mistral AI在23年10月发布的一种稀疏混合专家混合模型（SMoE）,在多项Benchmark测试中效果优于LLaMa2 70B和GPT3.5。基于Mixtral 8x7B得到的Mixtral 8x7B-Instruct指令跟随模型在Human Evaluation榜单上超过了GPT-3.5 Turbo、Claude-2.1、Gemini Pro和LLaMA 2-Chat 70B。同时，这2个模型免费供学术和商业使用。

Mixtral 8x7B同样是一个Decoder Only的模型，区别于传统的LLaMA等模型，FNN层由8个前馈神经网络（Expert）组成。如果我们从某一个Token的视角看，这个Token会经过其中的2个前馈神经网络或者说Expert。也就是说虽然整个模型的参数量是46B，但是在推理过程中激活的参数只有13B。

刚刚提到对于1个Token只有2个Expert被激活，这个机制是通过路由网络（Router）来控制。在Mixtral 8x7B中路由网络由1个前馈神经网络组成。整个混合专家层模型结构以及数学表示如图1^[1]所示。

图1 混合专家层

虽然Mixtral 8x7B在推理过程中同一时间激活的参数只有13B左右，但是为了保证推理性能，还是需要将全部参数（46B）读入显存，以A100-80GB为例，对于46B的参数的模型，按照FP16精度来估算，参数预期占用92GB显存。以batch为20、Context length=1024、Generate length=1024来看、KV Cache需要的显存为20GB^[2] ，也就是说理想态下（不考虑显存碎片，Activate output），至少需要2张A100-80GB完成推理。

多卡的LLM推理场景下，我们将LLM模型结构抽象为Attention和FFN两部分，因为Mixtral 8x7B模型不涉及Attention部分的模型结构改进，所以Attention部分的并行方案^[3]遵照标准方案，如图2所示。

图2 Attention层模型并行方案

接下来我们来看混合专家层的并行方案设计。

Mixtral 8x7B 2张卡的数据并行方案如图3所示。可以看到Expert 1-8的参数在每张GPU上都被拷贝了一份，对于不同的输入，每张卡单独进行推理。这样的方案不需要在Expert之间进行额外的通信开销。但是显而易见的问题是显存开销大；同时，每一张卡的输入是一个完整的任务，需要在每张卡任务的Batch和序列纬度进行MASK，从而拆解为不同的子任务，并单独完成模型推理，最后通过All Gather通信再次在每张卡上还原为完整的任务。

从上述过程中，可以看到算力除了特殊情况（Batch=1, Sequence length=1）之外并没有浪费，同时在保证Batch * Sequence length % GPU_CNT== 0的情况下也不存在负载均衡问题。

图3 Mixtral 8x7B 2卡并行推理-数据并行方案

从数据并行的方案中可以看到，最大的问题是因为每张卡都需要保留所有的Expert带来的显存浪费。既然这样，为什么不每张卡只保留若干Expert呢？专家并行就这样被提出了。如图4所示。

可以看到Expert 1-4被分配到了GPU 1上，Expert 5-8被分配到了GPU 2上，很好的解决了显存开销问题；但是同时引入了All2All通信（图4中红线）。和数据并行方案类似，每一张卡的输入是一个完整的任务，需要在每张卡任务的Batch和序列纬度进行MASK，从而拆解为不同的子任务，每个GPU上的子任务将各自的任务发送到对应的Expert上完成推理，之后再发送回自身所属的GPU，最后通过All Gather通信再次在每张卡上还原为完整的任务。

图4 Mixtral 8x7B 2卡并行推理-专家并行方案

另外，图4中的例子是一个理想态，Expert的激活刚好是每张卡激活2个Expert。但因为每个Token实际激活哪个Expert是不确定的，考虑更一般的情况如图5所示。可以看到这时候GPU 1上激活了3个Expert，GPU 2上激活了1个Expert，也就是说专家并行是存在负载不均衡的。另外考虑Batch=1的推理场景，Generate阶段同一时间只会有1张GPU卡被激活，造成算力的浪费。

图5 Mixtral 8x7B 2卡并行推理-专家并行方案

为了进一步验证Mixtral 8x7B的负载均衡问题，我们随机挑选了若干问题，并分别统计了32层的Expert激活情况，如图6（Heatmap）所示。

图6 Mixtral 8x7B 不同层Expert激活情况

从负载均衡角度，全局256个专家中可以明显看到Layer 14 的Expert 3 是最经常被路由到的，而Layer 13的Expert 4是最少被路由到的，前者比后者要勤奋5.72倍。从每层的8个Expert横向对比也可以发现，Layer 8 中最勤奋的Expert比最懒惰的Expert多工作4.3倍。

为了解决数据并行和专家并行带来的若干问题，我们讨论模型并行的方案，如图7所示。将Expert按照Embedding维度纵向切分。每张卡保留所有Expert的一个切片。

图7 Mixtral 8x7B 2卡并行推理-模型并行方案

GPU 1上保留了Expert 1-8的部分参数（左侧实线），GPU 2上保留了Expert 1-8的另外一部分参数（右侧实线）。同时在输入侧不需要按照Rank进行MASK操作，每张卡处理全部的数据，并根据各自的参数切片得到输出的切片，最后通过All Reduce通信将输出融合。

从上述的过程中，可以看到，模型并行的方案不存在负载均衡问题，同时受到的约束相对来说最少（Dimension % Rank == 0），同时不会引入额外的通信开销。

综上，我们发现基于稀疏混合专家混合模型的LLM推理主要限制因素包括显存、通信、计算3方面。而数据并行的方式在显存方面存在巨大的劣势，模型并行和专家并行在显存方面的显存占用一致。所以问题进一步简化为通信和计算两方面。

太初元碁Tecorigin基于上述分析，深度优化了Mixtral 8x7B模型推理，在Batch size=1，Context length=1024，Generate length=1024下，端到端推理速度百分位（越大越好）对比基于GPU A800 * 2硬件的不同开源LLM推理框架效果如图8所示。

图8 

Mixtral 8x7B 2卡并行推理TecoInference端到端速度对比

至此，本文简要介绍了Tecorigin在Mixtral 8x7B模型推理上的探索。未来，期待更多的大模型技术跟大家一起分享、交流、讨论。