PyTorch團(tuán)隊(duì)讓大模型推理速度加快了10倍。

且只用了不到1000行的純?cè)鶳yTorch代碼！

項(xiàng)目名為GPT-fast，加速效果觀感是這樣?jì)饍旱模?/p>

通暢，屬實(shí)通暢！

重點(diǎn)是，團(tuán)隊(duì)直接放出了代碼以及詳細(xì)“教程”。還是簡筆畫版的那種，特別好理解。

開發(fā)團(tuán)隊(duì)成員@Horace He表示：

我們不把它看作是庫或者框架，更希望大家能把它當(dāng)成個(gè)例子，根據(jù)自己的需求“復(fù)制粘貼”。

網(wǎng)友直接炸開鍋，英偉達(dá)AI科學(xué)家Jim Fan評(píng)價(jià)道：

這是自Andrej Karpathy發(fā)布的minGPT以來最棒的教程式repo之一!

開源世界需要更多minGPT、GPT-Fast這樣的項(xiàng)目!

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

開盒大模型“加速包”

總的來說，用到這幾種方法：

• Torch.compile：一個(gè)專門為PyTorch模型設(shè)計(jì)的編譯器，可以提升模型運(yùn)行效率。
• GPU量化：通過減少計(jì)算的精度來加速模型的運(yùn)算速度。
• 推測性解碼：使用一個(gè)較小的模型來預(yù)測較大模型的輸出，以此加快大語言模型的運(yùn)算。
• 張量并行性：通過在多個(gè)硬件設(shè)備上分布模型的運(yùn)算來加速處理速度。

下面我們來一一展開。

開發(fā)團(tuán)隊(duì)一開始使用簡單的PyTorch來實(shí)現(xiàn)，但效果不佳（25.5 tok/s）：

他們查看跟蹤后發(fā)現(xiàn)，一個(gè)原因是推理性能由于CPU過多占用而受限。

那么如何解決呢？

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

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

因此，開發(fā)團(tuán)隊(duì)建議給GPU更多的工作量，或者說一次性給它更大“塊”的任務(wù)來處理。

在推理過程中要做到這一點(diǎn)，可以引入torch.compile。

torch.compile能夠捕獲模型中更大的區(qū)域，并將其編譯成單一的編譯區(qū)域。特別是當(dāng)以“reduce-overhead”模式運(yùn)行時(shí)，它非常有效地減少了CPU的開銷。

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

接下來，開發(fā)團(tuán)隊(duì)想進(jìn)一步提升速度，但遇到了內(nèi)存帶寬瓶頸。

開發(fā)團(tuán)隊(duì)計(jì)算了模型的帶寬利用率，結(jié)果已經(jīng)達(dá)到了72%：

也就是說進(jìn)一步提高速度的空間可能有限。

重新審視上面的方程式，團(tuán)隊(duì)發(fā)現(xiàn)雖然實(shí)際上不能改變模型參數(shù)量，也不能改變GPU的內(nèi)存帶寬（至少在不花更多錢的情況下），但可以改變存儲(chǔ)每個(gè)參數(shù)所用的字節(jié)數(shù)。

這意味著，雖然無法改變模型的大小或者升級(jí)硬件來提高性能，但可以通過減少存儲(chǔ)模型參數(shù)所需的數(shù)據(jù)量來提高效率。

通常可以通過量化技術(shù)來實(shí)現(xiàn)，即減少表示每個(gè)參數(shù)所需的位數(shù)。

由此，開發(fā)團(tuán)隊(duì)引入了下一個(gè)技術(shù)——int8量化。

采用int8權(quán)重量化減少了內(nèi)存負(fù)載，進(jìn)一步提升了性能（157.4 tok/s）：

使用量化后還有一個(gè)問題：要生成100個(gè)token，必須加載（或調(diào)用）模型權(quán)重100次。頻繁加載模型權(quán)重也會(huì)導(dǎo)致效率低下。

乍一看，好像沒有什么解決的法子，因?yàn)樵谧曰貧w生成模式中存在著嚴(yán)格的序列依賴關(guān)系。

但開發(fā)團(tuán)隊(duì)指出，通過利用推測性解碼可以打破這種嚴(yán)格的序列依賴關(guān)系。

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

同時(shí)，還有一個(gè)初級(jí)工程師Drake，和Verity相反，不擅長技術(shù)決策，但編寫代碼的速度更快、成本也更低。

那么如何利用不同人的優(yōu)勢(shì)來提高整體效率？

方法很簡單，先讓Drake編寫代碼，并在此過程中做出技術(shù)決策。接下來，將代碼交給Verity進(jìn)行審查，不對(duì)的地方就讓Drake重做。

在Transformer模型推理中，大型的驗(yàn)證模型即為Verity角色，Drake則是一個(gè)更小的、能更快生成文本的草稿模型。

開發(fā)團(tuán)隊(duì)使用草稿模型生成8個(gè)token，然后使用驗(yàn)證模型并行處理，丟棄不匹配的部分。

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

值得一提的是，推測性解碼不會(huì)改變輸出的質(zhì)量。只要使用草稿模型生成token+驗(yàn)證這些token所需的時(shí)間少于單獨(dú)生成這些token所需的時(shí)間，這種方法就是有效的。

而且使用原生PyTorch實(shí)現(xiàn)這種技術(shù)實(shí)際上非常簡單，整個(gè)實(shí)現(xiàn)過程只需要大約50行原生PyTorch代碼。

由于AMD也支持Triton和torch.compile后端，因此之前在Nvidia GPU上應(yīng)用的所有優(yōu)化也可以在AMD GPU上重新應(yīng)用。

開發(fā)團(tuán)隊(duì)觀察到int8量化的加速從22 tok/s達(dá)到102 tok/s：

之后開發(fā)團(tuán)隊(duì)又用了int4量化，進(jìn)一步提升速度，但模型準(zhǔn)確性有所下降。

因此使用了分組量化和GPTQ降低權(quán)重大小。

最后在保證準(zhǔn)確性的前提下，速度提升至202.1 tok/s：

將以上技術(shù)結(jié)合使用，達(dá)到更高速度244.7 tok/s：

到目前為止，研發(fā)團(tuán)隊(duì)一直都是在單個(gè)GPU上提速。但其實(shí)很多情況下是可以使用多個(gè)GPU的。

而使用多個(gè)GPU可以增加內(nèi)存帶寬，從而提高模型的整體性能。

在選擇并行處理策略時(shí)，需要在多個(gè)設(shè)備上分割一個(gè)token的處理過程，所以需要使用張量并行性。

而PyTorch也提供了用于張量并行性的底層工具，可以與torch.compile結(jié)合使用。

開發(fā)團(tuán)隊(duì)還透露也正在開發(fā)用于表達(dá)張量并行性的更高級(jí)別的API。

然而，即使沒有更高級(jí)別的API，添加張量并行性也很容易，150行代碼即可實(shí)現(xiàn)，且不需要對(duì)模型進(jìn)行任何改變。

之前提到的所有優(yōu)化都可以與張量并行性相結(jié)合。將這些優(yōu)化結(jié)合起來，能夠以55 tokens/s的速度為Llama-70B提供int8量化。

最后總結(jié)成果，忽略量化，僅用766行代碼（model.py 244行代碼，generate.py 371行代碼，tp.py 151行代碼），就實(shí)現(xiàn)了快速推理、推測性解碼和張量并行性。

對(duì)于Llama-7B，使用compile+int4量化+推測性解碼速度達(dá)到241 tok/s。對(duì)于Llama-70B，通過加入張量并行性，達(dá)到80 tok/s。

這些性能都接近或超越了當(dāng)前SOTA。

參考鏈接：

[1]https://pytorch.org/blog/accelerating-generative-ai-2/?utm_cnotallow=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