深度揭秘爆火MoE!GPT-4關(guān)鍵架構(gòu),成開源模型逆襲殺手锏
Mistral上周末丟出的磁力鏈接震驚了開源圈子,這個(gè)7B×8E的開源MoE大模型性能已經(jīng)到達(dá)了LLaMA2 70B的級(jí)別!
而根據(jù)Jim Fan猜測(cè),如果Mistral內(nèi)部訓(xùn)練了34B×8E或者甚至100B+×8E級(jí)別的模型,那他們的能力很有可能已經(jīng)無限接近GPT-4了。
而在之前對(duì)于GPT-4結(jié)構(gòu)的曝料中,大部分的信息也指向GPT-4很可能是由8個(gè)或者是16個(gè)MoE構(gòu)成。
為什么MoE能成為高性能大模型的必選項(xiàng)?
簡(jiǎn)單來說,MoE是一種神經(jīng)網(wǎng)絡(luò)架構(gòu)設(shè)計(jì),在Transformer模塊中集成了專家/模型層。
當(dāng)數(shù)據(jù)流經(jīng)MoE層時(shí),每個(gè)輸入token都會(huì)動(dòng)態(tài)路由到專家子模型進(jìn)行處理。當(dāng)每個(gè)專家專門從事特定任務(wù)時(shí),這種方法可以實(shí)現(xiàn)更高效的計(jì)算并獲得更好的結(jié)果。
MoE最關(guān)鍵的組件:
- 專家(Expert):MoE層由許多專家、小型MLP或復(fù)雜的LLM(如 Mistral 7B)組成。
- 路由器(Router):路由器確定將哪些輸入token分配給哪些專家。
路由策略有兩種:token選擇路由器或路由器選擇token。
路由器使用softmax門控函數(shù)通過專家或token對(duì)概率分布進(jìn)行建模,并選擇前k個(gè)。
MoE能夠帶來的好處:
- 每個(gè)專家都可以專門處理不同的任務(wù)或數(shù)據(jù)的不同部分。
- MoE構(gòu)架能向LLM添加可學(xué)習(xí)參數(shù),而不增加推理成本。
- 可以利用稀疏矩陣的高效計(jì)算
- 并行計(jì)算所有專家層,以有效利用GPU的并行能力
- 幫助有效地?cái)U(kuò)展模型并減少訓(xùn)練時(shí)間。以更低的計(jì)算成本獲得更好的結(jié)果!
MoE:建立高效大模型的關(guān)鍵技術(shù)
大型語言模型(LLM)席卷了機(jī)器學(xué)習(xí)領(lǐng)域,而隨著現(xiàn)代數(shù)據(jù)集的復(fù)雜性和規(guī)模的增加,每個(gè)數(shù)據(jù)集都包含不同的模式,特征和標(biāo)注之間的關(guān)系截然不同。
——這時(shí),就需要MoE出手了。
專家混合(MoE)就像神經(jīng)網(wǎng)絡(luò)世界中的一種團(tuán)隊(duì)合作技術(shù)。
想象一下,將一項(xiàng)大任務(wù)分解成更小的部分,并讓不同的專家處理每個(gè)部分。然后,有一個(gè)聰明的法官會(huì)根據(jù)情況,決定遵循哪位專家的建議,所有這些建議都混合在一起。
——就像你利用不同的口味,組合成一道好吃的菜。
對(duì)于復(fù)雜的數(shù)據(jù)集,可以劃分為局部子集(local subsets),同樣,將需要預(yù)測(cè)的問題劃分為子任務(wù)(采用領(lǐng)域知識(shí)或者無監(jiān)督聚類算法)。
然后,針對(duì)每個(gè)數(shù)據(jù)子集訓(xùn)練專家模型(Expert Models),專家模型可以是任何模型,比如支持向量機(jī) (SVM) 或者神經(jīng)網(wǎng)絡(luò),每個(gè)專家模型接收相同的輸入模式并進(jìn)行預(yù)測(cè)。
MoE還包含門控模型(Gating Model),用于解釋每個(gè)專家做出的預(yù)測(cè),并根據(jù)輸入選擇信任哪個(gè)專家。
最后,MoE需要一種聚合機(jī)制(Pooling Method),根據(jù)門控模型和專家的輸出進(jìn)行預(yù)測(cè)。
在現(xiàn)實(shí)的應(yīng)用中,研究人員提出一種稱為「稀疏門控專家混合層」的方法,作為原始MoE的迭代,這個(gè)方法提供了一個(gè)通用的神經(jīng)網(wǎng)絡(luò)組件,可以適應(yīng)不同類型的任務(wù)。
稀疏門控專家混合層(Sparsely-Gated Mixture-of-Experts Layer)由許多專家網(wǎng)絡(luò)組成,每個(gè)專家網(wǎng)絡(luò)都是一個(gè)簡(jiǎn)單的前饋神經(jīng)網(wǎng)絡(luò)和一個(gè)可訓(xùn)練的門控網(wǎng)絡(luò)。門控網(wǎng)絡(luò)負(fù)責(zé)選擇這些專家的稀疏組合來處理每個(gè)輸入。
這里的重點(diǎn)是在門控功能中使用稀疏性,——這意味著對(duì)于每個(gè)輸入實(shí)例,門控網(wǎng)絡(luò)只選擇少數(shù)專家進(jìn)行處理,其余的專家保持非活動(dòng)狀態(tài)。
這種稀疏性和專家選擇是針對(duì)每個(gè)輸入動(dòng)態(tài)實(shí)現(xiàn)的,整個(gè)過程具有高度的靈活性和適應(yīng)性,而由于不需要處理網(wǎng)絡(luò)的非活動(dòng)部分,計(jì)算效率大大提高。
——簡(jiǎn)單來說就是,算得快、消耗少、省錢。
MoE層可以分層堆疊,其中主MoE選擇稀疏加權(quán)的「專家」組合。每個(gè)組合都使用一個(gè)MoE圖層。
此外,研究人員還提出了一種名為「Noisy Top-K Gating」的創(chuàng)新技術(shù)。
這種機(jī)制在門控函數(shù)中增加了一個(gè)可調(diào)的高斯噪聲,只保留前K個(gè)值,并將其余值分配給負(fù)無窮大,從而轉(zhuǎn)換為零門控值。
這種方法確保了門控網(wǎng)絡(luò)的稀疏性,同時(shí)保持了對(duì)門控函數(shù)輸出中潛在不連續(xù)性的魯棒性。另外,它還有助于跨專家網(wǎng)絡(luò)進(jìn)行負(fù)載平衡。
MoE和Transformer
下面我們來看一下MoE在Transformer ,也就是當(dāng)下大火的大語言模型中發(fā)揮了怎樣的作用。
MoE作為一種神經(jīng)網(wǎng)絡(luò)架構(gòu)設(shè)計(jì),可以集成進(jìn)Transformer的結(jié)構(gòu)中。
當(dāng)數(shù)據(jù)流經(jīng)MoE層時(shí),每個(gè)輸入(tokens)都會(huì)動(dòng)態(tài)路由到某個(gè)專家模型進(jìn)行計(jì)算,這樣每個(gè)專家都可以專注于特定的任務(wù),更好、更高效地給出結(jié)果。
上圖展示了具有MoE層的Transformer編碼器的演化(對(duì)于解碼器的修改也類似),MoE層取代了Transformer的前饋層。
上圖左邊是標(biāo)準(zhǔn)Transformer模型的編碼器,包括自注意力層和前饋層,交錯(cuò)著殘差連接和歸一化層。
中間部分通過用MoE層替換其他前饋層,得到了MoE Transformer Encoder的模型結(jié)構(gòu)。
上圖右側(cè)是當(dāng)模型擴(kuò)展到多個(gè)設(shè)備時(shí)的情況,MoE層將跨設(shè)備分片,而所有其他層都會(huì)被復(fù)制。
我們可以看到MoE的關(guān)鍵組件就是各種專家模型和路由模塊。
專家模型也可以是小型MLP, 或者復(fù)雜的LLM,如Mistral 7B。
而路由模塊用來確定將哪些輸入tokens分配給哪些專家。
一般有兩種路由策略:token選擇路由器,或路由器選擇token。這里使用softmax門控函數(shù),通過專家模型或者tokens,對(duì)概率分布進(jìn)行建模,并選擇top k。
由此,我們可知MoE層在Transformer中發(fā)揮了重要的作用。
每個(gè)專家都可以專門處理不同的任務(wù)或數(shù)據(jù)的不同部分;使用MoE可以在不增加推理成本的情況下向LLM添加可學(xué)習(xí)參數(shù);
此外,MoE還有助于對(duì)稀疏矩陣進(jìn)行高效計(jì)算;而MoE中的專家層可以并行計(jì)算,這樣就有效利用了GPU的并行能力。
最后,MoE在幫助減少訓(xùn)練時(shí)間的同時(shí),還可以做到有效地?cái)U(kuò)展模型,以更低的計(jì)算成本獲得更好的結(jié)果。
MoE開源再受關(guān)注
在Mistral放出這個(gè)開源的7B×8E的MoE之前,英偉達(dá)和谷歌也放出過其他完全開源的MoE。
曾在英偉達(dá)實(shí)習(xí)的新加坡國(guó)立大學(xué)博士生Fuzhao Xue表示,他們的團(tuán)隊(duì)在4個(gè)月前也開源了一個(gè)80億參數(shù)的MoE模型。
項(xiàng)目地址:https://github.com/XueFuzhao/OpenMoE
數(shù)據(jù)來源
- 一半來自The RedPajama,另一半來自The Stack Dedup
- 為提升模型的推理能力,采用了大量的編程相關(guān)數(shù)據(jù)
模型架構(gòu)
- OpenMoE模型基于「ST-MoE」,但采用了decoder-only架構(gòu)。
其它設(shè)計(jì)
- 采用umT5 tokenizer
- 使用RoPE技術(shù)
- 采用SwiGLU激活函數(shù)
- 設(shè)定2000 token的上下文長(zhǎng)度
BigBench評(píng)估
團(tuán)隊(duì)在BigBench-Lite上進(jìn)行了少樣本測(cè)試,其中包括與BIG-G、BIG-G-Sparse以及GPT-3的對(duì)比。
通過計(jì)算每個(gè)詞元激活的參數(shù)數(shù)量和訓(xùn)練詞元的數(shù)量來大致估計(jì)相對(duì)成本。圖中每個(gè)點(diǎn)的大小代表了相應(yīng)詞元激活的參數(shù)數(shù)量。特別需要指出的是,淺灰色的點(diǎn)表示MoE模型的總參數(shù)量。
對(duì)此,Jim Fan也表示,MoE并不新鮮,它只是沒有得到那么多關(guān)注而已......
比如,谷歌很早之前就開源了基于T5的MoE模型——Switch Transformer。
面臨的挑戰(zhàn)和機(jī)遇
MoE基礎(chǔ)設(shè)施建設(shè)
由于MoE擁有大量可訓(xùn)練參數(shù),理想的軟件環(huán)境應(yīng)該支持靈活組合的專家級(jí)、張量級(jí)、流水線級(jí)和數(shù)據(jù)并行,無論是節(jié)點(diǎn)內(nèi)還是節(jié)點(diǎn)間。
此外,如果能支持簡(jiǎn)單快速的激活卸載和權(quán)重量化,從而減輕MoE權(quán)重的內(nèi)存占用,就更好了。
MoE指令微調(diào)
FLAN-MoE研究提出:盡管將MoE的性能通過特定任務(wù)的微調(diào)轉(zhuǎn)移到下游任務(wù)上存在挑戰(zhàn),但指令微調(diào)卻能有效地與MoE模型協(xié)調(diào)一致。這展示了基于MoE的語言模型巨大的潛力。
MoE 評(píng)估
MoE模型的歸納偏置(Inductive bias)可能在困惑度(perplexity)之外還有其他效果,就像其他自適應(yīng)模型(如Universal Transformer和AdaTape)那樣。
硬件挑戰(zhàn)
值得一提的是,GPU在跨節(jié)點(diǎn)通信方面面臨挑戰(zhàn),因?yàn)槊總€(gè)節(jié)點(diǎn)通常只能配備有限數(shù)量的GPU。這使得專家并行中,通信成為瓶頸。
幸運(yùn)的是,NVIDIA最近推出了DGX GH200,將256個(gè)NVIDIA Grace Hopper Superchips集成到一個(gè)單一GPU中,很大程度上解決了通信帶寬問題,為開源領(lǐng)域的MoE模型的訓(xùn)練和部署提供了幫助。