PyTorch 團隊讓大模型推理速度加快了 10 倍。且只用了不到 1000 行的純原生 PyTorch 代碼!
項目名為 GPT-fast,加速效果觀感是這樣嬸兒的:
通暢,屬實通暢!
重點是,團隊直接放出了代碼以及詳細“教程”。還是簡筆畫版的那種,特別好理解。
開發(fā)團隊成員 @Horace He 表示:
我們不把它看作是庫或者框架,更希望大家能把它當(dāng)成個例子,根據(jù)自己的需求“復(fù)制粘貼”。
網(wǎng)友直接炸開鍋,英偉達 AI 科學(xué)家 Jim Fan 評價道:
這是自 Andrej Karpathy 發(fā)布的 minGPT 以來最棒的教程式 repo 之一!
開源世界需要更多 minGPT、GPT-Fast 這樣的項目!
那么 GPT-fast 究竟是如何給大模型提速的?
開盒大模型“加速包”
總的來說,用到這幾種方法:
Torch.compile:一個專門為 PyTorch 模型設(shè)計的編譯器,可以提升模型運行效率。
GPU 量化:通過減少計算的精度來加速模型的運算速度。
推測性解碼:使用一個較小的模型來預(yù)測較大模型的輸出,以此加快大語言模型的運算。
張量并行性:通過在多個硬件設(shè)備上分布模型的運算來加速處理速度。
下面我們來一一展開。
開發(fā)團隊一開始使用簡單的 PyTorch 來實現(xiàn),但效果不佳(25.5 tok / s):
他們查看跟蹤后發(fā)現(xiàn),一個原因是推理性能由于 CPU 過多占用而受限。
那么如何解決呢?
可以想象這樣一個場景,GPU 是一個龐大的工廠(擁有大量可用的算力),而 CPU 則是一個小推車,來回為工廠“供貨”。
在很多情況下,CPU 無法足夠快地“喂”GPU。
因此,開發(fā)團隊建議給 GPU 更多的工作量,或者說一次性給它更大“塊”的任務(wù)來處理。
在推理過程中要做到這一點,可以引入 torch.compile。
torch.compile 能夠捕獲模型中更大的區(qū)域,并將其編譯成單一的編譯區(qū)域。特別是當(dāng)以“reduce-overhead”模式運行時,它非常有效地減少了 CPU 的開銷。
效果立竿見影,性能直接提升了 4 倍,從 25 tok / s 提高到 107 tok / s:
接下來,開發(fā)團隊想進一步提升速度,但遇到了內(nèi)存帶寬瓶頸。
開發(fā)團隊計算了模型的帶寬利用率,結(jié)果已經(jīng)達到了 72%:
也就是說進一步提高速度的空間可能有限。
重新審視上面的方程式,團隊發(fā)現(xiàn)雖然實際上不能改變模型參數(shù)量,也不能改變 GPU 的內(nèi)存帶寬(至少在不花更多錢的情況下),但可以改變存儲每個參數(shù)所用的字節(jié)數(shù)。
這意味著,雖然無法改變模型的大小或者升級硬件來提高性能,但可以通過減少存儲模型參數(shù)所需的數(shù)據(jù)量來提高效率。
通??梢酝ㄟ^量化技術(shù)來實現(xiàn),即減少表示每個參數(shù)所需的位數(shù)。
由此,開發(fā)團隊引入了下一個技術(shù) ——int8 量化。
采用 int8 權(quán)重量化減少了內(nèi)存負載,進一步提升了性能(157.4 tok / s):
使用量化后還有一個問題:要生成 100 個 token,必須加載(或調(diào)用)模型權(quán)重 100 次。頻繁加載模型權(quán)重也會導(dǎo)致效率低下。
乍一看,好像沒有什么解決的法子,因為在自回歸生成模式中存在著嚴(yán)格的序列依賴關(guān)系。
但開發(fā)團隊指出,通過利用推測性解碼可以打破這種嚴(yán)格的序列依賴關(guān)系。
再來打個比方,想象有一個資深工程師 Verity,他在技術(shù)決策上總是正確,但編寫代碼的速度相對較慢。
同時,還有一個初級工程師 Drake,和 Verity 相反,不擅長技術(shù)決策,但編寫代碼的速度更快、成本也更低。
那么如何利用不同人的優(yōu)勢來提高整體效率?
方法很簡單,先讓 Drake 編寫代碼,并在此過程中做出技術(shù)決策。接下來,將代碼交給 Verity 進行審查,不對的地方就讓 Drake 重做。
在 Transformer 模型推理中,大型的驗證模型即為 Verity 角色,Drake 則是一個更小的、能更快生成文本的草稿模型。
開發(fā)團隊使用草稿模型生成 8 個 token,然后使用驗證模型并行處理,丟棄不匹配的部分。
由此一來,打破了串行依賴,再次提高速度。
值得一提的是,推測性解碼不會改變輸出的質(zhì)量。只要使用草稿模型生成 token + 驗證這些 token 所需的時間少于單獨生成這些 token 所需的時間,這種方法就是有效的。
而且使用原生 PyTorch 實現(xiàn)這種技術(shù)實際上非常簡單,整個實現(xiàn)過程只需要大約 50 行原生 PyTorch 代碼。
由于 AMD 也支持 Triton 和 torch.compile 后端,因此之前在 Nvidia GPU 上應(yīng)用的所有優(yōu)化也可以在 AMD GPU 上重新應(yīng)用。
開發(fā)團隊觀察到 int8 量化的加速從 22 tok / s 達到 102 tok / s:
之后開發(fā)團隊又用了 int4 量化,進一步提升速度,但模型準(zhǔn)確性有所下降。
因此使用了分組量化和 GPTQ 降低權(quán)重大小。
最后在保證準(zhǔn)確性的前提下,速度提升至 202.1 tok / s:
將以上技術(shù)結(jié)合使用,達到更高速度 244.7 tok / s:
到目前為止,研發(fā)團隊一直都是在單個 GPU 上提速。但其實很多情況下是可以使用多個 GPU 的。
而使用多個 GPU 可以增加內(nèi)存帶寬,從而提高模型的整體性能。
在選擇并行處理策略時,需要在多個設(shè)備上分割一個 token 的處理過程,所以需要使用張量并行性。
而 PyTorch 也提供了用于張量并行性的底層工具,可以與 torch.compile 結(jié)合使用。
開發(fā)團隊還透露也正在開發(fā)用于表達張量并行性的更高級別的 API。
然而,即使沒有更高級別的 API,添加張量并行性也很容易,150 行代碼即可實現(xiàn),且不需要對模型進行任何改變。
之前提到的所有優(yōu)化都可以與張量并行性相結(jié)合。將這些優(yōu)化結(jié)合起來,能夠以 55 tokens / s 的速度為 Llama-70B 提供 int8 量化。
最后總結(jié)成果,忽略量化,僅用 766 行代碼(model.py 244 行代碼,generate.py 371 行代碼,tp.py 151 行代碼),就實現(xiàn)了快速推理、推測性解碼和張量并行性。
對于 Llama-7B,使用 compile+int4 量化 + 推測性解碼速度達到 241 tok / s。對于 Llama-70B,通過加入張量并行性,達到 80 tok / s。
這些性能都接近或超越了當(dāng)前 SOTA。
參考鏈接:
[1]https://pytorch.org/blog/accelerating-generative-ai-2/?utm_content=273712248&utm_medium=social&utm_source=twitter&hss_channel=tw-776585502606721024
[2]https://twitter.com/DrJimFan/status/1730298947376443698
[3]https://twitter.com/cHHillee/status/1730293330213531844
廣告聲明:文內(nèi)含有的對外跳轉(zhuǎn)鏈接(包括不限于超鏈接、二維碼、口令等形式),用于傳遞更多信息,節(jié)省甄選時間,結(jié)果僅供參考,IT之家所有文章均包含本聲明。