編輯：蛋醬
一般來說，大語言模型的部署都會采用「預訓練 — 然后微調」的模式。但是，當針對眾多任務（如個性化助手）對 base 模型進行微調時，訓練和服務成本會變得非常高昂。低秩適配（LowRank Adaptation，LoRA）是一種參數效率高的微調方法，通常用于將 base 模型適配到多種任務中，從而產生了大量從一個 base 模型衍生出來的 LoRA 適配程序。
這種模式為服務過程中的批量推理提供了大量機會。LoRA 的研究表明了一點，只對適配器權重進行微調，就能獲得與全權重微調相當的性能。雖然這種方法可以實現單個適配器的低延遲推理和跨適配器的串行執行，但在同時為多個適配器提供服務時，會顯著降低整體服務吞吐量并增加總延遲。總之，如何大規模服務于這些微調變體的問題仍未得到解決。
在最近的一篇論文中，來自 UC 伯克利、斯坦福等高校的研究者提出了一種名為 S-LoRA 的新微調方式。
- 論文地址：https://arxiv.org/pdf/2311.03285.pdf
- 項目地址：https://Github.com/S-LoRA/S-LoRA
S-LoRA 是專為眾多 LoRA 適配程序的可擴展服務而設計的系統，它將所有適配程序存儲在主內存中，并將當前運行查詢所使用的適配程序取到 GPU 內存中。
S-LoRA 提出了「統一分頁」（Unified Paging）技術，即使用統一的內存池來管理不同等級的動態適配器權重和不同序列長度的 KV 緩存張量。此外，S-LoRA 還采用了新的張量并行策略和高度優化的定制 CUDA 內核，以實現 LoRA 計算的異構批處理。
這些功能使 S-LoRA 能夠以較小的開銷在單個 GPU 或多個 GPU 上為數千個 LoRA 適配器提供服務（同時為 2000 個適配器提供服務），并將增加的 LoRA 計算開銷降至最低。相比之下，vLLM-packed 需要維護多個權重副本，并且由于 GPU 內存限制，只能為少于 5 個適配器提供服務。
與 HuggingFace PEFT 和 vLLM（僅支持 LoRA 服務）等最先進的庫相比，S-LoRA 的吞吐量最多可提高 4 倍，服務的適配器數量可增加幾個數量級。因此，S-LoRA 能夠為許多特定任務的微調模型提供可擴展的服務，并為大規模定制微調服務提供了潛力。
S-LoRA 包含三個主要創新部分。論文的第 4 節介紹了批處理策略，該策略分解了 base 模型和 LoRA 適配器之間的計算。此外，研究者還解決了需求調度的難題，包括適配器集群和準入控制等方面。跨并發適配器的批處理能力給內存管理帶來了新的挑戰。第 5 節，研究者將 PagedAttention 推廣到 Unfied Paging，支持動態加載 LoRA 適配器。這種方法使用統一的內存池以分頁方式存儲 KV 緩存和適配器權重，可以減少碎片并平衡 KV 緩存和適配器權重的動態變化大小。最后，第 6 節介紹了新的張量并行策略，能夠高效地解耦 base 模型和 LoRA 適配器。
以下為重點內容：
批處理
對于單個適配器，Hu et al., 2021 推薦的方法是將適配器權重合并到 base 模型權重中，從而得到一個新模型（見公式 1）。這樣做的好處是在推理過程中沒有額外的適配器開銷，因為新模型的參數數與 base 模型相同。事實上，這也是最初 LoRA 工作的一個突出特點。
本文指出，將 LoRA 適配器合并到 base 模型中對于多 LoRA 高吞吐量服務設置來說效率很低。取而代之的是，研究者建議實時計算 LoRA 計算 xAB（如公式 2 所示）。
在 S-LoRA 中，計算 base 模型被批處理，然后使用定制的 CUDA 內核分別執行所有適配器的附加 xAB。這一過程如圖 1 所示。研究者沒有使用填充和 BLAS 庫中的批處理 GEMM 內核來計算 LoRA，而是實施了定制的 CUDA 內核，以便在不使用填充的情況下實現更高效的計算，實施細節在第 5.3 小節中。
如果將 LoRA 適配器存儲在主內存中，它們的數量可能會很大，但當前運行批所需的 LoRA 適配器數量是可控的，因為批大小受 GPU 內存的限制。為了利用這一優勢，研究者將所有的 LoRA 適配卡都存儲在主內存中，并在為當前正在運行的批進行推理時，僅將該批所需的 LoRA 適配卡取到 GPU RAM 中。在這種情況下，可服務的適配器最大數量受限于主內存大小。圖 2 展示了這一過程。第 5 節也討論了高效管理內存的技術。
內存管理
與為單個 base 模型提供服務相比，同時為多個 LoRA 適配卡提供服務會帶來新的內存管理挑戰。為了支持多個適配器，S-LoRA 將它們存儲在主內存中，并將當前運行批所需的適配器權重動態加載到 GPU RAM 中。
在這個過程中，有兩個明顯的挑戰。首先是內存碎片，這是由于動態加載和卸載不同大小的適配器權重造成的。其次是適配器加載和卸載帶來的延遲開銷。為了有效解決這些難題，研究者提出了 「Unfied Paging」，并通過預取適配器權重將 I/O 與計算重疊。
Unified Paging
研究者將 PagedAttention 的想法擴展為統一分頁（Unified Paging），后者除了管理 KV 緩存外，還管理適配器權重。統一分頁使用統一內存池來聯合管理 KV 緩存和適配器權重。為了實現這一點，他們首先為內存池靜態分配一個大緩沖區，除了 base 模型權重和臨時激活張量占用的空間外，該緩沖區使用所有可用空間。KV 緩存和適配器權重都以分頁方式存儲在內存池中，每頁對應一個 H 向量。因此，序列長度為 S 的 KV 緩存張量占用 S 頁，而 R 級的 LoRA 權重張量占用 R 頁。圖 3 展示了內存池布局，其中 KV 緩存和適配器權重以交錯和非連續方式存儲。這種方法大大減少了碎片，確保不同等級的適配器權重能以結構化和系統化的方式與動態 KV 緩存共存。
張量并行
此外，研究者為批量 LoRA 推斷設計了新穎的張量并行策略，以支持大型 Transformer 模型的多 GPU 推斷。張量并行是應用最廣泛的并行方法，因為它的單程序多數據模式簡化了其實施和與現有系統的集成。張量并行可以減少為大模型提供服務時每個 GPU 的內存使用量和延遲。在本文設置中，額外的 LoRA 適配器引入了新的權重矩陣和矩陣乘法，這就需要為這些新增項目制定新的分區策略。
評估
最后，研究者通過為 Llama-7B/13B/30B/70B 提供服務來評估 S-LoRA。
結果表明，S-LoRA 可以在單個 GPU 或多個 GPU 上為數千個 LoRA 適配器提供服務，而且開銷很小。與最先進的參數高效微調庫 Huggingface PEFT 相比，S-LoRA 的吞吐量最多可提高 30 倍。與使用支持 LoRA 服務的高吞吐量服務系統 vLLM 相比，S-LoRA 可將吞吐量提高 4 倍，并將服務適配器的數量增加幾個數量級。
