本文來自南京大學(xué)計(jì)算機(jī)學(xué)院軟件研究所徐經(jīng)緯DeepEngine團(tuán)隊(duì)，作者為徐經(jīng)緯準(zhǔn)聘助理教授、碩士生賴俊宇和黃云鵬。目前該論文已被 ICLR2025 接收。

在大語言模型領(lǐng)域中，預(yù)訓(xùn)練 + 微調(diào)范式已經(jīng)成為了部署各類下游應(yīng)用的重要基礎(chǔ)。在該框架下，通過使用搭低秩自適應(yīng)（LoRA）方法的大模型參數(shù)高效微調(diào)（PEFT）技術(shù)，已經(jīng)產(chǎn)生了大量針對(duì)特定任務(wù)、可重用的 LoRA 適配器。但使用 LoRA 適配器的微調(diào)方法需要明確的意圖選擇，因此在搭載多個(gè) LoRA 適配器的單一大語言模型上，自主任務(wù)感知和切換方面一直存在挑戰(zhàn)。

為此，團(tuán)隊(duì)提出了一個(gè)可擴(kuò)展、高效的多任務(wù)嵌入架構(gòu) MeteoRA。該框架通過全模式混合專家模型（MoE）將多個(gè)特定任務(wù)的 LoRA 適配器重用到了基礎(chǔ)模型上。該框架還包括了一個(gè)新穎的混合專家模型前向加速策略，以改善傳統(tǒng)實(shí)現(xiàn)的效率瓶頸。配備了 MeteoRA 框架的大語言模型在處理復(fù)合問題時(shí)取得了卓越的性能，可以在一次推理中高效地解決十個(gè)按次序輸入的不同問題，證明了該模型具備適配器及時(shí)切換的強(qiáng)化能力。

• 論文標(biāo)題：MeteoRA: Multiple-tasks Embedded LoRA for Large Language Models
• 論文地址：https://arxiv.org/abs/2405.13053
• 項(xiàng)目主頁(yè)：https://github.com/NJUDeepEngine/meteora

該工作的創(chuàng)新點(diǎn)主要有以下部分：

• 可擴(kuò)展的 LoRA 集成框架：MeteoRA 框架能夠整合現(xiàn)有的 LoRA 適配器，并提供了大語言模型自主按需選擇和切換不同 LoRA 的能力。
• 混合專家模型的前向加速策略：揭示了混合專家模型的效率問題，使用新的 GPU 內(nèi)核操作實(shí)現(xiàn)了前向傳播加速策略，在保持內(nèi)存開銷不變的同時(shí)實(shí)現(xiàn)了平均約 4 倍的加速。
• 卓越的模型性能：評(píng)估表明使用 MeteoRA 框架的大語言模型在復(fù)合任務(wù)中表現(xiàn)出了卓越的性能，因而增強(qiáng)了大語言模型在結(jié)合使用現(xiàn)成 LoRA 適配器上的實(shí)際效果。

相關(guān)工作

低秩適應(yīng) (LoRA)：低秩適應(yīng) [1] 提供了一種策略來減少下游任務(wù)微調(diào)所需要的可訓(xùn)練參數(shù)規(guī)模。對(duì)于一個(gè)基于 Transformer 的大語言模型，LoRA 會(huì)向每一個(gè)基本線性層的權(quán)重矩陣注入兩個(gè)可訓(xùn)練的低秩矩陣，并用兩個(gè)低秩矩陣的相乘來代表模型在原來權(quán)重矩陣上的微調(diào)。LoRA 可以用于 Transformer 中自注意力模塊和多層感知機(jī)模塊的 7 種權(quán)重矩陣，有效縮減了微調(diào)權(quán)重的規(guī)模。應(yīng)用低秩適應(yīng)技術(shù)可以在不改變預(yù)訓(xùn)練模型參數(shù)的前提下，使用自回歸語言模型的優(yōu)化目標(biāo)來訓(xùn)練 LoRA 適配器的參數(shù)。

多任務(wù) LoRA 融合：LoRA 適配器通常被微調(diào)來完成特定的下游任務(wù)。為了增強(qiáng)大語言模型處理多種任務(wù)的能力，主流的做法有兩種。一種方法是從多種任務(wù)中合并數(shù)據(jù)集，在此基礎(chǔ)上訓(xùn)練單一的 LoRA 適配器，但相關(guān)研究已經(jīng)證明同時(shí)學(xué)習(xí)不同領(lǐng)域的知識(shí)很困難。[2] 另一種方法是直接將多種 LoRA 適配器整合到一個(gè)模型中，如 PEFT [3]、S-LoRA [4]。但是流行的框架必須要明確指定注入的 LoRA 適配器，因此模型缺乏自主選擇和及時(shí)切換 LoRA 的能力。現(xiàn)有的工作如 LoRAHub [5] 可以在不人為指定的情況下完成 LoRA 適配器的選擇，但仍舊需要針對(duì)每一個(gè)下游任務(wù)進(jìn)行少量的學(xué)習(xí)。

混合專家模型（MoE）：混合專家模型是一種通過組合多個(gè)模型的預(yù)測(cè)結(jié)果來提高效率和性能的機(jī)器學(xué)習(xí)范式，該范式通過門控網(wǎng)絡(luò)將輸入動(dòng)態(tài)分配給最相關(guān)的 “專家” 來獲得預(yù)測(cè)結(jié)果。[6] 該模型利用來自不同專家的特定知識(shí)來改善在多樣、復(fù)雜的問題上的總體表現(xiàn)，已有研究證明 MoE 在大規(guī)模神經(jīng)網(wǎng)絡(luò)上的有效性 [7]。對(duì)于每個(gè)輸入，MoE 只會(huì)激活一部分專家，從而在不損害模型規(guī)模的基礎(chǔ)上顯著提高計(jì)算效率。該方法已被證實(shí)在擴(kuò)展基于 Transformer 的應(yīng)用架構(gòu)上十分有效，如 Mixtral [8]。

方法描述

該工作提出了一個(gè)可擴(kuò)展、高效的多任務(wù)嵌入架構(gòu) MeteoRA，它能夠直接使用開源社區(qū)中或面向特定下游任務(wù)微調(diào)好的 LoRA 適配器，并依據(jù)輸入自主選擇、及時(shí)切換合適的 LoRA 適配器。如圖 1 所示，MeteoRA 模塊可以被集成到注意力模塊和多層感知機(jī)模塊的所有基本線性層中，每一個(gè)模塊都包括了一系列低秩矩陣。通過 MoE 前向加速策略，大語言模型能夠高效解決廣泛的問題。

圖 1：集成了 MeteoRA 模塊實(shí)現(xiàn) MoE 架構(gòu)的大語言模型框架

圖 2 展示了 MeteoRA 模塊的內(nèi)部結(jié)構(gòu)。MeteoRA 模塊中嵌入了多個(gè) LoRA 適配器，并通過一個(gè)門控網(wǎng)絡(luò)來實(shí)現(xiàn) MoE 架構(gòu)。該門控網(wǎng)絡(luò)會(huì)依據(jù)輸入選擇 top-k 個(gè) LoRA 適配器，并將它們組合作為微調(diào)權(quán)重進(jìn)行前向傳播。通過這種架構(gòu)，門控網(wǎng)絡(luò)會(huì)執(zhí)行路由策略，從而根據(jù)輸入選取合適的 LoRA 適配器。每一個(gè) MeteoRA 模塊都包含一個(gè)獨(dú)立的門控網(wǎng)絡(luò)，不同門控網(wǎng)絡(luò)依據(jù)它們的輸入獨(dú)立決策，并在全部解碼器模塊的前向傳播過程中動(dòng)態(tài)選取不同 LoRA 適配器。

圖 2：應(yīng)用 MoE 模型集成 LoRA 嵌入的 MeteoRA 模塊的架構(gòu)

MeteoRA 模塊的訓(xùn)練遵循自回歸語言建模任務(wù)下的模型微調(diào)準(zhǔn)則，訓(xùn)練中需要保持基本的大語言模型權(quán)重和預(yù)訓(xùn)練的 LoRA 適配器參數(shù)不變。由于 MeteoRA 模塊支持選擇權(quán)重最高的若干個(gè) LoRA 適配器，團(tuán)隊(duì)引入了一種將自回歸語言建模損失和所有門控網(wǎng)絡(luò)損失組合起來的聯(lián)合優(yōu)化方案。該優(yōu)化函數(shù)綜合了自回歸模型中的預(yù)測(cè)損失和 MeteoRA 模塊中 LoRA 分類的交叉熵?fù)p失，實(shí)現(xiàn)門控網(wǎng)絡(luò)的訓(xùn)練。

MeteoRA 模塊的核心組件是整合了多個(gè) LoRA 適配器的 MoE 架構(gòu)。首先，團(tuán)隊(duì)將 LoRA 整合成在 HBM 上連續(xù)分配的張量。對(duì)于每個(gè)批次的輸入序列中的每個(gè)標(biāo)記，由于該模塊需要利用門控網(wǎng)絡(luò)找到 LoRA 適配器的編號(hào)索引集合，MeteoRA 模塊幾乎不可能與單獨(dú)的 LoRA 模塊保持相同的效率。基于原始循環(huán)（loop-original）的簡(jiǎn)單實(shí)現(xiàn) [8] 采用 for-loop 遍歷 LoRA 適配器，在每次遍歷中對(duì)該適配器的候選集合應(yīng)用 LoRA 前向傳遞。該方法簡(jiǎn)單地將所有批次的所有標(biāo)記拆成 LoRA 數(shù)量個(gè)集合，并且使得每個(gè)集合按順序傳遞。然而，考慮到輸入標(biāo)記相互獨(dú)立的性質(zhì)，這種方法無法充分利用并行化 GEMM 算子 [9] 加速，尤其是當(dāng)某些 LoRA 適配器僅被少數(shù)標(biāo)記選擇，或是待推理標(biāo)記小于 LoRA 數(shù)量時(shí)，實(shí)驗(yàn)中可能花費(fèi)最多 10 倍的運(yùn)行時(shí)間。

該工作采用了前向傳播加速策略 bmm-torch，直接索引全部批次標(biāo)記的 top-k 適配器，這會(huì)利用到兩次 bmm 計(jì)算。相比于原始的循環(huán)方法，bmm-torch 基于 PyTorch 提供的 bmm 操作符 [10] 并行化所有批次標(biāo)記的 top-k 個(gè)適配器，實(shí)現(xiàn)了約 4 倍的加速。在實(shí)驗(yàn)中該策略僅比它的上限單 LoRA 模塊慢 2.5 倍。由于 PyTorch 的索引約束（indexing constraints）[11]，bmm-torch 需要分配更大的內(nèi)存執(zhí)行批處理，當(dāng)批次或者輸入長(zhǎng)度過大時(shí) bmm-torch 的大內(nèi)存開銷可能成為瓶頸。為此，該工作使用了 Triton [12] 開發(fā)了自定義 GPU 內(nèi)核算子，不僅保持了 bmm-torch 約 80% 的運(yùn)算效率，而且還保持了原始循環(huán)級(jí)別的低內(nèi)存開銷。

實(shí)驗(yàn)結(jié)果

該工作在獨(dú)立任務(wù)和復(fù)合任務(wù)上，對(duì)所提出的架構(gòu)和設(shè)計(jì)進(jìn)行了廣泛的實(shí)驗(yàn)驗(yàn)證。實(shí)驗(yàn)使用了兩種知名的大語言模型 LlaMA2-13B 和 LlaMA3-8B，圖 3 顯示了集成 28 項(xiàng)任務(wù)的 MeteoRA 模塊與單 LoRA、參考模型 PEFT 在各項(xiàng)獨(dú)立任務(wù)上的預(yù)測(cè)表現(xiàn)雷達(dá)圖。評(píng)估結(jié)果表明，無論使用哪種大語言模型，MeteoRA 模型都具備與 PEFT 相近的性能，而 MeteoRA 不需要顯式地激活特定的 LoRA。表 1 展示了所有方法在各項(xiàng)任務(wù)上的平均分?jǐn)?shù)。

圖 3：MeteoRA 模型在 28 項(xiàng)選定任務(wù)上的評(píng)估表現(xiàn)

表 1：各模型在 28 項(xiàng)選定任務(wù)上的平均表現(xiàn)

為了驗(yàn)證該模型按順序解決復(fù)合問題的能力，該工作通過串行連接獨(dú)立的任務(wù)來構(gòu)建 3 個(gè)數(shù)據(jù)集，分別整合了 3、5、10 項(xiàng)任務(wù)，期望模型能夠按次序解決同一序列中輸入的不同類別問題。實(shí)驗(yàn)結(jié)果如表 2 所示，可見隨著復(fù)合任務(wù)數(shù)的提升，MeteoRA 模塊幾乎全面優(yōu)于參考 LoRA-B 模型。

表 2：復(fù)合任務(wù)的評(píng)估結(jié)果，左側(cè)為 MeteoRA 模塊，右側(cè)為 LoRA-B

為了更進(jìn)一步驗(yàn)證門控網(wǎng)絡(luò)在 MeteoRA 模塊中的功能，該工作展示了在復(fù)合任務(wù)推理過程中采用 top-2 策略的 LoRA 選擇模式。在兩個(gè)相鄰任務(wù)的連接處，門控網(wǎng)絡(luò)正確執(zhí)行了 LoRA 的切換操作。

圖 4：在 3 項(xiàng)任務(wù)復(fù)合中 LoRA 選取情況的例子

為了驗(yàn)證采用自定義 GPU 算子的前向傳播設(shè)計(jì)的運(yùn)算效率，該工作在 28 項(xiàng)任務(wù)上截取了部分樣本，將新的前向傳播策略與其上限和原始實(shí)現(xiàn)對(duì)比。評(píng)估結(jié)果如圖 5 所示，展現(xiàn)了新的加速策略卓越的性能。

圖 5：四種不同的前向傳播策略在 28 項(xiàng)任務(wù)上的整體運(yùn)行時(shí)間