作者 | 陳樂(lè)群

單位 | 華盛頓大學(xué)博士生

方向 | 機(jī)器學(xué)習(xí)系統(tǒng)及分布式系統(tǒng)

來(lái)自 | PaperWeekly

隨著開(kāi)源預(yù)訓(xùn)練大型語(yǔ)言模型（Large Language Model, LLM ）變得更加強(qiáng)大和開(kāi)放，越來(lái)越多的開(kāi)發(fā)者將大語(yǔ)言模型納入到他們的項(xiàng)目中。其中一個(gè)關(guān)鍵的適應(yīng)步驟是將領(lǐng)域特定的文檔集成到預(yù)訓(xùn)練模型中，這被稱為微調(diào)。

通常情況下，來(lái)自領(lǐng)域特定文檔的額外知識(shí)與預(yù)訓(xùn)練模型已經(jīng)知道的相比微不足道。在這種情況下，低秩適應(yīng)（Low-Rank Adaptation，LoRA ）技術(shù)證明是有價(jià)值的。

通過(guò) LoRA，微調(diào)模型僅向預(yù)訓(xùn)練模型添加不到 0.1% 的參數(shù)。具體來(lái)說(shuō)，這意味著 LoRA 微調(diào)模型僅增加了10~200MB 的存儲(chǔ)，具體取決于配置。從計(jì)算角度來(lái)看，考慮到與預(yù)訓(xùn)練模型相比參數(shù)的增加極少，額外的計(jì)算負(fù)載相對(duì)較小。

基于存儲(chǔ)和計(jì)算的額外開(kāi)銷都很小這一點(diǎn)，我相信構(gòu)建一個(gè)多租戶的大語(yǔ)言微調(diào)模型的推斷服務(wù)具有很大潛力。這個(gè)服務(wù)可以托管成千上萬(wàn)個(gè) LoRA 模型，它們都共享相同的預(yù)訓(xùn)練大語(yǔ)言模型。在每個(gè)批次的執(zhí)行中，每個(gè)用戶請(qǐng)求都會(huì)調(diào)用一個(gè)獨(dú)立的微調(diào)模型，從而分?jǐn)偞鎯?chǔ)和計(jì)算成本到各種不同的模型中。

在我的上一篇文章中，我深入探討了大語(yǔ)言模型推斷中的批處理效應(yīng)。在這篇文章中，我將詳細(xì)介紹為什么多租戶 LoRA 推斷服務(wù)具有巨大的潛力。

背景知識(shí)：文本生成

文本生成服務(wù)，如 ChatGPT，接受用戶文本輸入并提供文本響應(yīng)。這個(gè)輸入被稱為“提示”（prompt）。在內(nèi)部，當(dāng)大語(yǔ)言模型處理文本時(shí)，它在一系列“詞元”（token）上操作。我們可以大致將詞元視為幾個(gè)字符或一個(gè)單詞。文本生成過(guò)程有兩個(gè)主要階段：

預(yù)填充階段（或稱“編碼”，“初始化”）接受整個(gè)提示并生成隨后的詞元以及一個(gè)“鍵值緩存”（KV Cache）。

解碼階段處理新生成的詞元和鍵值緩存，然后生成下一個(gè)詞元，同時(shí)更新鍵值緩存。這個(gè)階段不斷重復(fù)，直到模型完成其輸出。

有趣的是，盡管預(yù)填充階段處理的詞元數(shù)量比解碼階段多 100 倍，但它們的計(jì)算延遲是在同一數(shù)量級(jí)的。由于解碼階段會(huì)重復(fù)執(zhí)行上百次，因此在這篇文章中，我將集中討論如何優(yōu)化解碼階段。

背景知識(shí)：大語(yǔ)言模型的架構(gòu)和批處理

在其核心，大語(yǔ)言模型的架構(gòu)非常簡(jiǎn)單。它主要包括多個(gè) Transformer 層，所有層都共享相同的架構(gòu)。每一層包括四個(gè)計(jì)算密集型組件：QKV 投影、自注意力、輸出投影和前饋網(wǎng)絡(luò)（Feed-Forward Network, FFN ）。

▲ Transformer 層在解碼階段的運(yùn)算圖示

概括地說(shuō)，其中包含了兩種算子：

自注意力（Self-Attention，黃色標(biāo)出）涉及矩陣-矩陣乘法。

密集投影（Dense Projection，綠色標(biāo)出）涉及向量-矩陣乘法。

考慮到每個(gè)批次中每個(gè)序列只有一個(gè)詞元，密集投影計(jì)算也非常微小，不足以充分利用 GPU 。因此，擴(kuò)大批處理大小幾乎不會(huì)影響密集投影的延遲，這使得大批處理大小對(duì)于構(gòu)建高吞吐、低延遲的推斷服務(wù)至關(guān)重要。

關(guān)于大語(yǔ)言模型推斷中的批處理的更詳細(xì)分析，請(qǐng)查看我的以前的博客文章。

背景知識(shí)：LoRA

給定形狀為 [H1, H2] 的預(yù)訓(xùn)練參數(shù)矩陣 W ，LoRA 微調(diào)訓(xùn)練一個(gè)形狀為。 [H1, R] 的小矩陣 A 和形狀為 [R, H2] 的 B 。我們使用（W+AB）作為微調(diào)模型的權(quán)重。這里的 R 是 LoRA 微調(diào)指定的秩，通常遠(yuǎn)小于原始維度（>= 4096），通常在 8~32 之間。

這種方法背后的邏輯是，與原始權(quán)重相比，新增加的知識(shí)只占一小部分，因此 LoRA 將增量壓縮成兩個(gè)低秩矩陣。

與完全微調(diào)相比，LoRA 預(yù)訓(xùn)練顯著降低了微調(diào)模型的存儲(chǔ)和內(nèi)存需求。

在大語(yǔ)言模型中，由于所有參數(shù)都位于密集投影中，LoRA 可以集成到 Transformer 層中的任何位置。雖然 HuggingFace PEFT 庫(kù)僅將 LoRA 加到 q_proj 和 v_proj ，但一些研究，如 QLoRA，主張將其包含在所有的密集投影中。

LoRA 延遲和批處理效應(yīng)

盡管在存儲(chǔ)方面，LoRA 矩陣明顯小于原始權(quán)重矩陣，但計(jì)算的延遲并不成比例地減少。我們可以使用以下代碼對(duì)骨干模型和 LoRA 增量的延遲進(jìn)行基準(zhǔn)測(cè)試：

h1=4096
h2=11008
r=16
forbsinrange(1,33):
w=torch.randn(h1,h2,dtype=torch.float16,device="cuda:0")
a=torch.randn(bs,h1,r,dtype=torch.float16,device="cuda:0")
b=torch.randn(bs,r,h2,dtype=torch.float16,device="cuda:0")
x=torch.randn(bs,1,h1,dtype=torch.float16,device="cuda:0")
bench(lambda:x@w)
bench(lambda:x@a@b)

▲ 骨干大語(yǔ)言模型和 LoRA 微調(diào)模型的推斷延遲比較

上圖表明，LoRA 增量?jī)H比骨干模型快 2.5 倍。

然而，顯然 LoRA 批處理效應(yīng)與骨干模型的效應(yīng)相似。批處理大小的增加僅在較小程度上影響延遲。這個(gè)特性使得多租戶 LoRA 非常可行。

在骨干模型中，所有一個(gè)批次中的所有請(qǐng)求都針對(duì)同一模型。而在多租戶 LoRA 服務(wù)中，一個(gè)批次中的請(qǐng)求可能會(huì)調(diào)用不同的 LoRA 微調(diào)模型。數(shù)學(xué)表達(dá)如下：

挑戰(zhàn)在于以批處理的方式將不同的 LoRA 增量應(yīng)用于一個(gè)批處理中的各個(gè)輸入，與此同時(shí)還要維持“白吃的午餐”式的批處理效應(yīng)。

批處理 LoRA 算子

我們想要的批處理 LoRA 算子具有以下函數(shù)簽名：

defadd_lora(
y:torch.Tensor,#(batch_size,1,out_features)
x:torch.Tensor,#(batch_size,1,in_features)
A:torch.Tensor,#(num_loras,in_features,lora_rank)
B:torch.Tensor,#(num_loras,lora_rank,out_features)
I:torch.LongTensor,#(batch_size,)
):
"""Semantics:y[i]+=x[i]@A[I[i]]@B[I[i]]"""
raiseNotImplementedError()

一個(gè)最簡(jiǎn)單的實(shí)現(xiàn)方法是在批處理維度上進(jìn)行循環(huán)：?

deflora_loop(
y:torch.Tensor,#(batch_size,1,out_features)
x:torch.Tensor,#(batch_size,1,in_features)
A:torch.Tensor,#(num_loras,in_features,lora_rank)
B:torch.Tensor,#(num_loras,lora_rank,out_features)
I:torch.LongTensor,#(batch_size,)
):
fori,idxinenumerate(I.cpu().numpy()):
y[i]+=x[i]@A[idx]@B[idx]

讓我們對(duì)循環(huán)版本進(jìn)行基準(zhǔn)測(cè)試。為了進(jìn)行比較，我們可以包括一個(gè)“作弊”實(shí)現(xiàn)，其中我們假設(shè)批處理中每個(gè)請(qǐng)求的 LoRA 矩陣已經(jīng)被合并在一起。這樣，我們只測(cè)量批量矩陣乘法（Batched Matrix Multiplication, bmm）的延遲。??

deflora_cheat_bmm(
y:torch.Tensor,#(batch_size,1,out_features)
x:torch.Tensor,#(batch_size,1,in_features)
cheat_A:torch.Tensor,#(batch_size,in_features,lora_rank)
cheat_B:torch.Tensor,#(batch_size,lora_rank,out_features)
):
y+=x@cheat_A@cheat_B


num_loras=50
h1=4096
h2=11008
r=16
A=torch.randn(num_loras,h1,r,dtype=torch.float16,device="cuda:0")
B=torch.randn(num_loras,r,h2,dtype=torch.float16,device="cuda:0")
forbsinrange(1,33):
x=torch.randn(bs,1,h1,dtype=torch.float16,device="cuda:0")
y=torch.randn(bs,1,h2,dtype=torch.float16,device="cuda:0")
I=torch.randint(num_loras,(bs,),dtype=torch.long,device="cuda:0")
cheat_A=A[I,:,:]
cheat_B=B[I,:,:]

bench(lambda:lora_loop(y,x,A,B,I))
bench(lambda:lora_cheat_bmm(y,x,cheat_A,cheat_B))

▲ LoRA 實(shí)現(xiàn)：for-loop vs bmm
可預(yù)見(jiàn)的是，循環(huán)版本明顯較慢，并且失去了批處理效應(yīng)。這是因?yàn)樗饌€(gè)處理輸入，而不是利用為批處理數(shù)據(jù)設(shè)計(jì)的高效 CUDA 核心。
然而， bmm 方法提供了一個(gè)很好的啟發(fā)。我們的目標(biāo)變得清晰起來(lái)：首先將所有 LoRA 矩陣匯總到一個(gè)臨時(shí)的張量中，然后使用 bmm 。經(jīng)過(guò)一番挖掘，我發(fā)現(xiàn)了 torch.index_select() 函數(shù)，它可以高效地執(zhí)行批量匯總（Batched Gather）。于是我們可以如下一個(gè) gbmm?（ gather-bmm ）實(shí)現(xiàn)：?

deflora_gbmm(
y:torch.Tensor,#(batch_size,1,out_features)
x:torch.Tensor,#(batch_size,1,in_features)
A:torch.Tensor,#(num_loras,in_features,lora_rank)
B:torch.Tensor,#(num_loras,lora_rank,out_features)
I:torch.LongTensor,#(batch_size,)
):
a=torch.index_select(A,0,I)#(batch_size,in_features,lora_rank)
b=torch.index_select(B,0,I)#(batch_size,lora_rank,out_features)
y+=x@a@b

BGMV 算子

雖然 gbmm 非常有效，但它并不是最終解決方案。我們沒(méi)有必要僅僅是因?yàn)?bmm 需要連續(xù)的存儲(chǔ)而將 LoRA 增量而匯總到一個(gè)連續(xù)的空間中。理想情況下，聚合可以在 CUDA 核心內(nèi)部進(jìn)行，于 bmm 操作同事進(jìn)行。如果可能的話，這將消除與 torch.index_select() 相關(guān)的 GPU 內(nèi)存讀寫(xiě)操作。

我請(qǐng)葉子豪大牛幫忙，他是精通高性能 CUDA 核心的編寫(xiě)。經(jīng)過(guò)幾輪迭代，子豪開(kāi)發(fā)了一個(gè)非常快的 CUDA 程序，把 LoRA 所需的計(jì)算分成兩半。我們將這個(gè)算子命名為 BGMV（Batched Gather Matrix-Vector Multiplication）：

▲ Batched Gather Matrix-Vector Multiplication (BGMV)

deflora_bgmv(
y:torch.Tensor,#(batch_size,1,out_features)
x:torch.Tensor,#(batch_size,1,in_features)
A:torch.Tensor,#(num_loras,in_features,lora_rank)
B:torch.Tensor,#(num_loras,lora_rank,out_features)
I:torch.LongTensor,#(batch_size,)
):
tmp=torch.zeros((x.size(0),A.size(-1)),dtype=x.dtype,device=x.device)
bgmv(tmp,x,A,I)
bgmv(y,tmp,B,I)

這是 gbmm 和 bgmv 的基準(zhǔn)結(jié)果：

▲ LoRA 實(shí)現(xiàn): for-loop vs bmm vs gbmm vs bgmv

如圖所示， gbmm 非常有效。聚合過(guò)程相對(duì)于 bmm 增加了約 20% 的延遲，但保持了令人滿意的批處理效應(yīng)。而子豪編寫(xiě)的 bgmv 算子更加高效，甚至超過(guò)了 bmm 。

多租戶 LoRA 文本生成性能

在 bgmv 算子的基礎(chǔ)上，我開(kāi)發(fā)了一個(gè)名為Punica的實(shí)驗(yàn)項(xiàng)目，支持多個(gè) LoRA 模型。Punica 的獨(dú)特能力是能把不同 LoRA 模型的請(qǐng)求合并在一個(gè)批處理中。我對(duì)比了 Punica 和一系列知名系統(tǒng)的性能，包括 HuggingFace Transformers、DeepSpeed、Faster Transformer 和 vLLM。

在測(cè)試中，每個(gè)請(qǐng)求都針對(duì)不同的 LoRA 模型。鑒于其他系統(tǒng)沒(méi)有明確針對(duì)多租戶 LoRA 服務(wù)進(jìn)行優(yōu)化，它們的批處理大小為 1。我使用了 HuggingFace PEFT庫(kù)來(lái)把 LoRA 加到 HuggingFace Transformers 和 DeepSpeed 中。我尚未調(diào)整 vLLM 和 Faster Transformer，所以它們?cè)跊](méi)有 LoRA 的情況下運(yùn)行。以下是結(jié)果：

▲ 多租戶 LoRA 文本生成的吞吐量

DeepSpeed、vLLM 和 Faster Transformers 都有高度優(yōu)化的 Transformer 實(shí)現(xiàn)，這些系統(tǒng)的吞吐量比標(biāo)準(zhǔn)的 HuggingFace Transformers 高出 3 倍。但由于批處理大小被限制在 1，它們?cè)诙嘧鈶?LoRA 服務(wù)效率方面表現(xiàn)不佳。

相反，Punica 在批處理大小為 16 時(shí)比這些系統(tǒng)高出 8 倍，與普通的 HuggingFace Transformers 相比，甚至高達(dá) 23 倍。值得注意的是，Punica 的吞吐量幾乎與批處理大小呈線性關(guān)系。

值得注意的是，Punica 仍然是一個(gè)早期的研究原型。它目前使用與普通的 HuggingFace Transformers 相同的 Transformer 實(shí)現(xiàn)，除了 LoRA 和自注意力運(yùn)算符。一些已知的 Transformer 層的優(yōu)化在 Punica 中尚未實(shí)現(xiàn)。這解釋了在批處理大小為 1 時(shí) Punica 與其他高度優(yōu)化的系統(tǒng)之間的性能差距。

吞吐量不錯(cuò)，那么從延遲方面表現(xiàn)如何？

▲ 多租戶 LoRA 文本生成的延遲

如圖所示，Punica 中的批處理不會(huì)引入顯著的延遲。

示例用途

在展示了多租戶 LoRA 服務(wù)的高效性后，讓我們?cè)O(shè)想一些潛在的應(yīng)用場(chǎng)景：

用一本新的小說(shuō)對(duì) LoRA 模型進(jìn)行細(xì)調(diào)，以幫助讀者總結(jié)每個(gè)角色的旅程。

針對(duì)迅速發(fā)展的新聞對(duì) LoRA 模型進(jìn)行調(diào)整，以讓讀者了解最新動(dòng)態(tài)。

基于網(wǎng)頁(yè)內(nèi)容對(duì) LoRA 模型進(jìn)行優(yōu)化，提高讀者的理解能力。

我將這種方法稱為“即時(shí)細(xì)調(diào)”（Just-in-time Fine-tuning），因?yàn)?LoRA 的訓(xùn)練速度非常快（在我的試驗(yàn)中，每個(gè)訓(xùn)練周期不到一秒）。

總結(jié)

本文展示了用批處理加速多個(gè) LoRA 微調(diào)模型并行推斷的可行性。我實(shí)現(xiàn)的 Punica 項(xiàng)目展現(xiàn)出了關(guān)于批處理大小幾乎線性的吞吐量擴(kuò)展，并且增加批處理大小并不顯著增加延遲。

這項(xiàng)研究仍然在進(jìn)行中。我正在積極開(kāi)展這項(xiàng)研究項(xiàng)目，預(yù)計(jì)很快會(huì)發(fā)布一個(gè)在線演示。我歡迎任何反饋或想法，請(qǐng)隨時(shí)在評(píng)論部分中分享。

最近這幾個(gè)月出來(lái)了很多非常棒的開(kāi)源大語(yǔ)言模型。要讓開(kāi)源預(yù)訓(xùn)練模型能更好的跑在我們要做的事情上，我們可以對(duì)這些預(yù)訓(xùn)練模型進(jìn)行微調(diào)。LoRA 是一種非常高效的微調(diào)技術(shù)，哪怕預(yù)訓(xùn)練模型要占好幾百 GB 空間，用 LoRA 進(jìn)行微調(diào)也只需要增加 1% 左右的空間。

這里我先簡(jiǎn)單概括一下 LoRA 是怎么做的。假設(shè)預(yù)訓(xùn)練模型有一個(gè)形狀是 [H1, H2] 的參數(shù)矩陣 W ，用 LoRA 進(jìn)行微調(diào)的時(shí)候，我們會(huì)增加兩個(gè)小矩陣， [H1, r] 的的矩陣 A ，以及 [r, H2] 的矩陣 B， H1 和 H2 遠(yuǎn)大于 r 。

舉個(gè) Llama2-7B 的例子， H1=4096, H2=11008 ，而我們只需要設(shè)置 r=16 就能獲得還不錯(cuò)的效果。在計(jì)算的時(shí)候，我們把 W+AB 作為微調(diào)后的參數(shù)。對(duì)于給定一個(gè)輸入 x ，計(jì)算微調(diào)模型就是 y := x @ (W + A@B) ，等價(jià)于 y := x@W + x@A@B 。

當(dāng)我們有 n 個(gè) LoRA 微調(diào)模型的時(shí)候，我們就會(huì)有一堆小矩陣， A1, B1, A2, B2, ..., An, Bn 。對(duì)于一個(gè)輸入 batch X := (x1,x2,..,xn) ，我們假設(shè) batch 其中的每一個(gè)輸入都對(duì)應(yīng)一個(gè) LoRA 微調(diào)模型，那么這個(gè) batch 的計(jì)算結(jié)果就是 Y := X@W + (x1@A1@B1, x2@A2@B2, ..., xn@An@Bn) 。

注意到這個(gè)加法的左邊就是直接把這個(gè) batch 輸入給原來(lái)的預(yù)訓(xùn)練模型。我們知道預(yù)訓(xùn)練模型 batch 起來(lái)跑其實(shí)是很快的，因?yàn)榇笳Z(yǔ)言模型有很強(qiáng)的批處理效應(yīng)（我之前有發(fā)過(guò)一篇知乎講這個(gè)事情： 剖析 GPT 推斷中的批處理效應(yīng)。）所以這里的難題就是，右手邊這一項(xiàng) LoRA 要怎么高效的計(jì)算？

我們最近做了一個(gè)項(xiàng)目叫做 Punica，能夠非常高效的計(jì)算右邊這一項(xiàng)，延遲非常低，而且保留了和預(yù)訓(xùn)練模型一樣的很強(qiáng)的批處理效應(yīng)。我們實(shí)現(xiàn)了一個(gè) CUDA kernel，叫做 SGMV，用來(lái)把 LoRA 這部分的計(jì)算 batch 起來(lái)算。下面這個(gè)圖介紹了 SGMV 的語(yǔ)義：

▲ SGMV 的語(yǔ)義

我們測(cè)了一下性能，從下面這張圖可以看出來(lái)三個(gè)事情。（1）預(yù)訓(xùn)練模型有很強(qiáng)的批處理效應(yīng)，增加批處理大小不會(huì)顯著增加延遲。（2）用簡(jiǎn)單的方法計(jì)算 LoRA 跑起來(lái)非常慢，而且破壞了這個(gè)批處理效應(yīng)。（3）我們用 SGMV 計(jì)算 LoRA 非常地快，而且保留了很強(qiáng)的批處理效應(yīng)。

▲ 算子的性能測(cè)試

我們還跟主流的大語(yǔ)言模型系統(tǒng)做了比較了一下文本生成的性能，包括 HuggingFace Transformers, DeepSpeed, FasterTransformer, vLLM。我們考慮了 4 種不同的運(yùn)行場(chǎng)景。

Distinct 是每一個(gè)輸入對(duì)應(yīng)一個(gè)不同的 LoRA 微調(diào)模型，Identical 是所有的輸入都指向同一個(gè) LoRA 微調(diào)模型（也就是整個(gè)系統(tǒng)只有一個(gè)模型），Uniform 和 Skewed 介于兩者之間，也就是有的 LoRA 微調(diào)模型更熱門(mén)一點(diǎn)，有的 LoRA 微調(diào)模型更冷門(mén)一點(diǎn)。從下面這個(gè)圖中我們可以看出，Punica 可以達(dá)到 12 倍于現(xiàn)有系統(tǒng)的吞吐量。

▲ 文本生成的性能比較

審核編輯：湯梓紅