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

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

打開網(wǎng)易新聞 查看精彩圖片

論文標題： MeteoRA: Multiple-tasks Embedded LoRA for Large Language Models 論文地址： https://arxiv.org/abs/2405.13053 項目主頁： https://github.com/NJUDeepEngine/meteora

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

• 可擴展的 LoRA 集成框架：MeteoRA 框架能夠整合現(xiàn)有的 LoRA 適配器，并提供了大語言模型自主按需選擇和切換不同 LoRA 的能力。

• 混合專家模型的前向加速策略：揭示了混合專家模型的效率問題，使用新的 GPU 內(nèi)核操作實現(xiàn)了前向傳播加速策略，在保持內(nèi)存開銷不變的同時實現(xiàn)了平均約 4 倍的加速。

• 卓越的模型性能：評估表明使用 MeteoRA 框架的大語言模型在復合任務中表現(xiàn)出了卓越的性能，因而增強了大語言模型在結(jié)合使用現(xiàn)成 LoRA 適配器上的實際效果。

一、相關工作

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

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

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

二、方法描述

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

打開網(wǎng)易新聞 查看精彩圖片

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

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

打開網(wǎng)易新聞 查看精彩圖片

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

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

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

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

三、實驗結(jié)果

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

打開網(wǎng)易新聞 查看精彩圖片

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

打開網(wǎng)易新聞 查看精彩圖片

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

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

打開網(wǎng)易新聞 查看精彩圖片

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

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

打開網(wǎng)易新聞 查看精彩圖片

圖 4：在 3 項任務復合中 LoRA 選取情況的例子

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

打開網(wǎng)易新聞 查看精彩圖片

圖 5：四種不同的前向傳播策略在 28 項任務上的整體運行時間
參考文獻

[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)緯、賴俊宇、黃云鵬 來源：公眾號【機器之心】

llustration From IconScout By IconScout Store

-The End-

掃碼觀看！

本周上新！

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

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

www.techbeat.net

投稿內(nèi)容

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

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

投稿須知

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

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

投稿方式

發(fā)送郵件到

melodybai@thejiangmen.com

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

關于我“門”

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

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

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

bp@thejiangmen.com

打開網(wǎng)易新聞 查看精彩圖片

點擊右上角，把文章分享到朋友圈