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

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

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

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

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

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

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

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

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

MoE模型量化

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

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

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

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

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

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

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

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

3. 計算并存儲Y中token的專家分配；

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

對于稀疏部分：

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

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

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

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

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

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

壓縮效果

精度

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

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

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

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

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

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

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

壓縮

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

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

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

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

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

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

運行時間

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

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

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

運行結(jié)果

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

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

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

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

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

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

盡管每層時序更快，但速度仍略有下降，這是因為編碼器有時會將多個token路由到同一個專家。

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

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

討論與局限性

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

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

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

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

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