隨著GPT-4的架構(gòu)被知名業(yè)內(nèi)大佬「開源」，混合專家架構(gòu)（MoE）再次成為了研究的重點(diǎn)。

GPT-4擁有16個(gè)專家模型，總共包含1.8萬億個(gè)參數(shù)。每生成一個(gè)token需要使用大約2800億參數(shù)和560TFLOPs。

然而，模型更快、更準(zhǔn)確的代價(jià)，則是巨大的參數(shù)量，和隨之而來的高昂成本。

比如，1.6萬億參數(shù)的SwitchTransformer-c2048模型，需要3.2TB的GPU顯存才能有效運(yùn)行。

為了解決這一問題，來自奧地利科技學(xué)院（ISTA）的研究人員提出了一種全新的壓縮和執(zhí)行框架——QMoE。

論文地址：https://arxiv.org/abs/2310.16795

通過采用專門設(shè)計(jì)的GPU解碼內(nèi)核，QMoE具備了高效的端到端壓縮推理——不僅可以實(shí)現(xiàn)高達(dá)20倍的壓縮率，而且只會(huì)產(chǎn)生輕微的精度損失。

具體而言，QMoE僅需單個(gè)GPU服務(wù)器，就可以在一天內(nèi)將1.6萬億參數(shù)的SwitchTransformer-c2048模型壓縮至不到160GB，相當(dāng)于每參數(shù)只有0.8位。

如此一來，就可以在4張英偉達(dá)RTX A6000或8張英偉達(dá)RTX 3090 GPU上運(yùn)行，而推理時(shí)的開銷還不到未壓縮模型的5%。

MoE模型量化

混合模型（MoE）的核心理念是通過增加模型參數(shù)量，來提高網(wǎng)絡(luò)的建模能力，同時(shí)與標(biāo)準(zhǔn)的前饋架構(gòu)相比，保持計(jì)算成本幾乎不變。

由于處理每個(gè)輸入token時(shí)僅需調(diào)用網(wǎng)絡(luò)中的一小部分，因此這種設(shè)計(jì)可以利用100個(gè)甚至1000個(gè)「專家」來構(gòu)建超大規(guī)模的模型，并進(jìn)行高效的訓(xùn)練和推理。

事實(shí)證明，在推理速度相當(dāng)?shù)那闆r下，MoE可以大幅提高準(zhǔn)確率和訓(xùn)練速度。但如此龐大的體積，也就意味著需要大量的顯存才能讓模型跑起來。

壓縮MoE的一個(gè)主要挑戰(zhàn)是需要維持龐大的激活集。

對(duì)此，可以通過精心安排模型執(zhí)行的方式，將需要計(jì)算的中間數(shù)據(jù)控制在一小部分。從而把主存儲(chǔ)從GPU卸載到價(jià)格更便宜、數(shù)量更多的CPU內(nèi)存中。

具體來說就是，維持一個(gè)大型緩沖區(qū)B，并按照以下步驟對(duì)Transformer塊的稠密部分進(jìn)行更新：

1. 從CPU到GPU，抓取一個(gè)包含有幾百個(gè)token的「樣本」X；

2. 通過對(duì)應(yīng)的稠密層，得到結(jié)果Y；

3. 計(jì)算并存儲(chǔ)Y中token的專家分配；

4. 將Y發(fā)送回CPU并覆蓋B中的X。

對(duì)于稀疏部分：

1. 從CPU到GPU，抓取B中所有已分配給專家E的token，用X_E表示。

2. 利用它們生成壓縮的專家E'（例如，使用GPTQ）。

3. 通過E'運(yùn)行X_E，得到Y(jié)_E'。

4. 將Y_E'送回CPU并覆蓋B中的X_E。

如圖2所示，這個(gè)過程最小化了內(nèi)存消耗和傳輸成本：只需一個(gè)B的副本，每個(gè)token在每個(gè)Transformer塊中只被讀寫了兩次。

更進(jìn)一步的，研究人員設(shè)計(jì)了一個(gè)編碼方案和一個(gè)CUDA內(nèi)核，實(shí)現(xiàn)了每權(quán)重低于1位的壓縮，并將推理的GPU執(zhí)行開銷降至最低。

壓縮效果

精度

首先，研究人員將所有SwitchTransformer模型量化到2位和三元精度，然后評(píng)估其驗(yàn)證損失。

對(duì)于128個(gè)專家，默認(rèn)的校準(zhǔn)樣本數(shù)為10K；對(duì)于2048個(gè)專家，默認(rèn)的校準(zhǔn)樣本數(shù)為160K。同時(shí)，研究人員也測(cè)試了0.5倍和2倍的樣本數(shù)。

結(jié)果顯示，使用數(shù)據(jù)依賴的量化，2位模型可以在最小的損失下實(shí)現(xiàn)（相對(duì)于c2048，損失為1.7%），而三元精度下的損失增加也很?。ㄏ鄬?duì)于c2048，損失為6.7%）。

這不僅證明了所提出的先進(jìn)量化方法的有效性，而且還表明極低位寬的壓縮確實(shí)適用于大規(guī)模的MoE。

此外，研究人員還在來自RedPajama的arXiv、GitHub、StackExchange和Wikipedia的數(shù)據(jù)上進(jìn)行了評(píng)估。

雖然校準(zhǔn)數(shù)據(jù)中只有<0.01%來自這些網(wǎng)站，但壓縮后的模型依然保持了幾乎與核心分布相同的性能。

就校準(zhǔn)數(shù)據(jù)而言，增加樣本數(shù)量通常會(huì)略微提高性能，在三元量化時(shí)最為明顯。但在此過程中也會(huì)出現(xiàn)一些噪聲，尤其是在2位時(shí)。

壓縮

測(cè)試中，研究人員同時(shí)考慮了僅MoE模塊的壓縮，以及相對(duì)于整個(gè)模型及其所有元數(shù)據(jù)的壓縮。

僅MoE本身的而言，所有規(guī)模都實(shí)現(xiàn)了>16倍的壓縮率，相當(dāng)于每個(gè)參數(shù)的存儲(chǔ)空間都<1位。

在c2048上，即使是包括所有未壓縮的稠密層在內(nèi)，整體的壓縮率也達(dá)到了19.81倍，相當(dāng)于每個(gè)參數(shù)0.807位，從而將檢查點(diǎn)大小從3142GB減少到158.6GB。

此外，還可以觀察到壓縮率隨模型大小的增加而增加，這有兩個(gè)原因：

（a）自然稀疏性增加，且研究人員針對(duì)c2048優(yōu)化了編碼字典；

（b）層越大，權(quán)重分布越接近獨(dú)立。

運(yùn)行時(shí)間

最后，研究人員評(píng)估了針對(duì)不同數(shù)量的校準(zhǔn)數(shù)據(jù)，在單個(gè)A6000 GPU上生成壓縮模型所需的時(shí)間。

結(jié)果顯示，較小的模型可以在一小時(shí)內(nèi)壓縮完成，即便是c2048也能在不到一天的時(shí)間內(nèi)完成，這證實(shí)了QMoE的高效性。

從large128到c2048，運(yùn)行時(shí)間的增加與大小的差異基本成正比，盡管后者使用了多16倍的樣本。這是因?yàn)槊總€(gè)專家的樣本數(shù)量保持不變，而專家規(guī)模僅略有增加。

運(yùn)行結(jié)果

首先，將壓縮的矩陣-向量積內(nèi)核與PyTorch標(biāo)準(zhǔn)的（未壓縮）bfloat16 cuBLAS內(nèi)核進(jìn)行直接（孤立）比較。

圖 5（左）顯示了壓縮內(nèi)核與bfloat16內(nèi)核相比，在兩款不同的GPU上，MoE發(fā)現(xiàn)矩陣形狀所耗費(fèi)的時(shí)間。

雖然研究人員使用的儲(chǔ)存性能較差，但執(zhí)行壓縮內(nèi)核所需的時(shí)間，依然比接近理想的bfloat16基線少。在特定矩陣形狀下，速度最多可提高35%。

而這些操作的延遲也非常低，其中，最小的矩陣耗時(shí)<0.02毫秒，最大的耗時(shí)<0.05毫秒。

隨后，研究人員在HuggingFace中，利用壓縮MoE模型的實(shí)際權(quán)重，對(duì)內(nèi)核進(jìn)行了端到端的基準(zhǔn)測(cè)試。

結(jié)果如圖5（右）所示，壓縮模型的端到端執(zhí)行速度只比標(biāo)準(zhǔn)（未壓縮）的慢了<5%。

盡管每層時(shí)序更快，但速度仍略有下降，這是因?yàn)榫幋a器有時(shí)會(huì)將多個(gè)token路由到同一個(gè)專家。

目前的實(shí)現(xiàn)方式是，對(duì)每個(gè)token執(zhí)行單獨(dú)的矩陣向量乘積，而基線執(zhí)行的是更高效的聯(lián)合矩陣乘法。

在一些應(yīng)用中，這是一個(gè)很大的瓶頸。對(duì)此，可以在內(nèi)核中引入token內(nèi)循環(huán)，或者在token數(shù)量較多的情況下，先進(jìn)行完全解壓縮，然后再執(zhí)行標(biāo)準(zhǔn)的矩陣乘法。

討論與局限性

總結(jié)而言， QMoE是一個(gè)開源的端到端壓縮和推理框架，用于解決MoE在推理過程中，內(nèi)存開銷過大的問題。

研究人員首次證明了，像SwitchTransformer c2048這樣的萬億參數(shù)模型，可以精確壓縮到每個(gè)參數(shù)小于1位，壓縮率接近20倍。并且，首次在單個(gè)消費(fèi)級(jí)GPU服務(wù)器上，實(shí)現(xiàn)了此類模型的高效端到端執(zhí)行。

不過，由于只有少數(shù)大規(guī)模且精確的MoE可以被公開獲得，因此研究的模型集十分有限。

此外，由于其規(guī)模龐大，大多數(shù)MoE都是在不同的定制框架中訓(xùn)練和部署的，這就需要復(fù)雜的手動(dòng)集成才能用于進(jìn)一步研究。

盡管如此，研究人員還是涵蓋了一些規(guī)模最大、精度最高的MoE，特別是SwitchTransformer。