正文共：2285字 8圖

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

作者：思成

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

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

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

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

圖1 混合專(zhuān)家層

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

多卡的LLM推理場(chǎng)景下，我們將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層模型并行方案

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

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

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

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

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

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

圖4 Mixtral 8x7B 2卡并行推理-專(zhuān)家并行方案

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

圖5 Mixtral 8x7B 2卡并行推理-專(zhuān)家并行方案

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

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

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

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

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

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

從上述的過(guò)程中，可以看到，模型并行的方案不存在負(fù)載均衡問(wèn)題，同時(shí)受到的約束相對(duì)來(lái)說(shuō)最少（Dimension % Rank == 0），同時(shí)不會(huì)引入額外的通信開(kāi)銷(xiāo)。

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

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

圖8 

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

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