PyTorch造大模型“加速包”,不到1000行代碼提速10倍!英偉達科學家:minGPT以來最好的教程式repo之一

2023-12-2 27 12/2

項目名爲GPT-fast

PyTorch團隊讓大模型推理速度加快了
10倍

且只用了不到
1000行的純原生PyTorch代碼

項目名爲
GPT-fast
,加速效果觀感是這樣嬸兒的:

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

通暢,屬實通暢!

重點是,團隊直接放出了代碼以及詳細“教程”。還是簡筆畫版的那種,特別好理解。

開發團隊成員@Horace He表示:

我們不把它看作是庫或者框架,更希望大家能把它當成個例子,根據自己的需求“複製粘貼”。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

網友直接炸開鍋,英偉達AI科學家Jim Fan評價道:

這是自Andrej Karpathy發佈的minGPT以來最棒的教程式repo之一!

開源世界需要更多minGPT、GPT-Fast這樣的項目!

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

那麼GPT-fast究竟是如何給大模型提速的?

開盒大模型“加速包”

總的來說,用到這幾種方法:

  • Torch.compile
    :一個專門爲PyTorch模型設計的編譯器,可以提升模型運行效率。
  • GPU量化
    :通過減少計算的精度來加速模型的運算速度。
  • 推測性解碼
    :使用一個較小的模型來預測較大模型的輸出,以此加快大語言模型的運算。
  • 張量並行性
    :通過在多個硬件設備上分佈模型的運算來加速處理速度。

下面我們來一一展開。

開發團隊一開始使用簡單的PyTorch來實現,但效果不佳(25.5 tok/s):

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

他們查看跟蹤後發現,一個原因是推理性能由於CPU過多佔用而受限。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

那麼如何解決呢?

可以想象這樣一個場景,GPU是一個龐大的工廠(擁有大量可用的算力),而CPU則是一個小推車,來回爲工廠“供貨”。

在很多情況下,CPU無法足夠快地“喂”GPU。

因此,開發團隊建議給GPU更多的工作量,或者說
一次性給它更大“塊”的任務
來處理。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

在推理過程中要做到這一點,可以引入
torch.compile

torch.compile能夠捕獲模型中更大的區域,並將其編譯成單一的編譯區域。特別是當以“reduce-overhead”模式運行時,它非常有效地減少了CPU的開銷。

效果立竿見影,性能直接提升了4倍,從25 tok/s提高到107 tok/s:

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

接下來,開發團隊想進一步提升速度,但遇到了
內存帶寬
瓶頸。

開發團隊計算了模型的帶寬利用率,結果已經達到了72%:

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

也就是說進一步提高速度的空間可能有限。

重新審視上面的方程式,團隊發現雖然實際上不能改變模型參數量,也不能改變GPU的內存帶寬(至少在不花更多錢的情況下),但可以改變存儲每個參數所用的
字節數

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

這意味着,雖然無法改變模型的大小或者升級硬件來提高性能,但可以通過減少存儲模型參數所需的數據量來提高效率。

通常可以通過量化技術來實現,即減少表示每個參數所需的位數。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

由此,開發團隊引入了下一個技術——
int8量化

採用int8權重量化減少了內存負載,進一步提升了性能(157.4 tok/s):

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

使用量化後還有0-一個問題:要生成100個token,必須加載(或調用)模型權重100次。頻繁加載模型權重也會導致效率低下。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

乍一看,好像沒有什麼解決的法子,因爲在自迴歸生成模式中存在着嚴格的序列依賴關係。

但開發團隊指出,通過利用推測性解碼可以打破這種嚴格的序列依賴關係。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

再來打個比方,想象有一個資深工程師Verity,他在技術決策上總是正確,但編寫代碼的速度相對較慢。

同時,還有一個初級工程師Drake,和Verity相反,不擅長技術決策,但編寫代碼的速度更快、成本也更低。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

那麼如何利用不同人的優勢來提高整體效率?

方法很簡單,先讓Drake編寫代碼,並在此過程中做出技術決策。接下來,將代碼交給Verity進行審查,不對的地方就讓Drake重做。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

在Transformer模型推理中,大型的驗證模型即爲Verity角色,Drake則是一個更小的、能更快生成文本的草稿模型。

開發團隊使用草稿模型生成8個token,然後使用驗證模型
並行處理
,丟棄不匹配的部分。

由此一來,打破了串行依賴,再次提高速度。

值得一提的是,推測性解碼不會改變輸出的質量。只要使用草稿模型生成token+驗證這些token所需的時間少於單獨生成這些token所需的時間,這種方法就是有效的。

而且使用原生PyTorch實現這種技術實際上非常簡單,整個實現過程只需要大約50行原生PyTorch代碼。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

由於AMD也支持Triton和torch.compile後端,因此之前在Nvidia GPU上應用的所有優化也可以在AMD GPU上重新應用。

開發團隊觀察到int8量化的加速從22 tok/s達到102 tok/s:

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

之後開發團隊又用了int4量化,進一步提升速度,但模型準確性有所下降。

因此使用了分組量化和GPTQ降低權重大小。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

最後在保證準確性的前提下,速度提升至202.1 tok/s:

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

將以上技術結合使用,達到更高速度244.7 tok/s:

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

到目前爲止,研發團隊一直都是在單個GPU上提速。但其實很多情況下是可以使用多個GPU的。

而使用多個GPU可以增加內存帶寬,從而提高模型的整體性能。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

在選擇並行處理策略時,需要在多個設備上分割一個token的處理過程,所以需要使用張量並行性。

而PyTorch也提供了用於張量並行性的底層工具,可以與torch.compile結合使用。

開發團隊還透露也正在開發用於表達張量並行性的更高級別的API。

然而,即使沒有更高級別的API,添加張量並行性也很容易,150行代碼即可實現,且不需要對模型進行任何改變。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

之前提到的所有優化都可以與張量並行性相結合。將這些優化結合起來,能夠以55 tokens/s的速度爲Llama-70B提供int8量化。

PyTorch造大模型“加速包”,不到1000行代码提速10倍!英伟达科学家:minGPT以来最好的教程式repo之一

最後總結成果,忽略量化,僅用766行代碼(model.py 244行代碼,generate.py 371行代碼,tp.py 151行代碼),就實現了快速推理、推測性解碼和張量並行性。

對於Llama-7B,使用compile+int4量化+推測性解碼速度達到241 tok/s。對於Llama-70B,通過加入張量並行性,達到80 tok/s。

這些性能都接近或超越了當前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