AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報道了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報道。投稿郵箱：[email protected]；[email protected]
本文來自南京大學計算機學院軟體研究所徐經緯DeepEngine團隊，作者為徐經緯準聘助理教授、碩士生賴俊宇和黃雲鵬。目前該論文已被 ICLR2025 接收。
在大語言模型領域中，預訓練 + 微調正規化已經成為了部署各類下游應用的重要基礎。在該框架下，透過使用搭低秩自適應（LoRA）方法的大模型引數高效微調（PEFT）技術，已經產生了大量針對特定任務、可重用的 LoRA 介面卡。但使用 LoRA 介面卡的微調方法需要明確的意圖選擇，因此在搭載多個 LoRA 介面卡的單一大語言模型上，自主任務感知和切換方面一直存在挑戰。
為此，團隊提出了一個可擴充套件、高效的多工嵌入架構 MeteoRA。該框架透過全模式混合專家模型（MoE）將多個特定任務的 LoRA 介面卡重用到了基礎模型上。該框架還包括了一個新穎的混合專家模型前向加速策略，以改善傳統實現的效率瓶頸。配備了 MeteoRA 框架的大語言模型在處理複合問題時取得了卓越的效能，可以在一次推理中高效地解決十個按次序輸入的不同問題，證明了該模型具備介面卡及時切換的強化能力。
-
論文標題：MeteoRA: Multiple-tasks Embedded LoRA for Large Language Models -
論文地址：https://arxiv.org/abs/2405.13053 -
專案主頁：https://github.com/NJUDeepEngine/meteora
該工作的創新點主要有以下部分：
-
可擴充套件的 LoRA 整合框架：MeteoRA 框架能夠整合現有的 LoRA 介面卡，並提供了大語言模型自主按需選擇和切換不同 LoRA 的能力。 -
混合專家模型的前向加速策略：揭示了混合專家模型的效率問題，使用新的 GPU 核心操作實現了前向傳播加速策略，在保持記憶體開銷不變的同時實現了平均約 4 倍的加速。 -
卓越的模型效能：評估表明使用 MeteoRA 框架的大語言模型在複合任務中表現出了卓越的效能，因而增強了大語言模型在結合使用現成 LoRA 介面卡上的實際效果。
相關工作
低秩適應 (LoRA)：低秩適應 [1] 提供了一種策略來減少下游任務微調所需要的可訓練引數規模。對於一個基於 Transformer 的大語言模型，LoRA 會向每一個基本線性層的權重矩陣注入兩個可訓練的低秩矩陣，並用兩個低秩矩陣的相乘來代表模型在原來權重矩陣上的微調。LoRA 可以用於 Transformer 中自注意力模組和多層感知機模組的 7 種權重矩陣，有效縮減了微調權重的規模。應用低秩適應技術可以在不改變預訓練模型引數的前提下，使用自迴歸語言模型的最佳化目標來訓練 LoRA 介面卡的引數。
多工 LoRA 融合：LoRA 介面卡通常被微調來完成特定的下游任務。為了增強大語言模型處理多種任務的能力，主流的做法有兩種。一種方法是從多種任務中合併資料集，在此基礎上訓練單一的 LoRA 介面卡，但相關研究已經證明同時學習不同領域的知識很困難。[2] 另一種方法是直接將多種 LoRA 介面卡整合到一個模型中，如 PEFT [3]、S-LoRA [4]。但是流行的框架必須要明確指定注入的 LoRA 介面卡，因此模型缺乏自主選擇和及時切換 LoRA 的能力。現有的工作如 LoRAHub [5] 可以在不人為指定的情況下完成 LoRA 介面卡的選擇，但仍舊需要針對每一個下游任務進行少量的學習。
混合專家模型（MoE）：混合專家模型是一種透過組合多個模型的預測結果來提高效率和效能的機器學習正規化，該正規化透過門控網路將輸入動態分配給最相關的 “專家” 來獲得預測結果。[6] 該模型利用來自不同專家的特定知識來改善在多樣、複雜的問題上的總體表現，已有研究證明 MoE 在大規模神經網路上的有效性 [7]。對於每個輸入，MoE 只會啟用一部分專家，從而在不損害模型規模的基礎上顯著提高計算效率。該方法已被證實在擴充套件基於 Transformer 的應用架構上十分有效，如 Mixtral [8]。
方法描述
該工作提出了一個可擴充套件、高效的多工嵌入架構 MeteoRA，它能夠直接使用開源社群中或面向特定下游任務微調好的 LoRA 介面卡，並依據輸入自主選擇、及時切換合適的 LoRA 介面卡。如圖 1 所示，MeteoRA 模組可以被整合到注意力模組和多層感知機模組的所有基本線性層中，每一個模組都包括了一系列低秩矩陣。透過 MoE 前向加速策略，大語言模型能夠高效解決廣泛的問題。
圖 1：集成了 MeteoRA 模組實現 MoE 架構的大語言模型框架
圖 2 展示了 MeteoRA 模組的內部結構。MeteoRA 模組中嵌入了多個 LoRA 介面卡，並透過一個門控網路來實現 MoE 架構。該門控網路會依據輸入選擇 top-k 個 LoRA 介面卡，並將它們組合作為微調權重進行前向傳播。透過這種架構，門控網路會執行路由策略，從而根據輸入選取合適的 LoRA 介面卡。每一個 MeteoRA 模組都包含一個獨立的門控網路，不同門控網路依據它們的輸入獨立決策，並在全部解碼器模組的前向傳播過程中動態選取不同 LoRA 介面卡。
圖 2：應用 MoE 模型整合 LoRA 嵌入的 MeteoRA 模組的架構
MeteoRA 模組的訓練遵循自迴歸語言建模任務下的模型微調準則，訓練中需要保持基本的大語言模型權重和預訓練的 LoRA 介面卡引數不變。由於 MeteoRA 模組支援選擇權重最高的若干個 LoRA 介面卡，團隊引入了一種將自迴歸語言建模損失和所有門控網路損失組合起來的聯合最佳化方案。該最佳化函式綜合了自迴歸模型中的預測損失和 MeteoRA 模組中 LoRA 分類的交叉熵損失，實現門控網路的訓練。
MeteoRA 模組的核心元件是整合了多個 LoRA 介面卡的 MoE 架構。首先，團隊將 LoRA 整合成在 HBM 上連續分配的張量。對於每個批次的輸入序列中的每個標記，由於該模組需要利用門控網路找到 LoRA 介面卡的編號索引集合，MeteoRA 模組幾乎不可能與單獨的 LoRA 模組保持相同的效率。基於原始迴圈（loop-original）的簡單實現 [8] 採用 for-loop 遍歷 LoRA 介面卡，在每次遍歷中對該介面卡的候選集合應用 LoRA 前向傳遞。該方法簡單地將所有批次的所有標記拆成 LoRA 數量個集合，並且使得每個集合按順序傳遞。然而，考慮到輸入標記相互獨立的性質，這種方法無法充分利用並行化 GEMM 運算元 [9] 加速，尤其是當某些 LoRA 介面卡僅被少數標記選擇，或是待推理標記小於 LoRA 數量時，實驗中可能花費最多 10 倍的執行時間。
該工作採用了前向傳播加速策略 bmm-torch，直接索引全部批次標記的 top-k 介面卡，這會利用到兩次 bmm 計算。相比於原始的迴圈方法，bmm-torch 基於 PyTorch 提供的 bmm 運算子 [10] 並行化所有批次標記的 top-k 個介面卡，實現了約 4 倍的加速。在實驗中該策略僅比它的上限單 LoRA 模組慢 2.5 倍。由於 PyTorch 的索引約束（indexing constraints）[11]，bmm-torch 需要分配更大的記憶體執行批處理，當批次或者輸入長度過大時 bmm-torch 的大記憶體開銷可能成為瓶頸。為此，該工作使用了 Triton [12] 開發了自定義 GPU 核心運算元，不僅保持了 bmm-torch 約 80% 的運算效率，而且還保持了原始迴圈級別的低記憶體開銷。
實驗結果
該工作在獨立任務和複合任務上，對所提出的架構和設計進行了廣泛的實驗驗證。實驗使用了兩種知名的大語言模型 LlaMA2-13B 和 LlaMA3-8B，圖 3 顯示了整合 28 項任務的 MeteoRA 模組與單 LoRA、參考模型 PEFT 在各項獨立任務上的預測表現雷達圖。評估結果表明，無論使用哪種大語言模型，MeteoRA 模型都具備與 PEFT 相近的效能，而 MeteoRA 不需要顯式地啟用特定的 LoRA。表 1 展示了所有方法在各項任務上的平均分數。
圖 3：MeteoRA 模型在 28 項選定任務上的評估表現
表 1：各模型在 28 項選定任務上的平均表現
為了驗證該模型按順序解決複合問題的能力，該工作透過序列連線獨立的任務來構建 3 個數據集，分別整合了 3、5、10 項任務，期望模型能夠按次序解決同一序列中輸入的不同類別問題。實驗結果如表 2 所示，可見隨著複合任務數的提升，MeteoRA 模組幾乎全面優於參考 LoRA-B 模型。
表 2：複合任務的評估結果，左側為 MeteoRA 模組，右側為 LoRA-B
為了更進一步驗證門控網路在 MeteoRA 模組中的功能，該工作展示了在複合任務推理過程中採用 top-2 策略的 LoRA 選擇模式。在兩個相鄰任務的連線處，門控網路正確執行了 LoRA 的切換操作。
圖 4：在 3 項任務複合中 LoRA 選取情況的例子
為了驗證採用自定義 GPU 運算元的前向傳播設計的運算效率，該工作在 28 項任務上截取了部分樣本，將新的前向傳播策略與其上限和原始實現對比。評估結果如圖 5 所示，展現了新的加速策略卓越的效能。
圖 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, Lélio Renard Lavaud, Lucile Saulnier, MarieAnne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothé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.
© THE END
轉載請聯絡本公眾號獲得授權
投稿或尋求報道：[email protected]