正文共：2285字 8圖

預(yù)計(jì)閱讀時間：6分鐘

作者：思成

軟件生態(tài)中心·應(yīng)用平臺部·模型應(yīng)用組

Mixtral 8x7B是由Mistral AI在23年10月發(fā)布的一種稀疏混合專家混合模型（SMoE）,在多項(xiàng)Benchmark測試中效果優(yōu)于LLaMa2 70B和GPT3.5?；贛ixtral 8x7B得到的Mixtral 8x7B-Instruct指令跟隨模型在Human Evaluation榜單上超過了GPT-3.5 Turbo、Claude-2.1、Gemini Pro和LLaMA 2-Chat 70B。同時，這2個模型免費(fèi)供學(xué)術(shù)和商業(yè)使用。

Mixtral 8x7B同樣是一個Decoder Only的模型，區(qū)別于傳統(tǒng)的LLaMA等模型，F(xiàn)NN層由8個前饋神經(jīng)網(wǎng)絡(luò)（Expert）組成。如果我們從某一個Token的視角看，這個Token會經(jīng)過其中的2個前饋神經(jīng)網(wǎng)絡(luò)或者說Expert。也就是說雖然整個模型的參數(shù)量是46B，但是在推理過程中激活的參數(shù)只有13B。

剛剛提到對于1個Token只有2個Expert被激活，這個機(jī)制是通過路由網(wǎng)絡(luò)（Router）來控制。在Mixtral 8x7B中路由網(wǎng)絡(luò)由1個前饋神經(jīng)網(wǎng)絡(luò)組成。整個混合專家層模型結(jié)構(gòu)以及數(shù)學(xué)表示如圖1^[1]所示。

圖1 混合專家層

雖然Mixtral 8x7B在推理過程中同一時間激活的參數(shù)只有13B左右，但是為了保證推理性能，還是需要將全部參數(shù)（46B）讀入顯存，以A100-80GB為例，對于46B的參數(shù)的模型，按照FP16精度來估算，參數(shù)預(yù)期占用92GB顯存。以batch為20、Context length=1024、Generate length=1024來看、KV Cache需要的顯存為20GB^[2] ，也就是說理想態(tài)下（不考慮顯存碎片，Activate output），至少需要2張A100-80GB完成推理。

多卡的LLM推理場景下，我們將LLM模型結(jié)構(gòu)抽象為Attention和FFN兩部分，因?yàn)镸ixtral 8x7B模型不涉及Attention部分的模型結(jié)構(gòu)改進(jìn)，所以Attention部分的并行方案^[3]遵照標(biāo)準(zhǔn)方案，如圖2所示。

圖2 Attention層模型并行方案

接下來我們來看混合專家層的并行方案設(shè)計(jì)。

Mixtral 8x7B 2張卡的數(shù)據(jù)并行方案如圖3所示。可以看到Expert 1-8的參數(shù)在每張GPU上都被拷貝了一份，對于不同的輸入，每張卡單獨(dú)進(jìn)行推理。這樣的方案不需要在Expert之間進(jìn)行額外的通信開銷。但是顯而易見的問題是顯存開銷大；同時，每一張卡的輸入是一個完整的任務(wù)，需要在每張卡任務(wù)的Batch和序列緯度進(jìn)行MASK，從而拆解為不同的子任務(wù)，并單獨(dú)完成模型推理，最后通過All Gather通信再次在每張卡上還原為完整的任務(wù)。

從上述過程中，可以看到算力除了特殊情況（Batch=1, Sequence length=1）之外并沒有浪費(fèi)，同時在保證Batch * Sequence length % GPU_CNT== 0的情況下也不存在負(fù)載均衡問題。

圖3 Mixtral 8x7B 2卡并行推理-數(shù)據(jù)并行方案

從數(shù)據(jù)并行的方案中可以看到，最大的問題是因?yàn)槊繌埧ǘ夹枰Ａ羲械腅xpert帶來的顯存浪費(fèi)。既然這樣，為什么不每張卡只保留若干Expert呢？專家并行就這樣被提出了。如圖4所示。

可以看到Expert 1-4被分配到了GPU 1上，Expert 5-8被分配到了GPU 2上，很好的解決了顯存開銷問題；但是同時引入了All2All通信（圖4中紅線）。和數(shù)據(jù)并行方案類似，每一張卡的輸入是一個完整的任務(wù)，需要在每張卡任務(wù)的Batch和序列緯度進(jìn)行MASK，從而拆解為不同的子任務(wù)，每個GPU上的子任務(wù)將各自的任務(wù)發(fā)送到對應(yīng)的Expert上完成推理，之后再發(fā)送回自身所屬的GPU，最后通過All Gather通信再次在每張卡上還原為完整的任務(wù)。

圖4 Mixtral 8x7B 2卡并行推理-專家并行方案

另外，圖4中的例子是一個理想態(tài)，Expert的激活剛好是每張卡激活2個Expert。但因?yàn)槊總€Token實(shí)際激活哪個Expert是不確定的，考慮更一般的情況如圖5所示?？梢钥吹竭@時候GPU 1上激活了3個Expert，GPU 2上激活了1個Expert，也就是說專家并行是存在負(fù)載不均衡的。另外考慮Batch=1的推理場景，Generate階段同一時間只會有1張GPU卡被激活，造成算力的浪費(fèi)。

圖5 Mixtral 8x7B 2卡并行推理-專家并行方案

為了進(jìn)一步驗(yàn)證Mixtral 8x7B的負(fù)載均衡問題，我們隨機(jī)挑選了若干問題，并分別統(tǒng)計(jì)了32層的Expert激活情況，如圖6（Heatmap）所示。

圖6 Mixtral 8x7B 不同層Expert激活情況

從負(fù)載均衡角度，全局256個專家中可以明顯看到Layer 14 的Expert 3 是最經(jīng)常被路由到的，而Layer 13的Expert 4是最少被路由到的，前者比后者要勤奮5.72倍。從每層的8個Expert橫向?qū)Ρ纫部梢园l(fā)現(xiàn)，Layer 8 中最勤奮的Expert比最懶惰的Expert多工作4.3倍。

為了解決數(shù)據(jù)并行和專家并行帶來的若干問題，我們討論模型并行的方案，如圖7所示。將Expert按照Embedding維度縱向切分。每張卡保留所有Expert的一個切片。

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

GPU 1上保留了Expert 1-8的部分參數(shù)（左側(cè)實(shí)線），GPU 2上保留了Expert 1-8的另外一部分參數(shù)（右側(cè)實(shí)線）。同時在輸入側(cè)不需要按照Rank進(jìn)行MASK操作，每張卡處理全部的數(shù)據(jù)，并根據(jù)各自的參數(shù)切片得到輸出的切片，最后通過All Reduce通信將輸出融合。

從上述的過程中，可以看到，模型并行的方案不存在負(fù)載均衡問題，同時受到的約束相對來說最少（Dimension % Rank == 0），同時不會引入額外的通信開銷。

綜上，我們發(fā)現(xiàn)基于稀疏混合專家混合模型的LLM推理主要限制因素包括顯存、通信、計(jì)算3方面。而數(shù)據(jù)并行的方式在顯存方面存在巨大的劣勢，模型并行和專家并行在顯存方面的顯存占用一致。所以問題進(jìn)一步簡化為通信和計(jì)算兩方面。

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

圖8 

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

至此，本文簡要介紹了Tecorigin在Mixtral 8x7B模型推理上的探索。未來，期待更多的大模型技術(shù)跟大家一起分享、交流、討論。