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

為此，南京大學(xué)計(jì)算機(jī)學(xué)院軟件研究所的DeepEngine團(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)目主頁： 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 倍的加速。

• 卓越的模型性能：評估表明使用 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ī)模。對于一個(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 適配器的選擇，但仍舊需要針對每一個(gè)下游任務(wù)進(jìn)行少量的學(xué)習(xí)。

混合專家模型（MoE）：混合專家模型是一種通過組合多個(gè)模型的預(yù)測結(jié)果來提高效率和性能的機(jī)器學(xué)習(xí)范式，該范式通過門控網(wǎng)絡(luò)將輸入動(dòng)態(tài)分配給最相關(guān)的 “專家” 來獲得預(yù)測結(jié)果。[6] 該模型利用來自不同專家的特定知識(shí)來改善在多樣、復(fù)雜的問題上的總體表現(xiàn)，已有研究證明 MoE 在大規(guī)模神經(jīng)網(wǎng)絡(luò)上的有效性 [7]。對于每個(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ù)測損失和 MeteoRA 模塊中 LoRA 分類的交叉熵?fù)p失，實(shí)現(xiàn)門控網(wǎng)絡(luò)的訓(xùn)練。

MeteoRA 模塊的核心組件是整合了多個(gè) LoRA 適配器的 MoE 架構(gòu)。首先，團(tuán)隊(duì)將 LoRA 整合成在 HBM 上連續(xù)分配的張量。對于每個(gè)批次的輸入序列中的每個(gè)標(biāo)記，由于該模塊需要利用門控網(wǎng)絡(luò)找到 LoRA 適配器的編號索引集合，MeteoRA 模塊幾乎不可能與單獨(dú)的 LoRA 模塊保持相同的效率。基于原始循環(huán)（loop-original）的簡單實(shí)現(xiàn) [8] 采用 for-loop 遍歷 LoRA 適配器，在每次遍歷中對該適配器的候選集合應(yīng)用 LoRA 前向傳遞。該方法簡單地將所有批次的所有標(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)批次或者輸入長度過大時(shí) bmm-torch 的大內(nèi)存開銷可能成為瓶頸。為此，該工作使用了 Triton [12] 開發(fā)了自定義 GPU 內(nèi)核算子，不僅保持了 bmm-torch 約 80% 的運(yùn)算效率，而且還保持了原始循環(huán)級別的低內(nèi)存開銷。

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

該工作在獨(dú)立任務(wù)和復(fù)合任務(wù)上，對所提出的架構(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ù)測表現(xiàn)雷達(dá)圖。評估結(jié)果表明，無論使用哪種大語言模型，MeteoRA 模型都具備與 PEFT 相近的性能，而 MeteoRA 不需要顯式地激活特定的 LoRA。表 1 展示了所有方法在各項(xiàng)任務(wù)上的平均分?jǐn)?shù)。

圖 3：MeteoRA 模型在 28 項(xiàng)選定任務(wù)上的評估表現(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ù)的評估結(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)對比。評估結(jié)果如圖 5 所示，展現(xiàn)了新的加速策略卓越的性能。

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

[1] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685, 2021.

[2] Chen Ling, Xujiang Zhao, Jiaying Lu, Chengyuan Deng, Can Zheng, Junxiang Wang, Tanmoy Chowdhury, Yun Li, Hejie Cui, Xuchao Zhang, Tianjiao Zhao, Amit Panalkar, Dhagash Mehta, Stefano Pasquali, Wei Cheng, Haoyu Wang, Yanchi Liu, Zhengzhang Chen, Haifeng Chen, Chris White, Quanquan Gu, Jian Pei, Carl Yang, and Liang Zhao. Domain specialization as the key to make large language models disruptive: A comprehensive survey, 2024.

[3] Sourab Mangrulkar, Sylvain Gugger, Lysandre Debut, Younes Belkada, Sayak Paul, and Benjamin Bossan. Peft: State-of-the-art parameter-efficient fine-tuning methods. https://github.com/huggingface/peft, 2022.

[4] Ying Sheng, Shiyi Cao, Dacheng Li, Coleman Hooper, Nicholas Lee, Shuo Yang, Christopher Chou, Banghua Zhu, Lianmin Zheng, Kurt Keutzer, et al. S-lora: Serving thousands of concurrent lora adapters. arXiv preprint arXiv:2311.03285, 2023.

[5] Chengsong Huang, Qian Liu, Bill Yuchen Lin, Tianyu Pang, Chao Du, and Min Lin. Lorahub: Efficient cross-task generalization via dynamic lora composition. arXiv preprint arXiv:2307.13269, 2023.

[6] Robert A Jacobs, Michael I Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive mixtures of local experts. Neural computation, 3 (1):79–87, 1991.

[7] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538, 2017.

[8] Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour, Guillaume Lample, Le?lio Renard Lavaud, Lucile Saulnier, MarieAnne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, The?ophile Gervet, Thibaut Lavril, Thomas Wang, Timothe?e Lacroix, and William El Sayed. Mixtral of experts, 2024.

[9] Shixun Wu, Yujia Zhai, Jinyang Liu, Jiajun Huang, Zizhe Jian, Bryan Wong, and Zizhong Chen. Anatomy of high-performance gemm with online fault tolerance on gpus. In Proceedings of the 37th International Conference on Supercomputing, pages 360–372, 2023b.

[10] PyTorch. torch.bmm — pytorch 2.3 documentation. https://pytorch.org/docs/stable/generated/torch.bmm.html, 2024a. Accessed: 2024-05-23.

[11] PyTorch. Tensor indexing api — pytorch documentation. https://pytorch.org/cppdocs/notes/tensor_indexing.html, 2024b. Accessed: 2024-05-23.

[12] Philippe Tillet, Hsiang-Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19, 2019.

作者：徐經(jīng)緯、賴俊宇、黃云鵬 來源：公眾號【機(jī)器之心】

llustration From IconScout By IconScout Store

-The End-

掃碼觀看！

本周上新！

“AI技術(shù)流”原創(chuàng)投稿計(jì)劃

TechBeat是由將門創(chuàng)投建立的AI學(xué)習(xí)社區(qū)（

www.techbeat.net

投稿內(nèi)容

// 最新技術(shù)解讀/系統(tǒng)性知識(shí)分享 //

// 前沿資訊解說/心得經(jīng)歷講述 //

投稿須知

稿件需要為原創(chuàng)文章，并標(biāo)明作者信息。

我們會(huì)選擇部分在深度技術(shù)解析及科研心得方向，對用戶啟發(fā)更大的文章，做原創(chuàng)性內(nèi)容獎(jiǎng)勵(lì)

投稿方式

發(fā)送郵件到

melodybai@thejiangmen.com

或添加工作人員微信（yellowsubbj）投稿，溝通投稿詳情；還可以關(guān)注“將門創(chuàng)投”公眾號，后臺(tái)回復(fù)“投稿”二字，獲得投稿說明。

關(guān)于我“門”

將門是一家以專注于數(shù)智核心科技領(lǐng)域的新型創(chuàng)投機(jī)構(gòu)，也是北京市標(biāo)桿型孵化器。 公司致力于通過連接技術(shù)與商業(yè)，發(fā)掘和培育具有全球影響力的科技創(chuàng)新企業(yè)，推動(dòng)企業(yè)創(chuàng)新發(fā)展與產(chǎn)業(yè)升級。

將門成立于2015年底，創(chuàng)始團(tuán)隊(duì)由微軟創(chuàng)投在中國的創(chuàng)始團(tuán)隊(duì)原班人馬構(gòu)建而成，曾為微軟優(yōu)選和深度孵化了126家創(chuàng)新的技術(shù)型創(chuàng)業(yè)公司。

如果您是技術(shù)領(lǐng)域的初創(chuàng)企業(yè)，不僅想獲得投資，還希望獲得一系列持續(xù)性、有價(jià)值的投后服務(wù)，歡迎發(fā)送或者推薦項(xiàng)目給我“門”:

bp@thejiangmen.com

點(diǎn)擊右上角，把文章分享到朋友圈