隨著 Mixtral 8x7B (announcement, model card) 的推出,一種稱為混合專家模型 (Mixed Expert Models,簡(jiǎn)稱 MoEs) 的 Transformer 模型在開源人工智能社區(qū)引起了廣泛關(guān)注。在本篇博文中,我們將深入探討 MoEs 的核心組件、訓(xùn)練方法,以及在推理過(guò)程中需要考量的各種因素。
讓我們開始吧!
簡(jiǎn)短總結(jié)
混合專家模型 (MoEs):
與稠密模型相比,預(yù)訓(xùn)練速度更快
與具有相同參數(shù)數(shù)量的模型相比,具有更快的推理速度
需要大量顯存,因?yàn)樗袑<蚁到y(tǒng)都需要加載到內(nèi)存中
在微調(diào)方面存在諸多挑戰(zhàn),但 近期的研究 表明,對(duì)混合專家模型進(jìn)行指令調(diào)優(yōu)具有很大的潛力。
讓我們開始吧!
什么是混合專家模型?
模型規(guī)模是提升模型性能的關(guān)鍵因素之一。在有限的計(jì)算資源預(yù)算下,用更少的訓(xùn)練步數(shù)訓(xùn)練一個(gè)更大的模型,往往比用更多的步數(shù)訓(xùn)練一個(gè)較小的模型效果更佳。
混合專家模型 (MoE) 的一個(gè)顯著優(yōu)勢(shì)是它們能夠在遠(yuǎn)少于稠密模型所需的計(jì)算資源下進(jìn)行有效的預(yù)訓(xùn)練。這意味著在相同的計(jì)算預(yù)算條件下,您可以顯著擴(kuò)大模型或數(shù)據(jù)集的規(guī)模。特別是在預(yù)訓(xùn)練階段,與稠密模型相比,混合專家模型通常能夠更快地達(dá)到相同的質(zhì)量水平。
那么,究竟什么是一個(gè)混合專家模型 (MoE) 呢?作為一種基于 Transformer 架構(gòu)的模型,混合專家模型主要由兩個(gè)關(guān)鍵部分組成:
稀疏 MoE 層: 這些層代替了傳統(tǒng) Transformer 模型中的前饋網(wǎng)絡(luò) (FFN) 層。MoE 層包含若干“專家”(例如 8 個(gè)),每個(gè)專家本身是一個(gè)獨(dú)立的神經(jīng)網(wǎng)絡(luò)。在實(shí)際應(yīng)用中,這些專家通常是前饋網(wǎng)絡(luò) (FFN),但它們也可以是更復(fù)雜的網(wǎng)絡(luò)結(jié)構(gòu),甚至可以是 MoE 層本身,從而形成層級(jí)式的 MoE 結(jié)構(gòu)。
門控網(wǎng)絡(luò)或路由: 這個(gè)部分用于決定哪些令牌 (token) 被發(fā)送到哪個(gè)專家。例如,在下圖中,“More”這個(gè)令牌可能被發(fā)送到第二個(gè)專家,而“Parameters”這個(gè)令牌被發(fā)送到第一個(gè)專家。有時(shí),一個(gè)令牌甚至可以被發(fā)送到多個(gè)專家。令牌的路由方式是 MoE 使用中的一個(gè)關(guān)鍵點(diǎn),因?yàn)槁酚善饔蓪W(xué)習(xí)的參數(shù)組成,并且與網(wǎng)絡(luò)的其他部分一同進(jìn)行預(yù)訓(xùn)練。
Switch Transformers paper 論文中的 MoE layer
總結(jié)來(lái)說(shuō),在混合專家模型 (MoE) 中,我們將傳統(tǒng) Transformer 模型中的每個(gè)前饋網(wǎng)絡(luò) (FFN) 層替換為 MoE 層,其中 MoE 層由兩個(gè)核心部分組成: 一個(gè)門控網(wǎng)絡(luò)和若干數(shù)量的專家。
盡管混合專家模型 (MoE) 提供了若干顯著優(yōu)勢(shì),例如更高效的預(yù)訓(xùn)練和與稠密模型相比更快的推理速度,但它們也伴隨著一些挑戰(zhàn):
訓(xùn)練挑戰(zhàn): 雖然 MoE 能夠?qū)崿F(xiàn)更高效的計(jì)算預(yù)訓(xùn)練,但它們?cè)谖⒄{(diào)階段往往面臨泛化能力不足的問(wèn)題,長(zhǎng)期以來(lái)易于引發(fā)過(guò)擬合現(xiàn)象。
推理挑戰(zhàn): MoE 模型雖然可能擁有大量參數(shù),但在推理過(guò)程中只使用其中的一部分,這使得它們的推理速度快于具有相同數(shù)量參數(shù)的稠密模型。然而,這種模型需要將所有參數(shù)加載到內(nèi)存中,因此對(duì)內(nèi)存的需求非常高。以 Mixtral 8x7B 這樣的 MoE 為例,需要足夠的 VRAM 來(lái)容納一個(gè) 47B 參數(shù)的稠密模型。之所以是 47B 而不是 8 x 7B = 56B,是因?yàn)樵?MoE 模型中,只有 FFN 層被視為獨(dú)立的專家,而模型的其他參數(shù)是共享的。此外,假設(shè)每個(gè)令牌只使用兩個(gè)專家,那么推理速度 (以 FLOPs 計(jì)算) 類似于使用 12B 模型 (而不是 14B 模型),因?yàn)殡m然它進(jìn)行了 2x7B 的矩陣乘法計(jì)算,但某些層是共享的。
了解了 MoE 的基本概念后,讓我們進(jìn)一步探索推動(dòng)這類模型發(fā)展的研究。
混合專家模型簡(jiǎn)史
混合專家模型 (MoE) 的理念起源于 1991 年的論文 Adaptive Mixture of Local Experts。這個(gè)概念與集成學(xué)習(xí)方法相似,旨在為由多個(gè)單獨(dú)網(wǎng)絡(luò)組成的系統(tǒng)建立一個(gè)監(jiān)管機(jī)制。在這種系統(tǒng)中,每個(gè)網(wǎng)絡(luò) (被稱為“專家”) 處理訓(xùn)練樣本的不同子集,專注于輸入空間的特定區(qū)域。那么,如何選擇哪個(gè)專家來(lái)處理特定的輸入呢?這就是門控網(wǎng)絡(luò)發(fā)揮作用的地方,它決定了分配給每個(gè)專家的權(quán)重。在訓(xùn)練過(guò)程中,這些專家和門控網(wǎng)絡(luò)都同時(shí)接受訓(xùn)練,以優(yōu)化它們的性能和決策能力。
在 2010 至 2015 年間,兩個(gè)獨(dú)立的研究領(lǐng)域?yàn)榛旌蠈<夷P?(MoE) 的后續(xù)發(fā)展做出了顯著貢獻(xiàn):
組件專家: 在傳統(tǒng)的 MoE 設(shè)置中,整個(gè)系統(tǒng)由一個(gè)門控網(wǎng)絡(luò)和多個(gè)專家組成。在支持向量機(jī) (SVMs) 、高斯過(guò)程和其他方法的研究中,MoE 通常被視為整個(gè)模型的一部分。然而,Eigen、Ranzato 和 Ilya 的研究 探索了將 MoE 作為更深層網(wǎng)絡(luò)的一個(gè)組件。這種方法允許將 MoE 嵌入到多層網(wǎng)絡(luò)中的某一層,使得模型既大又高效。
條件計(jì)算: 傳統(tǒng)的神經(jīng)網(wǎng)絡(luò)通過(guò)每一層處理所有輸入數(shù)據(jù)。在這一時(shí)期,Yoshua Bengio 等研究人員開始探索基于輸入令牌動(dòng)態(tài)激活或停用網(wǎng)絡(luò)組件的方法。
這些研究的融合促進(jìn)了在自然語(yǔ)言處理 (NLP) 領(lǐng)域?qū)旌蠈<夷P偷奶剿鳌L貏e是在 2017 年,Shazeer 等人 (團(tuán)隊(duì)包括 Geoffrey Hinton 和 Jeff Dean,后者有時(shí)被戲稱為 “谷歌的 Chuck Norris”) 將這一概念應(yīng)用于 137B 的 LSTM (當(dāng)時(shí)被廣泛應(yīng)用于 NLP 的架構(gòu),由 Schmidhuber 提出)。通過(guò)引入稀疏性,這項(xiàng)工作在保持極高規(guī)模的同時(shí)實(shí)現(xiàn)了快速的推理速度。這項(xiàng)工作主要集中在翻譯領(lǐng)域,但面臨著如高通信成本和訓(xùn)練不穩(wěn)定性等多種挑戰(zhàn)。
Outrageously Large Neural Network 論文中的 MoE layer
混合專家模型 (MoE) 的引入使得訓(xùn)練具有數(shù)千億甚至萬(wàn)億參數(shù)的模型成為可能,如開源的 1.6 萬(wàn)億參數(shù)的 Switch Transformers 等。這種技術(shù)不僅在自然語(yǔ)言處理 (NLP) 領(lǐng)域得到了廣泛應(yīng)用,也開始在計(jì)算機(jī)視覺(jué)領(lǐng)域進(jìn)行探索。然而,本篇博客文章將主要聚焦于自然語(yǔ)言處理領(lǐng)域的應(yīng)用和探討。
什么是稀疏性?
稀疏性的概念采用了條件計(jì)算的思想。在傳統(tǒng)的稠密模型中,所有的參數(shù)都會(huì)對(duì)所有輸入數(shù)據(jù)進(jìn)行處理。相比之下,稀疏性允許我們僅針對(duì)整個(gè)系統(tǒng)的某些特定部分執(zhí)行計(jì)算。這意味著并非所有參數(shù)都會(huì)在處理每個(gè)輸入時(shí)被激活或使用,而是根據(jù)輸入的特定特征或需求,只有部分參數(shù)集合被調(diào)用和運(yùn)行。
讓我們深入分析 Shazeer 對(duì)混合專家模型 (MoE) 在翻譯應(yīng)用中的貢獻(xiàn)。條件計(jì)算的概念 (即僅在每個(gè)樣本的基礎(chǔ)上激活網(wǎng)絡(luò)的不同部分) 使得在不增加額外計(jì)算負(fù)擔(dān)的情況下擴(kuò)展模型規(guī)模成為可能。這一策略在每個(gè) MoE 層中實(shí)現(xiàn)了數(shù)以千計(jì)甚至更多的專家的有效利用。
這種稀疏性設(shè)置確實(shí)帶來(lái)了一些挑戰(zhàn)。例如,在混合專家模型 (MoE) 中,盡管較大的批量大小通常有利于提高性能,但當(dāng)數(shù)據(jù)通過(guò)激活的專家時(shí),實(shí)際的批量大小可能會(huì)減少。比如,假設(shè)我們的輸入批量包含 10 個(gè)令牌,可能會(huì)有五個(gè)令牌被路由到同一個(gè)專家,而剩下的五個(gè)令牌分別被路由到不同的專家。這導(dǎo)致了批量大小的不均勻分配和資源利用效率不高的問(wèn)題。在接下來(lái)的部分中,將會(huì)討論讓 MoE 高效運(yùn)行的其他挑戰(zhàn)以及相應(yīng)的解決方案。
那我們應(yīng)該如何解決這個(gè)問(wèn)題呢?一個(gè)可學(xué)習(xí)的門控網(wǎng)絡(luò) (G) 決定將輸入的哪一部分發(fā)送給哪些專家 (E):
在這種設(shè)置下,雖然所有專家都會(huì)對(duì)所有輸入進(jìn)行運(yùn)算,但通過(guò)門控網(wǎng)絡(luò)的輸出進(jìn)行加權(quán)乘法操作。但是,如果 G (門控網(wǎng)絡(luò)的輸出) 為 0 會(huì)發(fā)生什么呢?如果是這種情況,就沒(méi)有必要計(jì)算相應(yīng)的專家操作,因此我們可以節(jié)省計(jì)算資源。那么一個(gè)典型的門控函數(shù)是什么呢?一個(gè)典型的門控函數(shù)通常是一個(gè)帶有 softmax 函數(shù)的簡(jiǎn)單的網(wǎng)絡(luò)。這個(gè)網(wǎng)絡(luò)將學(xué)習(xí)將輸入發(fā)送給哪個(gè)專家。
Shazeer 等人的工作還探索了其他的門控機(jī)制,其中包括帶噪聲的 TopK 門控 (Noisy Top-K Gating)。這種門控方法引入了一些可調(diào)整的噪聲,然后保留前 k 個(gè)值。具體來(lái)說(shuō):
添加一些噪聲
選擇保留前 K 個(gè)值
應(yīng)用 Softmax 函數(shù)
這種稀疏性引入了一些有趣的特性。通過(guò)使用較低的 k 值 (例如 1 或 2),我們可以比激活多個(gè)專家時(shí)更快地進(jìn)行訓(xùn)練和推理。為什么不僅選擇最頂尖的專家呢?最初的假設(shè)是,需要將輸入路由到不止一個(gè)專家,以便門控學(xué)會(huì)如何進(jìn)行有效的路由選擇,因此至少需要選擇兩個(gè)專家。Switch Transformers 就這點(diǎn)進(jìn)行了更多的研究。
我們?yōu)槭裁匆砑釉肼暷兀窟@是為了專家間的負(fù)載均衡!
混合專家模型中令牌的負(fù)載均衡
正如之前討論的,如果所有的令牌都被發(fā)送到只有少數(shù)幾個(gè)受歡迎的專家,那么訓(xùn)練效率將會(huì)降低。在通常的混合專家模型 (MoE) 訓(xùn)練中,門控網(wǎng)絡(luò)往往傾向于主要激活相同的幾個(gè)專家。這種情況可能會(huì)自我加強(qiáng),因?yàn)槭軞g迎的專家訓(xùn)練得更快,因此它們更容易被選擇。為了緩解這個(gè)問(wèn)題,引入了一個(gè)輔助損失,旨在鼓勵(lì)給予所有專家相同的重要性。這個(gè)損失確保所有專家接收到大致相等數(shù)量的訓(xùn)練樣本,從而平衡了專家之間的選擇。接下來(lái)的部分還將探討專家容量的概念,它引入了一個(gè)關(guān)于專家可以處理多少令牌的閾值。在 transformers 庫(kù)中,可以通過(guò) aux_loss 參數(shù)來(lái)控制輔助損失。
MoEs and Transformers
Transformer 類模型明確表明,增加參數(shù)數(shù)量可以提高性能,因此谷歌使用 GShard 嘗試將 Transformer 模型的參數(shù)量擴(kuò)展到超過(guò) 6000 億并不令人驚訝。
GShard 將在編碼器和解碼器中的每個(gè)前饋網(wǎng)絡(luò) (FFN) 層中的替換為使用 Top-2 門控的混合專家模型 (MoE) 層。下圖展示了編碼器部分的結(jié)構(gòu)。這種架構(gòu)對(duì)于大規(guī)模計(jì)算非常有效: 當(dāng)擴(kuò)展到多個(gè)設(shè)備時(shí),MoE 層在不同設(shè)備間共享,而其他所有層則在每個(gè)設(shè)備上復(fù)制。我們將在 “讓 MoE 起飛”部分對(duì)這一點(diǎn)進(jìn)行更詳細(xì)的討論。
GShard 論文中的 MoE Transformer Encoder
為了保持負(fù)載平衡和訓(xùn)練效率,GShard 的作者除了引入了上一節(jié)中討論的類似輔助損失外,還引入了一些關(guān)鍵變化:
隨機(jī)路由: 在 Top-2 設(shè)置中,我們始終選擇排名最高的專家,但第二個(gè)專家是根據(jù)其權(quán)重比例隨機(jī)選擇的。
專家容量: 我們可以設(shè)定一個(gè)閾值,定義一個(gè)專家能處理多少令牌。如果兩個(gè)專家的容量都達(dá)到上限,令牌就會(huì)溢出,并通過(guò)殘差連接傳遞到下一層,或在某些情況下被完全丟棄。專家容量是 MoE 中最重要的概念之一。為什么需要專家容量呢?因?yàn)樗袕埩康男螤钤诰幾g時(shí)是靜態(tài)確定的,我們無(wú)法提前知道多少令牌會(huì)分配給每個(gè)專家,因此需要一個(gè)固定的容量因子。
GShard 的工作對(duì)適用于 MoE 的并行計(jì)算模式也做出了重要貢獻(xiàn),但這些內(nèi)容的討論超出了這篇博客的范圍。
注意: 在推理過(guò)程中,只有部分專家被激活。同時(shí),有些計(jì)算過(guò)程是共享的,例如自注意力 (self-attention) 機(jī)制,它適用于所有令牌。這就解釋了為什么我們可以使用相當(dāng)于 12B 稠密模型的計(jì)算資源來(lái)運(yùn)行一個(gè)包含 8 個(gè)專家的 47B 模型。如果我們采用 Top-2 門控,模型會(huì)使用高達(dá) 14B 的參數(shù)。但是,由于自注意力操作 (專家間共享) 的存在,實(shí)際上模型運(yùn)行時(shí)使用的參數(shù)數(shù)量是 12B。
Switch Transformers
盡管混合專家模型 (MoE) 顯示出了很大的潛力,但它們?cè)谟?xùn)練和微調(diào)過(guò)程中存在穩(wěn)定性問(wèn)題。Switch Transformers 是一項(xiàng)非常激動(dòng)人心的工作,它深入研究了這些話題。作者甚至在 Hugging Face 上發(fā)布了一個(gè) 1.6 萬(wàn)億參數(shù)的 MoE,擁有 2048 個(gè)專家,你可以使用 transformers 庫(kù)來(lái)運(yùn)行它。Switch Transformers 實(shí)現(xiàn)了與 T5-XXL 相比 4 倍的預(yù)訓(xùn)練速度提升。
Switch Transformer 論文中的 Switch Transformer Layer
就像在 GShard 中一樣,作者用混合專家模型 (MoE) 層替換了前饋網(wǎng)絡(luò) (FFN) 層。Switch Transformers 提出了一個(gè) Switch Transformer 層,它接收兩個(gè)輸入 (兩個(gè)不同的令牌) 并擁有四個(gè)專家。
與最初使用至少兩個(gè)專家的想法相反,Switch Transformers 采用了簡(jiǎn)化的單專家策略。這種方法的效果包括:
減少門控網(wǎng)絡(luò) (路由) 計(jì)算負(fù)擔(dān)
每個(gè)專家的批量大小至少可以減半
降低通信成本
保持模型質(zhì)量
Switch Transformers 也對(duì)專家容量這個(gè)概念進(jìn)行了研究。
上述建議的容量是將批次中的令牌數(shù)量均勻分配到各個(gè)專家。如果我們使用大于 1 的容量因子,我們?yōu)榱钆品峙洳煌耆胶鈺r(shí)提供了一個(gè)緩沖。增加容量因子會(huì)導(dǎo)致更高的設(shè)備間通信成本,因此這是一個(gè)需要考慮的權(quán)衡。特別值得注意的是,Switch Transformers 在低容量因子 (例如 1 至 1.25) 下表現(xiàn)出色。
Switch Transformer 的作者還重新審視并簡(jiǎn)化了前面章節(jié)中提到的負(fù)載均衡損失。在訓(xùn)練期間,對(duì)于每個(gè) Switch 層的輔助損失被添加到總模型損失中。這種損失鼓勵(lì)均勻路由,并可以使用超參數(shù)進(jìn)行加權(quán)。
作者還嘗試了混合精度的方法,例如用 bfloat16 精度訓(xùn)練專家,同時(shí)對(duì)其余計(jì)算使用全精度進(jìn)行。較低的精度可以減少處理器間的通信成本、計(jì)算成本以及存儲(chǔ)張量的內(nèi)存。然而,在最初的實(shí)驗(yàn)中,當(dāng)專家和門控網(wǎng)絡(luò)都使用 bfloat16 精度訓(xùn)練時(shí),出現(xiàn)了不穩(wěn)定的訓(xùn)練現(xiàn)象。這種不穩(wěn)定性特別是由路由計(jì)算引起的,因?yàn)槁酚缮婕爸笖?shù)函數(shù)等操作,這些操作對(duì)精度要求較高。因此,為了保持計(jì)算的穩(wěn)定性和精確性,保持更高的精度是重要的。為了減輕不穩(wěn)定性,路由過(guò)程也使用了全精度。
使用混合精度不會(huì)降低模型質(zhì)量并可實(shí)現(xiàn)更快的訓(xùn)練
這個(gè) Jupyter Notebook 展示了如何對(duì) Switch Transformers 進(jìn)行微調(diào)以進(jìn)行摘要生成的詳細(xì)指南。然而,在開始微調(diào) Switch Transformers 之前,強(qiáng)烈建議您先閱讀關(guān)于微調(diào)混合專家模型部分的內(nèi)容。
Switch Transformers 采用了編碼器 - 解碼器的架構(gòu),實(shí)現(xiàn)了與 T5 類似的混合專家模型 (MoE) 版本。GLaM 這篇工作探索了如何使用僅為原來(lái) 1/3 的計(jì)算資源 (因?yàn)?MoE 模型在訓(xùn)練時(shí)需要的計(jì)算量較少,從而能夠顯著降低碳足跡) 來(lái)訓(xùn)練與 GPT-3 質(zhì)量相匹配的模型來(lái)提高這些模型的規(guī)模。作者專注于僅解碼器 (decoder-only) 的模型以及少樣本和單樣本評(píng)估,而不是微調(diào)。他們使用了 Top-2 路由和更大的容量因子。此外,他們探討了將容量因子作為一個(gè)動(dòng)態(tài)度量,根據(jù)訓(xùn)練和評(píng)估期間所使用的計(jì)算量進(jìn)行調(diào)整。
用 Router z-loss 穩(wěn)定模型訓(xùn)練
之前討論的平衡損失可能會(huì)導(dǎo)致穩(wěn)定性問(wèn)題。我們可以使用許多方法來(lái)穩(wěn)定稀疏模型的訓(xùn)練,但這可能會(huì)犧牲模型質(zhì)量。例如,引入 dropout 可以提高穩(wěn)定性,但會(huì)導(dǎo)致模型質(zhì)量下降。另一方面,增加更多的乘法分量可以提高質(zhì)量,但會(huì)降低模型穩(wěn)定性。
ST-MoE 引入的 Router z-loss 在保持了模型性能的同時(shí)顯著提升了訓(xùn)練的穩(wěn)定性。這種損失機(jī)制通過(guò)懲罰門控網(wǎng)絡(luò)輸入的較大 logits 來(lái)起作用,目的是促使數(shù)值的絕對(duì)大小保持較小,這樣可以有效減少計(jì)算中的舍入誤差。這一點(diǎn)對(duì)于那些依賴指數(shù)函數(shù)進(jìn)行計(jì)算的門控網(wǎng)絡(luò)尤其重要。為了深入了解這一機(jī)制,建議參考原始論文以獲得更全面的細(xì)節(jié)。
專家如何學(xué)習(xí)?
ST-MoE 的研究者們發(fā)現(xiàn),編碼器中不同的專家傾向于專注于特定類型的令牌或淺層概念。例如,某些專家可能專門處理標(biāo)點(diǎn)符號(hào),而其他專家則專注于專有名詞等。與此相反,解碼器中的專家通常具有較低的專業(yè)化程度。此外,研究者們還對(duì)這一模型進(jìn)行了多語(yǔ)言訓(xùn)練。盡管人們可能會(huì)預(yù)期每個(gè)專家處理一種特定語(yǔ)言,但實(shí)際上并非如此。由于令牌路由和負(fù)載均衡的機(jī)制,沒(méi)有任何專家被特定配置以專門處理某一特定語(yǔ)言。
ST-MoE 論文中顯示了哪些令牌組被發(fā)送給了哪個(gè)專家的表格
專家的數(shù)量對(duì)預(yù)訓(xùn)練有何影響?
增加更多專家可以提升處理樣本的效率和加速模型的運(yùn)算速度,但這些優(yōu)勢(shì)隨著專家數(shù)量的增加而遞減 (尤其是當(dāng)專家數(shù)量達(dá)到 256 或 512 之后更為明顯)。同時(shí),這也意味著在推理過(guò)程中,需要更多的顯存來(lái)加載整個(gè)模型。值得注意的是,Switch Transformers 的研究表明,其在大規(guī)模模型中的特性在小規(guī)模模型下也同樣適用,即便是每層僅包含 2、4 或 8 個(gè)專家。
微調(diào)混合專家模型
4.36.0 版本的 transformers 庫(kù)支持 Mixtral 模型。你可以用以下命令進(jìn)行安裝: pip install "transformers==4.36.0 --upgrade
稠密模型和稀疏模型在過(guò)擬合的動(dòng)態(tài)表現(xiàn)上存在顯著差異。稀疏模型更易于出現(xiàn)過(guò)擬合現(xiàn)象,因此在處理這些模型時(shí),嘗試更強(qiáng)的內(nèi)部正則化措施是有益的,比如使用更高比例的 dropout。例如,我們可以為稠密層設(shè)定一個(gè)較低的 dropout 率,而為稀疏層設(shè)置一個(gè)更高的 dropout 率,以此來(lái)優(yōu)化模型性能。
在微調(diào)過(guò)程中是否使用輔助損失是一個(gè)需要決策的問(wèn)題。ST-MoE 的作者嘗試關(guān)閉輔助損失,發(fā)現(xiàn)即使高達(dá) 11% 的令牌被丟棄,模型的質(zhì)量也沒(méi)有顯著受到影響。令牌丟棄可能是一種正則化形式,有助于防止過(guò)擬合。
Switch Transformers 的作者觀察到,在相同的預(yù)訓(xùn)練困惑度下,稀疏模型在下游任務(wù)中的表現(xiàn)不如對(duì)應(yīng)的稠密模型,特別是在重理解任務(wù) (如 SuperGLUE) 上。另一方面,對(duì)于知識(shí)密集型任務(wù) (如 TriviaQA),稀疏模型的表現(xiàn)異常出色。作者還觀察到,在微調(diào)過(guò)程中,較少的專家的數(shù)量有助于改善性能。另一個(gè)關(guān)于泛化問(wèn)題確認(rèn)的發(fā)現(xiàn)是,模型在小型任務(wù)上表現(xiàn)較差,但在大型任務(wù)上表現(xiàn)良好。
在小任務(wù) (左圖) 中,我們可以看到明顯的過(guò)擬合,因?yàn)橄∈枘P驮隍?yàn)證集中的表現(xiàn)要差得多。在較大的任務(wù) (右圖) 中,MoE 則表現(xiàn)良好。該圖來(lái)自 ST-MoE 論文
一種可行的微調(diào)策略是嘗試凍結(jié)所有非專家層的權(quán)重。實(shí)踐中,這會(huì)導(dǎo)致性能大幅下降,但這符合我們的預(yù)期,因?yàn)榛旌蠈<夷P?(MoE) 層占據(jù)了網(wǎng)絡(luò)的主要部分。我們可以嘗試相反的方法: 僅凍結(jié) MoE 層的參數(shù)。實(shí)驗(yàn)結(jié)果顯示,這種方法幾乎與更新所有參數(shù)的效果相當(dāng)。這種做法可以加速微調(diào)過(guò)程,并降低顯存需求。
通過(guò)僅凍結(jié) MoE 層,我們可以在保持質(zhì)量的同時(shí)加快訓(xùn)練速度。該圖來(lái)自 ST-MoE 論文
在微調(diào)稀疏混合專家模型 (MoE) 時(shí)需要考慮的最后一個(gè)問(wèn)題是,它們有特別的微調(diào)超參數(shù)設(shè)置——例如,稀疏模型往往更適合使用較小的批量大小和較高的學(xué)習(xí)率,這樣可以獲得更好的訓(xùn)練效果。
降低學(xué)習(xí)率和調(diào)大批量可以提升稀疏模型微調(diào)質(zhì)量。該圖來(lái)自 ST-MoE 論文
此時(shí),您可能會(huì)對(duì)人們微調(diào) MoE 中遇到的這些挑戰(zhàn)而感到沮喪,但最近的一篇論文 《MoEs Meets Instruction Tuning》 (2023 年 7 月) 帶來(lái)了令人興奮的發(fā)現(xiàn)。這篇論文進(jìn)行了以下實(shí)驗(yàn):
單任務(wù)微調(diào)
多任務(wù)指令微調(diào)
多任務(wù)指令微調(diào)后接單任務(wù)微調(diào)
當(dāng)研究者們對(duì) MoE 和對(duì)應(yīng)性能相當(dāng)?shù)?T5 模型進(jìn)行微調(diào)時(shí),他們發(fā)現(xiàn) T5 的對(duì)應(yīng)模型表現(xiàn)更為出色。然而,當(dāng)研究者們對(duì) Flan T5 (一種 T5 的指令優(yōu)化版本) 的 MoE 版本進(jìn)行微調(diào)時(shí),MoE 的性能顯著提升。更值得注意的是,F(xiàn)lan-MoE 相比原始 MoE 的性能提升幅度超過(guò)了 Flan T5 相對(duì)于原始 T5 的提升,這意味著 MoE 模型可能從指令式微調(diào)中獲益更多,甚至超過(guò)了稠密模型。此外,MoE 在多任務(wù)學(xué)習(xí)中表現(xiàn)更佳。與之前關(guān)閉輔助損失函數(shù)的做法相反,實(shí)際上這種損失函數(shù)可以幫助防止過(guò)擬合。
與稠密模型相比,稀疏模型從指令微調(diào)中受益更多。該圖來(lái)自 MoEs Meets instructions Tuning 論文
稀疏 VS 稠密,如何選擇?
稀疏混合專家模型 (MoE) 適用于擁有多臺(tái)機(jī)器且要求高吞吐量的場(chǎng)景。在固定的預(yù)訓(xùn)練計(jì)算資源下,稀疏模型往往能夠?qū)崿F(xiàn)更優(yōu)的效果。相反,在顯存較少且吞吐量要求不高的場(chǎng)景,稠密模型則是更合適的選擇。
注意: 直接比較稀疏模型和稠密模型的參數(shù)數(shù)量是不恰當(dāng)?shù)模驗(yàn)檫@兩類模型基于的概念和參數(shù)量的計(jì)算方法完全不同。
讓 MoE 起飛
最初的混合專家模型 (MoE) 設(shè)計(jì)采用了分支結(jié)構(gòu),這導(dǎo)致了計(jì)算效率低下。這種低效主要是因?yàn)?GPU 并不是為處理這種結(jié)構(gòu)而設(shè)計(jì)的,而且由于設(shè)備間需要傳遞數(shù)據(jù),網(wǎng)絡(luò)帶寬常常成為性能瓶頸。在接下來(lái)的討論中,我們會(huì)討論一些現(xiàn)有的研究成果,旨在使這些模型在預(yù)訓(xùn)練和推理階段更加高效和實(shí)用。我們來(lái)看看如何優(yōu)化 MoE 模型,讓 MoE 起飛。
并行計(jì)算
讓我們簡(jiǎn)要回顧一下并行計(jì)算的幾種形式:
數(shù)據(jù)并行: 相同的權(quán)重在所有節(jié)點(diǎn)上復(fù)制,數(shù)據(jù)在節(jié)點(diǎn)之間分割。
模型并行: 模型在節(jié)點(diǎn)之間分割,相同的數(shù)據(jù)在所有節(jié)點(diǎn)上復(fù)制。
模型和數(shù)據(jù)并行: 我們可以在節(jié)點(diǎn)之間同時(shí)分割模型和數(shù)據(jù)。注意,不同的節(jié)點(diǎn)處理不同批次的數(shù)據(jù)。
專家并行: 專家被放置在不同的節(jié)點(diǎn)上。如果與數(shù)據(jù)并行結(jié)合,每個(gè)節(jié)點(diǎn)擁有不同的專家,數(shù)據(jù)在所有節(jié)點(diǎn)之間分割。
在專家并行中,專家被放置在不同的節(jié)點(diǎn)上,每個(gè)節(jié)點(diǎn)處理不同批次的訓(xùn)練樣本。對(duì)于非 MoE 層,專家并行的行為與數(shù)據(jù)并行相同。對(duì)于 MoE 層,序列中的令牌被發(fā)送到擁有所需專家的節(jié)點(diǎn)。
Switch Transformers 論文中展示如何使用不同的并行技術(shù)在節(jié)點(diǎn)上分割數(shù)據(jù)和模型的插圖
容量因子和通信開銷
提高容量因子 (Capacity Factor, CF) 可以增強(qiáng)模型的性能,但這也意味著更高的通信成本和對(duì)保存激活值的顯存的需求。在設(shè)備通信帶寬有限的情況下,選擇較小的容量因子可能是更佳的策略。一個(gè)合理的初始設(shè)置是采用 Top-2 路由、1.25 的容量因子,同時(shí)每個(gè)節(jié)點(diǎn)配置一個(gè)專家。在評(píng)估性能時(shí),應(yīng)根據(jù)需要調(diào)整容量因子,以在設(shè)備間的通信成本和計(jì)算成本之間找到一個(gè)平衡點(diǎn)。
部署技術(shù)
您可以在 Inference Endpoints 部署 mistralai/Mixtral-8x7B-Instruct-v0.1。
部署混合專家模型 (MoE) 的一個(gè)關(guān)鍵挑戰(zhàn)是其龐大的參數(shù)規(guī)模。對(duì)于本地使用情況,我們可能希望使用更小的模型。為了使模型更適合部署,下面是幾種有用的技術(shù):
預(yù)先蒸餾實(shí)驗(yàn): Switch Transformers 的研究者們進(jìn)行了預(yù)先蒸餾的實(shí)驗(yàn)。他們通過(guò)將 MoE 模型蒸餾回其對(duì)應(yīng)的稠密模型,成功保留了 30-40%的由稀疏性帶來(lái)的性能提升。預(yù)先蒸餾不僅加快了預(yù)訓(xùn)練速度,還使得在推理中使用更小型的模型成為可能。
任務(wù)級(jí)別路由: 最新的方法中,路由器被修改為將整個(gè)句子或任務(wù)直接路由到一個(gè)專家。這樣做可以提取出一個(gè)用于服務(wù)的子網(wǎng)絡(luò),有助于簡(jiǎn)化模型的結(jié)構(gòu)。
專家網(wǎng)絡(luò)聚合: 這項(xiàng)技術(shù)通過(guò)合并各個(gè)專家的權(quán)重,在推理時(shí)減少了所需的參數(shù)數(shù)量。這樣可以在不顯著犧牲性能的情況下降低模型的復(fù)雜度。
高效訓(xùn)練
FasterMoE (2022 年 3 月) 深入分析了 MoE 在不同并行策略下的理論性能極限,并且探索了一系列創(chuàng)新技術(shù),包括用于專家權(quán)重調(diào)整的方法、減少延遲的細(xì)粒度通信調(diào)度技術(shù),以及一個(gè)基于最低延遲進(jìn)行專家選擇的拓?fù)涓兄T控機(jī)制。這些技術(shù)的結(jié)合使得 MoE 運(yùn)行速度提升高達(dá) 17 倍。
Megablocks (2022 年 11 月) 則專注于通過(guò)開發(fā)新的 GPU kernel 來(lái)處理 MoE 模型中的動(dòng)態(tài)性,以實(shí)現(xiàn)更高效的稀疏預(yù)訓(xùn)練。其核心優(yōu)勢(shì)在于,它不會(huì)丟棄任何令牌,并能高效地適應(yīng)現(xiàn)代硬件架構(gòu) (支持塊稀疏矩陣乘),從而達(dá)到顯著的加速效果。Megablocks 的創(chuàng)新之處在于,它不像傳統(tǒng) MoE 那樣使用批量矩陣乘法 (這通常假設(shè)所有專家形狀相同且處理相同數(shù)量的令牌),而是將 MoE 層表示為塊稀疏操作,可以靈活適應(yīng)不均衡的令牌分配。
針對(duì)不同規(guī)模的專家和令牌數(shù)量的塊稀疏矩陣乘法。該圖來(lái)自 MegaBlocks 論文
開源混合專家模型
目前,下面這些開源項(xiàng)目可以用于訓(xùn)練混合專家模型 (MoE):
Megablocks: https://github.com/stanford-futuredata/megablocks
Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
OpenMoE: https://github.com/XueFuzhao/OpenMoE
對(duì)于開源的混合專家模型 (MoE),你可以關(guān)注下面這些:
Switch Transformers (Google): 基于 T5 的 MoE 集合,專家數(shù)量從 8 名到 2048 名。最大的模型有 1.6 萬(wàn)億個(gè)參數(shù)。
NLLB MoE (Meta): NLLB 翻譯模型的一個(gè) MoE 變體。
OpenMoE: 社區(qū)對(duì)基于 Llama 的模型的 MoE 嘗試。
Mixtral 8x7B (Mistral): 一個(gè)性能超越了 Llama 2 70B 的高質(zhì)量混合專家模型,并且具有更快的推理速度。此外,還發(fā)布了一個(gè)經(jīng)過(guò)指令微調(diào)的模型。有關(guān)更多信息,可以在 Mistral 的 公告博客文章 中了解。
一些有趣的研究方向
首先是嘗試將稀疏混合專家模型 (SMoE)蒸餾回到具有更少實(shí)際參數(shù)但相似等價(jià)參數(shù)量的稠密模型。
MoE 的量化也是一個(gè)有趣的研究領(lǐng)域。例如,QMoE (2023 年 10 月) 通過(guò)將 MoE 量化到每個(gè)參數(shù)不到 1 位,將 1.6 萬(wàn)億參數(shù)的 Switch Transformer 所需的存儲(chǔ)從 3.2TB 壓縮到僅 160GB。
簡(jiǎn)而言之,一些值得探索的有趣領(lǐng)域包括:
將 Mixtral 蒸餾成一個(gè)稠密模型。
探索合并專家模型的技術(shù)及其對(duì)推理時(shí)間的影響。
嘗試對(duì) Mixtral 進(jìn)行極端量化的實(shí)驗(yàn)。
-
人工智能
+關(guān)注
關(guān)注
1791文章
47183瀏覽量
238266 -
開源
+關(guān)注
關(guān)注
3文章
3310瀏覽量
42473 -
模型
+關(guān)注
關(guān)注
1文章
3226瀏覽量
48809
原文標(biāo)題:混合專家模型 (MoE) 詳解
文章出處:【微信號(hào):OSC開源社區(qū),微信公眾號(hào):OSC開源社區(qū)】歡迎添加關(guān)注!文章轉(zhuǎn)載請(qǐng)注明出處。
發(fā)布評(píng)論請(qǐng)先 登錄
相關(guān)推薦
評(píng)論