大語言模型推斷中的批處理效應(yīng)

作者 | 陳樂群

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

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

來自 | PaperWeekly

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

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

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

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

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

背景知識：文本生成

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

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

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

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

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

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

▲ Transformer 層在解碼階段的運算圖示

概括地說，其中包含了兩種算子：

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

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

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

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

背景知識：LoRA

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

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

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

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

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

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

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)

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

上圖表明，LoRA 增量僅比骨干模型快 2.5 倍。

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

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

挑戰(zhàn)在于以批處理的方式將不同的 LoRA 增量應(yīng)用于一個批處理中的各個輸入，與此同時還要維持“白吃的午餐”式的批處理效應(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()

一個最簡單的實現(xià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]

讓我們對循環(huán)版本進行基準測試。為了進行比較，我們可以包括一個“作弊”實現(xiàn)，其中我們假設(shè)批處理中每個請求的 LoRA 矩陣已經(jīng)被合并在一起。這樣，我們只測量批量矩陣乘法（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 實現(xiàn)：for-loop vs bmm
可預(yù)見的是，循環(huán)版本明顯較慢，并且失去了批處理效應(yīng)。這是因為它逐個處理輸入，而不是利用為批處理數(shù)據(jù)設(shè)計的高效 CUDA 核心。
然而， bmm 方法提供了一個很好的啟發(fā)。我們的目標變得清晰起來：首先將所有 LoRA 矩陣匯總到一個臨時的張量中，然后使用 bmm 。經(jīng)過一番挖掘，我發(fā)現(xiàn)了 torch.index_select() 函數(shù)，它可以高效地執(zhí)行批量匯總（Batched Gather）。于是我們可以如下一個 gbmm?（ gather-bmm ）實現(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 非常有效，但它并不是最終解決方案。我們沒有必要僅僅是因為 bmm 需要連續(xù)的存儲而將 LoRA 增量而匯總到一個連續(xù)的空間中。理想情況下，聚合可以在 CUDA 核心內(nèi)部進行，于 bmm 操作同事進行。如果可能的話，這將消除與 torch.index_select() 相關(guān)的 GPU 內(nèi)存讀寫操作。

我請葉子豪大牛幫忙，他是精通高性能 CUDA 核心的編寫。經(jīng)過幾輪迭代，子豪開發(fā)了一個非常快的 CUDA 程序，把 LoRA 所需的計算分成兩半。我們將這個算子命名為 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 的基準結(jié)果：

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

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

多租戶 LoRA 文本生成性能

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

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

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

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

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

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

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

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

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

示例用途

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

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

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

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

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

總結(jié)

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

這項研究仍然在進行中。我正在積極開展這項研究項目，預(yù)計很快會發(fā)布一個在線演示。我歡迎任何反饋或想法，請隨時在評論部分中分享。

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

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

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

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

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

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

▲ SGMV 的語義

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

▲ 算子的性能測試

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

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

▲ 文本生成的性能比較

審核編輯：湯梓紅