#AIGC創(chuàng)新先鋒者征文大賽# 部署 LLMs 前如何計(jì)算與優(yōu)化 GPU 內(nèi)存需求? 原創(chuàng) 精華
??【本文正在參與 AI.x社區(qū)AIGC創(chuàng)新先鋒者征文大賽】??
??http://scjtxx.cn/aigc/2223.html??
編者按:想要部署大語言模型(LLMs),卻不知該如何估算所需的 GPU 內(nèi)存?在項(xiàng)目預(yù)算有限的情況下,是否曾因?yàn)?GPU 內(nèi)存估算不準(zhǔn)而導(dǎo)致資源浪費(fèi)或性能不足?這些問題不僅影響項(xiàng)目進(jìn)度,還可能直接導(dǎo)致成本超支或服務(wù)質(zhì)量下降。
本文作者憑借豐富的 LLM 部署經(jīng)驗(yàn),深入剖析了 GPU 內(nèi)存需求的計(jì)算方法。從模型參數(shù)到 KV 緩存,從激活值到系統(tǒng)開銷,文章全面而詳細(xì)地講解了各個(gè)組成部分的內(nèi)存占用。文章還討論了內(nèi)存管理面臨的挑戰(zhàn),如內(nèi)存碎片、過度分配和高級(jí)解碼算法帶來的額外需求。為解決這些問題,作者介紹了 PagedAttention 和 vLLM 等優(yōu)化技術(shù),當(dāng) GPU 內(nèi)存不足時(shí),還可以采用 Swapping 和 Recomputation 等優(yōu)化策略。
作者 | Muhammad Saad Uddin
編譯 | 岳揚(yáng)
將 LLMs 投入生產(chǎn)環(huán)境使用,會(huì)遇到諸多挑戰(zhàn),尤其是想要為 LLMs 高效提供所需的計(jì)算資源時(shí)。有過此類經(jīng)驗(yàn)的人可能深有體會(huì),GPU 內(nèi)存是支持 LLMs 運(yùn)行的一項(xiàng)關(guān)鍵資源。由于這些模型體積巨大,且推理過程具有動(dòng)態(tài)性質(zhì),因此對(duì) GPU 內(nèi)存使用的規(guī)劃和優(yōu)化提出了更高的要求。
Image by Author via DallE
出于以下幾個(gè)原因,準(zhǔn)確估算 GPU 內(nèi)存的需求至關(guān)重要:
- 成本效益:GPU資源成本高昂。高估內(nèi)存需求會(huì)導(dǎo)致不必要的硬件支出,而低估內(nèi)存需求則會(huì)導(dǎo)致系統(tǒng)故障或性能下降。
- 性能優(yōu)化:合理的內(nèi)存管理能夠保障模型的高效運(yùn)行,從而快速響應(yīng)用戶需求,并提高并發(fā)處理能力。
- 可擴(kuò)展性:隨著業(yè)務(wù)需求的增長,準(zhǔn)確掌握內(nèi)存需求對(duì)于在不影響性能和不產(chǎn)生過高成本的情況下擴(kuò)展服務(wù)至關(guān)重要。
然而,計(jì)算 LLMs 服務(wù)所需的 GPU 內(nèi)存并非一件簡單的事。模型的大小(model size)、序列長度(sequence lengths)、批處理數(shù)量(batch sizes)以及解碼算法(decoding algorithms)等多種因素,都會(huì)以復(fù)雜的方式影響內(nèi)存使用。而且,傳統(tǒng)的內(nèi)存分配方法常因內(nèi)存碎片和鍵值(KV)緩存等動(dòng)態(tài)內(nèi)存組件的低效管理而造成大量浪費(fèi)。
在本文中,我將盡可能詳細(xì)地解釋如何計(jì)算 LLMs 服務(wù)所需的 GPU 內(nèi)存。我將分析影響內(nèi)存使用的各部分,并根據(jù)模型參數(shù)和工作負(fù)載特征(workload characteristics),逐步介紹如何估算內(nèi)存占用大小。同時(shí),我還會(huì)探討 Paged Attention 和 vLLM 等先進(jìn)的優(yōu)化技術(shù),這些技術(shù)能顯著降低內(nèi)存消耗并提升處理能力。通過閱讀本文,你將能夠全面了解如何規(guī)劃和優(yōu)化 LLMs 的 GPU 內(nèi)存使用,從而在實(shí)際應(yīng)用中實(shí)現(xiàn)高效且低成本的 LLMs 部署。
01 了解 LLM 推理過程中,主要消耗 GPU 內(nèi)存的幾部分
要想掌握 GPU 內(nèi)存的計(jì)算方法,最關(guān)鍵的是了解各部分如何占用 GPU 內(nèi)存的。了解內(nèi)存的使用去向有助于我們更好地規(guī)劃與優(yōu)化資源。在 LLMs 推理過程中,主要消耗 GPU 內(nèi)存的幾部分包括權(quán)重(模型參數(shù))、鍵值緩存內(nèi)存(Key-Value Cache Memory)、激活值(Activations)與臨時(shí)緩沖區(qū)(Temporary Buffers),以及系統(tǒng)開銷(Overheads)(如果你對(duì)并行處理或分布式計(jì)算有所研究,可能對(duì)這個(gè)概念已有一定的認(rèn)識(shí))。
1.1 模型參數(shù)(權(quán)重)
模型參數(shù)是神經(jīng)網(wǎng)絡(luò)在訓(xùn)練過程中學(xué)到的數(shù)值(權(quán)重(weights)和偏置(biases))。這些參數(shù)定義了模型如何處理輸入數(shù)據(jù)生成輸出。
模型大小對(duì) GPU 內(nèi)存的影響
- 直接關(guān)系:模型越大(參數(shù)越多),存儲(chǔ)這些權(quán)重所需的 GPU 內(nèi)存就越多。
- 內(nèi)存計(jì)算:在使用半精度(FP16)格式時(shí),每個(gè)參數(shù)通常需要 2 個(gè)字節(jié),這在推理過程中很常見,可以節(jié)省內(nèi)存而不會(huì)明顯損失精度。
讓我們看看這些模型:
- 擁有 34.5 億參數(shù)的小型 LLM:
- 所需內(nèi)存:34.5 億 × 2 字節(jié) = 69 MB。單 GPU 即可輕松支持。
- 現(xiàn)在如果使用 llama2-13b 模型:
- 擁有 130 億參數(shù),所需內(nèi)存將是:130 億 × 2 字節(jié) = 26 GB。這種情況下,需要一個(gè)擁有 40 GB內(nèi)存的 A100 GPU。
- 如果我們看看據(jù)說擁有 1750 億參數(shù)的 GPT-3 模型:
- 所需內(nèi)存:1750 億 × 2 字節(jié) = 350 GB,我們至少需要 9 個(gè) GPU 來僅存放模型權(quán)重。
請(qǐng)記住,對(duì)于 GPT-3 及其之后的模型,使用模型并行化(model parallelism)將模型參數(shù)分布到多個(gè) GPU 上是十分必要的。
1.2 鍵值(KV)緩存內(nèi)存
KV緩存存儲(chǔ)生成序列中每個(gè) token 所需的中間表示。簡單來說,當(dāng)模型每次生成一個(gè) token 時(shí),它需要記住之前的 tokens 以保持上下文。KV緩存存儲(chǔ)了目前為止生成的每個(gè) token 的鍵(key)和值(value)向量,使模型能夠高效地處理過去的 tokens ,而無需重新計(jì)算。
工作原理:
- Key 和 Values:在注意力機(jī)制中,模型為每個(gè) token 計(jì)算一個(gè)鍵向量和一個(gè)值向量。
- Storage:這些向量存儲(chǔ)在 KV 緩存中,并在后續(xù)步驟中用于生成新 tokens 。
序列長度(Sequence Length)和并發(fā)請(qǐng)求(Concurrent Requests)的影響:
- Longer Sequences:tokens 越多,KV緩存中的條目就越多,內(nèi)存使用量增加。
- Multiple Users:同時(shí)服務(wù)多個(gè)請(qǐng)求會(huì)成倍增加所需的KV緩存內(nèi)存。
計(jì)算每個(gè) token 的 KV 緩存大小
讓我們來分析一下如何得出每個(gè) token 的 KV 緩存大?。?/p>
- 每個(gè) token 的 KV 緩存組件:
鍵向量(每層一個(gè)鍵向量)和值向量(每層一個(gè)值向量)
- 每個(gè) token 的向量總數(shù):
模型層數(shù)(L)× 隱藏層大小(H):模型的深度 × 每個(gè)向量的維度。
再次以 llama-13b 模型為例,假設(shè)模型具有以下特征:
- 模型層數(shù)(L):40層
- 隱藏層大?。℉):5120維度(這種大小的模型中的常見維度)
- 計(jì)算每個(gè) token 占用的內(nèi)存:
i. 鍵向量:
- 總數(shù)量:L層 × H維度 = 40 × 5120 = 204,800 個(gè)
- 內(nèi)存大小:204,800 個(gè) × 2字節(jié)(FP16)= 409,600字節(jié)(或400 KB)
ii. 值向量:
- 與鍵向量相同:也是400 KB
iii. 每個(gè) token 的總KV緩存:
- 鍵向量 + 值向量:400 KB + 400 KB = 800 KB
現(xiàn)在考慮輸出內(nèi)容為2000個(gè) tokens 的情況:
800 KB/token × 2000 tokens = 每個(gè)輸出序列 1.6 GB
假如有 10 個(gè)并發(fā)請(qǐng)求(模型同時(shí)為 10 個(gè)用戶服務(wù)):1.6 GB/輸出序列 × 10 輸出序列 = 16 GB 的 KV 緩存內(nèi)存
KV緩存隨著序列長度和并發(fā)請(qǐng)求數(shù)量的增加而線性增長。我從一篇論文[1]中了解到,KV緩存可以消耗多達(dá) 30% 甚至更多的GPU內(nèi)存。
1.3 激活值和臨時(shí)緩沖區(qū)
激活值(Activations)是指推理過程中神經(jīng)網(wǎng)絡(luò)層的輸出,而臨時(shí)緩沖區(qū)(temporary buffers)用于中間計(jì)算。激活值和緩沖區(qū)通常消耗的內(nèi)存比模型權(quán)重和 KV 緩存要少。
它們可能使用大約 5-10% 的總 GPU 內(nèi)存。
盡管它們的容量較小,但激活值對(duì)于模型計(jì)算每一層的輸出是十分必要的。它們?cè)谇跋騻鬟f過程(forward pass)中被創(chuàng)建和丟棄,但仍需要足夠的內(nèi)存分配。
1.4 內(nèi)存開銷
額外的內(nèi)存使用開銷來自于內(nèi)存分配和使用的低效率。下面是對(duì)其的簡要介紹:
內(nèi)存碎片:
- Internal Fragmentation:當(dāng)分配的內(nèi)存塊沒有被完全利用時(shí)產(chǎn)生。
- External Fragmentation:隨著時(shí)間的推移,空閑內(nèi)存被分割成小塊,使得在需要時(shí)難以分配較大的連續(xù)內(nèi)存塊。
計(jì)算過程中產(chǎn)生的中間步驟:
- 臨時(shí)數(shù)據(jù):像矩陣乘法這樣的操作可能會(huì)創(chuàng)建消耗內(nèi)存的臨時(shí)張量。
低效內(nèi)存管理的影響:
- 性能降低:浪費(fèi)的內(nèi)存可能會(huì)限制系統(tǒng)可以處理的并發(fā)請(qǐng)求數(shù)量。
- 吞吐量降低:低效的內(nèi)存管理可能會(huì)導(dǎo)致延遲并降低向用戶提供服務(wù)響應(yīng)的整體速度。
示例:如果內(nèi)存碎片在 40 GB GPU 上浪費(fèi)了 20 %的內(nèi)存,那么就有 8 GB 的內(nèi)存本可以用來處理更多請(qǐng)求,但是現(xiàn)在被浪費(fèi)了。
02 計(jì)算 GPU 內(nèi)存需求
既然我們已經(jīng)對(duì)關(guān)鍵內(nèi)容有了足夠的了解,那么就不再拖延,直接計(jì)算完整的 GPU 內(nèi)存需求!
逐步計(jì)算:
要計(jì)算任何模型的內(nèi)存需求,幾乎以下內(nèi)容都需要:了解模型權(quán)重、KV緩存、激活值和臨時(shí)緩沖區(qū)以及系統(tǒng)內(nèi)存開銷。以 llama-2 13B 模型為例,公式為:
所需內(nèi)存總量:模型權(quán)重 + KV緩存 + 激活值和系統(tǒng)內(nèi)存開銷
對(duì)于 13 B 模型來說:
模型權(quán)重 = 參數(shù)數(shù)量 × 每個(gè)參數(shù)的字節(jié)數(shù)
總 KV 緩存內(nèi)存 = 每個(gè) token 的 KV 緩存內(nèi)存 × 輸出序列長度 × 輸出序列數(shù)量
激活值和系統(tǒng)內(nèi)存開銷 = GPU總內(nèi)存的 5–10 %
激活值和系統(tǒng)內(nèi)存開銷通常消耗模型參數(shù)和 KV 緩存使用的 GPU 總內(nèi)存的大約 5–10 %。你可以額外分配目前計(jì)算出的總內(nèi)存的 10 %作為這部分的內(nèi)存消耗預(yù)留量。
模型權(quán)重 = 130 億 × 2 字節(jié) = 26 GB
總 KV 緩存內(nèi)存 = 800 KB × 8192* tokens × 10* 并發(fā)請(qǐng)求 = 66 GB
激活值和系統(tǒng)內(nèi)存開銷 = 0.1 × (26 GB + 66GB) = 9.2 GB
*假設(shè)模型的輸出系列長度為 8192,有 10 個(gè)并行請(qǐng)求。
所需內(nèi)存總量:26 GB + 66 GB + 9.2 GB = 101.2 GB
所以,運(yùn)行 llama-2 7B 模型至少需要 3 個(gè) A100 40GB GPU。
如果我想要托管一個(gè) GPT-3 模型(我知道這很瘋狂;D),計(jì)算方法與此類似,但這次我會(huì)假設(shè)每次只處理一個(gè)請(qǐng)求,并使用 OPT-175B[2] 模型的大小( 96 層和每層 12288 維度)作為參考。
模型權(quán)重 = 1750 億 × 2 字節(jié) = 350 GB
總 KV 緩存內(nèi)存 = 4.5 MB × 8192 token × 1 并發(fā)請(qǐng)求 = 36 GB
激活值和系統(tǒng)內(nèi)存開銷 = 0.1 × (350 GB + 36GB) = 38.6 GB
所需總內(nèi)存:350 GB + 36 GB + 38.6 GB = 424.6 GB 幾乎需要 11 個(gè) A100 ??。
如果假設(shè) GPT-4 是一個(gè)擁有 1 萬億參數(shù)的模型,那么將需要 2.3 TB的內(nèi)存。
根據(jù)有關(guān)模型大小和參數(shù)的公開信息,計(jì)算出的內(nèi)存計(jì)算表如下所示:
Table calculated by Author
同樣,如果我將模型部署給許多用戶(比如 10 個(gè))同時(shí)使用,計(jì)算出的內(nèi)存計(jì)算表如下所示:
Table calculated by Author
在處理多個(gè)請(qǐng)求時(shí),內(nèi)存消耗明顯增加。主要是KV緩存大量增加,因?yàn)槟P蜋?quán)重和系統(tǒng)內(nèi)存開銷保持不變,KV 緩存會(huì)隨著 tokens 數(shù)量和并發(fā)請(qǐng)求的增加而大幅增加,矩陣的行數(shù)就會(huì)越多,從而直接增加內(nèi)存消耗。
現(xiàn)在想象一下 OpenAI[3] 或 Anthropic[4] 的大模型擁有數(shù)百萬用戶的情況吧?。?/p>
03 使用 GPU 內(nèi)存過程中遇到的挑戰(zhàn)及其優(yōu)化策略
經(jīng)過上述計(jì)算,我意識(shí)到如果不探討在部署大語言模型(LLMs)時(shí)遇到的一些挑戰(zhàn),以及目前的研究是如何針對(duì)這些問題進(jìn)行優(yōu)化的,那么這篇文章將顯得不夠完整。對(duì)于我們?cè)S多人來說,了解、掌握高效的 GPU 內(nèi)存管理技巧至關(guān)重要。下面我們簡要分析一下。
3.1 挑戰(zhàn)一:內(nèi)存碎片與內(nèi)存過度分配
在部署過程中,我們通常會(huì)靜態(tài)地為 KV cache 分配內(nèi)存,為每個(gè)請(qǐng)求預(yù)留盡可能大的內(nèi)存空間。這樣往往會(huì)導(dǎo)致內(nèi)存過度分配,因?yàn)楸M管實(shí)際的輸出序列往往更短,系統(tǒng)會(huì)為可能的最長輸出序列預(yù)留內(nèi)存空間。
此外,內(nèi)存碎片會(huì)降低有效的可用內(nèi)存,從而限制系統(tǒng)同時(shí)處理的請(qǐng)求數(shù)量。內(nèi)存碎片分為內(nèi)部(Internal)和外部(External)兩種。內(nèi)部內(nèi)存碎片是指分配的內(nèi)存塊未被充分利用,留下未使用的內(nèi)存空間。而外部內(nèi)存碎片則是指隨著時(shí)間推移,空閑內(nèi)存被分割成多個(gè)小的且不連續(xù)的內(nèi)存塊,這樣就難以在需要時(shí)分配足夠大的連續(xù)內(nèi)存塊。
內(nèi)存使用效率低下意味著 GPU 的計(jì)算資源并未得到充分利用。結(jié)果,系統(tǒng)受內(nèi)存限制而非計(jì)算能力的限制,浪費(fèi)了處理器性能。(這也是我們?cè)诓⑿谢蚍植际较到y(tǒng)中力求避免的問題)
3.2 挑戰(zhàn)二:解碼算法
大量的 LLM 應(yīng)用都傾向于采用先進(jìn)的解碼算法來優(yōu)化輸出質(zhì)量或是產(chǎn)生多樣化的輸出結(jié)果。盡管這些方法效果顯著,但它們也對(duì)內(nèi)存管理提出了新的挑戰(zhàn)。以束搜索(Beam Search)為例,該算法會(huì)生成多個(gè)備選輸出序列(即“束(beams)”),并根據(jù)評(píng)分標(biāo)準(zhǔn)保留得分最高的輸出序列。這意味著,每個(gè)“束(beams)”都需要專屬的 KV 緩存空間,從而增加了內(nèi)存使用量。同樣,Parallel Sampling 通過從模型生成的概率分布(probability distribution)中抽取樣本,一次性生成多個(gè)獨(dú)立輸出,每個(gè)輸出同樣需要獨(dú)立的 KV 緩存,這無疑進(jìn)一步增加了內(nèi)存消耗。
在動(dòng)態(tài)內(nèi)存分配這種情況下,解碼過程中“束(beams)”或樣本的數(shù)量可能會(huì)發(fā)生變化,從而導(dǎo)致不可預(yù)測(cè)的內(nèi)存需求。在不產(chǎn)生內(nèi)存碎片和過度內(nèi)存開銷的情況下,動(dòng)態(tài)地分配和釋放內(nèi)存,成為一項(xiàng)技術(shù)挑戰(zhàn)。此外,這些解碼方法可能會(huì)成倍增加內(nèi)存需求,有時(shí)甚至超出了 GPU 的處理能力。如果 GPU 內(nèi)存不足,系統(tǒng)可能不得不將數(shù)據(jù)轉(zhuǎn)移到速度較慢的 CPU 內(nèi)存或硬盤上,這無疑會(huì)延長處理時(shí)間。
面對(duì)這些問題,我們可能會(huì)思考:
我們?cè)撊绾瓮黄七@些限制?
3.3 PagedAttention
受操作系統(tǒng)內(nèi)存管理方式的啟發(fā),PagedAttention 技術(shù)將虛擬內(nèi)存的分頁原理應(yīng)用于 KV 緩存的管理。這種方法使得 KV 緩存數(shù)據(jù)不必占據(jù)一大塊連續(xù)的內(nèi)存空間,而是可以分散存儲(chǔ)于多個(gè)不連續(xù)的內(nèi)存頁面上。PagedAttention 采用動(dòng)態(tài)內(nèi)存分配策略,即根據(jù)實(shí)際需求為 KV 緩存分配內(nèi)存,無需提前預(yù)留出最大輸出序列長度所需的內(nèi)存。這樣的注意力機(jī)制能夠順暢地從不同內(nèi)存地址中檢索 KV 緩存數(shù)據(jù)。
PagedAttention 的優(yōu)勢(shì)在于,通過使用較小的內(nèi)存塊,有效減少了內(nèi)存碎片,降低了因內(nèi)存碎片導(dǎo)致的內(nèi)存浪費(fèi),從而提升了內(nèi)存的整體使用率。
3.4 vLLM
簡單來說,vLLM 是一個(gè)基于 PagedAttention 構(gòu)建的高吞吐量 LLM 服務(wù)系統(tǒng)。其核心目的是在推理過程中高效管理 GPU 內(nèi)存,尤其是 KV 緩存。從理論上看,vLLM 幾乎實(shí)現(xiàn)了零內(nèi)存浪費(fèi)的解決方案。通過內(nèi)存的動(dòng)態(tài)分配和非連續(xù)存儲(chǔ),它幾乎消除了內(nèi)存浪費(fèi)。理論上,它還支持在單個(gè)請(qǐng)求內(nèi)部和跨請(qǐng)求之間共享 KV 緩存數(shù)據(jù),這對(duì)于高級(jí)解碼方法尤其有用。在此基礎(chǔ)上,vLLM 能夠處理更大的批處理數(shù)量和更多的并發(fā)請(qǐng)求,從而提升整體性能。
即使進(jìn)行了優(yōu)化,有時(shí)也可能出現(xiàn) GPU 內(nèi)存不足的情況。vLLM 可通過 swapping 和 recomputation 來應(yīng)對(duì)這個(gè)問題。讓我們進(jìn)一步了解這一機(jī)制。
3.5 Swapping KV Cache to CPU Memory
- Swapping:當(dāng) GPU 內(nèi)存滿載時(shí),系統(tǒng)會(huì)將 KV 緩存數(shù)據(jù)從 GPU 內(nèi)存臨時(shí)轉(zhuǎn)移到 CPU 內(nèi)存中。
- 優(yōu)點(diǎn):
- 內(nèi)存釋放(Memory Relief):通過將數(shù)據(jù)移出 GPU 內(nèi)存,可以為新的請(qǐng)求騰出空間,確保 GPU 內(nèi)存不會(huì)因?yàn)橘Y源不足而阻礙新任務(wù)的執(zhí)行。
- 代價(jià):
- 延遲增加:由于 CPU 內(nèi)存的訪問速度通常低于 GPU 內(nèi)存,因此從 CPU 內(nèi)存讀取數(shù)據(jù)會(huì)比從 GPU 內(nèi)存讀取數(shù)據(jù)更加耗時(shí)。
- 數(shù)據(jù)傳輸開銷:在 GPU 內(nèi)存和 CPU 內(nèi)存之間轉(zhuǎn)移數(shù)據(jù)需要消耗帶寬和處理器時(shí)間。
3.6 Recomputation
不存儲(chǔ)所有 KV 緩存數(shù)據(jù),而是在需要時(shí)按需重新計(jì)算。
- 優(yōu)點(diǎn):
- 減少內(nèi)存使用:在內(nèi)存中需要存儲(chǔ)的數(shù)據(jù)量減少。
- 代價(jià):
- 增加計(jì)算量:重新計(jì)算數(shù)據(jù)需要額外的處理能力。
- 延遲影響:由于增加的額外計(jì)算量,可能會(huì)導(dǎo)致響應(yīng)時(shí)間變長。
Swapping 和 Recomputation 兩種方法比較表
Table by Author
單獨(dú)使用 Swapping 或 Recomputation 可能各有優(yōu)缺點(diǎn),但是將兩者結(jié)合起來使用,可以相互彌補(bǔ)對(duì)方的不足,從而在節(jié)省內(nèi)存、減少計(jì)算量、降低延遲等方面達(dá)到一個(gè)較為理想的平衡狀態(tài)。
Thanks for reading!
Hope you have enjoyed and learned new things from this blog!
About the authors
Muhammad Saad Uddin
Data Science is poetry of math where data is your love and its on you how you write your verses to show world your poetic expressions with utmost clarity?.
END
本期互動(dòng)內(nèi)容 ??
?你是否遇到過因?yàn)槟P瓦^大而導(dǎo)致的 GPU 內(nèi)存不足問題?
??文中鏈接??
[1]??https://arxiv.org/pdf/2309.06180??
[2]??https://medium.com/@plienhar/llm-inference-series-4-kv-caching-a-deeper-look-4ba9a77746c8??
[4]??https://www.anthropic.com/??
原文鏈接:
