一般來(lái)說(shuō)，大語(yǔ)言模型的部署都會(huì)采用「預(yù)訓(xùn)練 — 然后微調(diào)」的模式。但是，當(dāng)針對(duì)眾多任務(wù)（如個(gè)性化助手）對(duì) base 模型進(jìn)行微調(diào)時(shí)，訓(xùn)練和服務(wù)成本會(huì)變得非常高昂。低秩適配（LowRank Adaptation，LoRA）是一種參數(shù)效率高的微調(diào)方法，通常用于將 base 模型適配到多種任務(wù)中，從而產(chǎn)生了大量從一個(gè) base 模型衍生出來(lái)的 LoRA 適配程序。

這種模式為服務(wù)過(guò)程中的批量推理提供了大量機(jī)會(huì)。LoRA 的研究表明了一點(diǎn)，只對(duì)適配器權(quán)重進(jìn)行微調(diào)，就能獲得與全權(quán)重微調(diào)相當(dāng)?shù)男阅?。雖然這種方法可以實(shí)現(xiàn)單個(gè)適配器的低延遲推理和跨適配器的串行執(zhí)行，但在同時(shí)為多個(gè)適配器提供服務(wù)時(shí)，會(huì)顯著降低整體服務(wù)吞吐量并增加總延遲?？傊?，如何大規(guī)模服務(wù)于這些微調(diào)變體的問(wèn)題仍未得到解決。

在最近的一篇論文中，來(lái)自 UC 伯克利、斯坦福等高校的研究者提出了一種名為 S-LoRA 的新微調(diào)方式。

• 論文地址：https://arxiv.org/pdf/2311.03285.pdf
• 項(xiàng)目地址：https://github.com/S-LoRA/S-LoRA

S-LoRA 是專(zhuān)為眾多 LoRA 適配程序的可擴(kuò)展服務(wù)而設(shè)計(jì)的系統(tǒng)，它將所有適配程序存儲(chǔ)在主內(nèi)存中，并將當(dāng)前運(yùn)行查詢(xún)所使用的適配程序取到 GPU 內(nèi)存中。

S-LoRA 提出了「統(tǒng)一分頁(yè)」（Unified Paging）技術(shù)，即使用統(tǒng)一的內(nèi)存池來(lái)管理不同等級(jí)的動(dòng)態(tài)適配器權(quán)重和不同序列長(zhǎng)度的 KV 緩存張量。此外，S-LoRA 還采用了新的張量并行策略和高度優(yōu)化的定制 CUDA 內(nèi)核，以實(shí)現(xiàn) LoRA 計(jì)算的異構(gòu)批處理。

這些功能使 S-LoRA 能夠以較小的開(kāi)銷(xiāo)在單個(gè) GPU 或多個(gè) GPU 上為數(shù)千個(gè) LoRA 適配器提供服務(wù)（同時(shí)為 2000 個(gè)適配器提供服務(wù)），并將增加的 LoRA 計(jì)算開(kāi)銷(xiāo)降至最低。相比之下，vLLM-packed 需要維護(hù)多個(gè)權(quán)重副本，并且由于 GPU 內(nèi)存限制，只能為少于 5 個(gè)適配器提供服務(wù)。

與 HuggingFace PEFT 和 vLLM（僅支持 LoRA 服務(wù)）等最先進(jìn)的庫(kù)相比，S-LoRA 的吞吐量最多可提高 4 倍，服務(wù)的適配器數(shù)量可增加幾個(gè)數(shù)量級(jí)。因此，S-LoRA 能夠?yàn)樵S多特定任務(wù)的微調(diào)模型提供可擴(kuò)展的服務(wù)，并為大規(guī)模定制微調(diào)服務(wù)提供了潛力。

S-LoRA 包含三個(gè)主要?jiǎng)?chuàng)新部分。論文的第 4 節(jié)介紹了批處理策略，該策略分解了 base 模型和 LoRA 適配器之間的計(jì)算。此外，研究者還解決了需求調(diào)度的難題，包括適配器集群和準(zhǔn)入控制等方面。跨并發(fā)適配器的批處理能力給內(nèi)存管理帶來(lái)了新的挑戰(zhàn)。第 5 節(jié)，研究者將 PagedAttention 推廣到 Unfied Paging，支持動(dòng)態(tài)加載 LoRA 適配器。這種方法使用統(tǒng)一的內(nèi)存池以分頁(yè)方式存儲(chǔ) KV 緩存和適配器權(quán)重，可以減少碎片并平衡 KV 緩存和適配器權(quán)重的動(dòng)態(tài)變化大小。最后，第 6 節(jié)介紹了新的張量并行策略，能夠高效地解耦 base 模型和 LoRA 適配器。

以下為重點(diǎn)內(nèi)容：

批處理 

對(duì)于單個(gè)適配器，Hu et al., 2021 推薦的方法是將適配器權(quán)重合并到 base 模型權(quán)重中，從而得到一個(gè)新模型（見(jiàn)公式 1）。這樣做的好處是在推理過(guò)程中沒(méi)有額外的適配器開(kāi)銷(xiāo)，因?yàn)樾履Ｐ偷膮?shù)數(shù)與 base 模型相同。事實(shí)上，這也是最初 LoRA 工作的一個(gè)突出特點(diǎn)。

本文指出，將 LoRA 適配器合并到 base 模型中對(duì)于多 LoRA 高吞吐量服務(wù)設(shè)置來(lái)說(shuō)效率很低。取而代之的是，研究者建議實(shí)時(shí)計(jì)算 LoRA 計(jì)算 xAB（如公式 2 所示）。

在 S-LoRA 中，計(jì)算 base 模型被批處理，然后使用定制的 CUDA 內(nèi)核分別執(zhí)行所有適配器的附加 xAB。這一過(guò)程如圖 1 所示。研究者沒(méi)有使用填充和 BLAS 庫(kù)中的批處理 GEMM 內(nèi)核來(lái)計(jì)算 LoRA，而是實(shí)施了定制的 CUDA 內(nèi)核，以便在不使用填充的情況下實(shí)現(xiàn)更高效的計(jì)算，實(shí)施細(xì)節(jié)在第 5.3 小節(jié)中。

如果將 LoRA 適配器存儲(chǔ)在主內(nèi)存中，它們的數(shù)量可能會(huì)很大，但當(dāng)前運(yùn)行批所需的 LoRA 適配器數(shù)量是可控的，因?yàn)榕笮∈?GPU 內(nèi)存的限制。為了利用這一優(yōu)勢(shì)，研究者將所有的 LoRA 適配卡都存儲(chǔ)在主內(nèi)存中，并在為當(dāng)前正在運(yùn)行的批進(jìn)行推理時(shí)，僅將該批所需的 LoRA 適配卡取到 GPU RAM 中。在這種情況下，可服務(wù)的適配器最大數(shù)量受限于主內(nèi)存大小。圖 2 展示了這一過(guò)程。第 5 節(jié)也討論了高效管理內(nèi)存的技術(shù)。

內(nèi)存管理

與為單個(gè) base 模型提供服務(wù)相比，同時(shí)為多個(gè) LoRA 適配卡提供服務(wù)會(huì)帶來(lái)新的內(nèi)存管理挑戰(zhàn)。為了支持多個(gè)適配器，S-LoRA 將它們存儲(chǔ)在主內(nèi)存中，并將當(dāng)前運(yùn)行批所需的適配器權(quán)重動(dòng)態(tài)加載到 GPU RAM 中。

在這個(gè)過(guò)程中，有兩個(gè)明顯的挑戰(zhàn)。首先是內(nèi)存碎片，這是由于動(dòng)態(tài)加載和卸載不同大小的適配器權(quán)重造成的。其次是適配器加載和卸載帶來(lái)的延遲開(kāi)銷(xiāo)。為了有效解決這些難題，研究者提出了 「Unfied Paging」，并通過(guò)預(yù)取適配器權(quán)重將 I/O 與計(jì)算重疊。

Unified Paging

研究者將 PagedAttention 的想法擴(kuò)展為統(tǒng)一分頁(yè)（Unified Paging），后者除了管理 KV 緩存外，還管理適配器權(quán)重。統(tǒng)一分頁(yè)使用統(tǒng)一內(nèi)存池來(lái)聯(lián)合管理 KV 緩存和適配器權(quán)重。為了實(shí)現(xiàn)這一點(diǎn)，他們首先為內(nèi)存池靜態(tài)分配一個(gè)大緩沖區(qū)，除了 base 模型權(quán)重和臨時(shí)激活張量占用的空間外，該緩沖區(qū)使用所有可用空間。KV 緩存和適配器權(quán)重都以分頁(yè)方式存儲(chǔ)在內(nèi)存池中，每頁(yè)對(duì)應(yīng)一個(gè) H 向量。因此，序列長(zhǎng)度為 S 的 KV 緩存張量占用 S 頁(yè)，而 R 級(jí)的 LoRA 權(quán)重張量占用 R 頁(yè)。圖 3 展示了內(nèi)存池布局，其中 KV 緩存和適配器權(quán)重以交錯(cuò)和非連續(xù)方式存儲(chǔ)。這種方法大大減少了碎片，確保不同等級(jí)的適配器權(quán)重能以結(jié)構(gòu)化和系統(tǒng)化的方式與動(dòng)態(tài) KV 緩存共存。

張量并行 

此外，研究者為批量 LoRA 推斷設(shè)計(jì)了新穎的張量并行策略，以支持大型 Transformer 模型的多 GPU 推斷。張量并行是應(yīng)用最廣泛的并行方法，因?yàn)樗膯纬绦蚨鄶?shù)據(jù)模式簡(jiǎn)化了其實(shí)施和與現(xiàn)有系統(tǒng)的集成。張量并行可以減少為大模型提供服務(wù)時(shí)每個(gè) GPU 的內(nèi)存使用量和延遲。在本文設(shè)置中，額外的 LoRA 適配器引入了新的權(quán)重矩陣和矩陣乘法，這就需要為這些新增項(xiàng)目制定新的分區(qū)策略。

評(píng)估

最后，研究者通過(guò)為 Llama-7B/13B/30B/70B 提供服務(wù)來(lái)評(píng)估 S-LoRA。

結(jié)果表明，S-LoRA 可以在單個(gè) GPU 或多個(gè) GPU 上為數(shù)千個(gè) LoRA 適配器提供服務(wù)，而且開(kāi)銷(xiāo)很小。與最先進(jìn)的參數(shù)高效微調(diào)庫(kù) Huggingface PEFT 相比，S-LoRA 的吞吐量最多可提高 30 倍。與使用支持 LoRA 服務(wù)的高吞吐量服務(wù)系統(tǒng) vLLM 相比，S-LoRA 可將吞吐量提高 4 倍，并將服務(wù)適配器的數(shù)量增加幾個(gè)數(shù)量級(jí)。

更多研究細(xì)節(jié)，可參考原論文。