PyTorch 團隊讓大模型推理速度加快了10 倍。
且隻用了不到1000 行的純原生 PyTorch 代碼！
項目名爲GPT-fast，加速效果觀感是這樣嬸兒的：
通暢，屬實通暢！
重點是，團隊直接放出了代碼以及詳細 " 教程 "。還是簡筆畫版的那種，特别好理解。
開發團隊成員 @Horace He 表示：
我們不把它看作是庫或者框架，更希望大家能把它當成個例子，根據自己的需求 " 複制粘貼 "。
網友直接炸開鍋，英偉達 AI 科學家 Jim Fan 評價道：
這是自 Andrej Karpathy 發布的 minGPT 以來最棒的教程式 repo 之一！
開源世界需要更多 minGPT、GPT-Fast 這樣的項目！
那麽 GPT-fast 究竟是如何給大模型提速的？
總的來說，用到這幾種方法：
Torch.compile：一個專門爲 PyTorch 模型設計的編譯器，可以提升模型運行效率。
GPU 量化：通過減少計算的精度來加速模型的運算速度。
推測性解碼：使用一個較小的模型來預測較大模型的輸出，以此加快大語言模型的運算。
張量并行性：通過在多個硬件設備上分布模型的運算來加速處理速度。
下面我們來一一展開。
開發團隊一開始使用簡單的 PyTorch 來實現，但效果不佳（25.5 tok/s）：
他們查看跟蹤後發現，一個原因是推理性能由于 CPU 過多占用而受限。
那麽如何解決呢？
可以想象這樣一個場景，GPU 是一個龐大的工廠（擁有大量可用的算力），而 CPU 則是一個小推車，來回爲工廠 " 供貨 "。
在很多情況下，CPU 無法足夠快地 " 喂 "GPU。
因此，開發團隊建議給 GPU 更多的工作量，或者說一次性給它更大 " 塊 " 的任務來處理。
在推理過程中要做到這一點，可以引入torch.compile。
torch.compile 能夠捕獲模型中更大的區域，并将其編譯成單一的編譯區域。特别是當以 "reduce-overhead" 模式運行時，它非常有效地減少了 CPU 的開銷。
效果立竿見影，性能直接提升了 4 倍，從 25 tok/s 提高到 107 tok/s：
接下來，開發團隊想進一步提升速度，但遇到了内存帶寬瓶頸。
開發團隊計算了模型的帶寬利用率，結果已經達到了 72%：
也就是說進一步提高速度的空間可能有限。
重新審視上面的方程式，團隊發現雖然實際上不能改變模型參數量，也不能改變 GPU 的内存帶寬（至少在不花更多錢的情況下），但可以改變存儲每個參數所用的字節數。
這意味着，雖然無法改變模型的大小或者升級硬件來提高性能，但可以通過減少存儲模型參數所需的數據量來提高效率。
通常可以通過量化技術來實現，即減少表示每個參數所需的位數。
由此，開發團隊引入了下一個技術——int8 量化。
采用 int8 權重量化減少了内存負載，進一步提升了性能（157.4 tok/s）：
使用量化後還有一個問題：要生成 100 個 token，必須加載（或調用）模型權重 100 次。頻繁加載模型權重也會導緻效率低下。
乍一看，好像沒有什麽解決的法子，因爲在自回歸生成模式中存在着嚴格的序列依賴關系。
但開發團隊指出，通過利用推測性解碼可以打破這種嚴格的序列依賴關系。
再來打個比方，想象有一個資深工程師 Verity，他在技術決策上總是正确，但編寫代碼的速度相對較慢。
同時，還有一個初級工程師 Drake，和 Verity 相反，不擅長技術決策，但編寫代碼的速度更快、成本也更低。
那麽如何利用不同人的優勢來提高整體效率？
方法很簡單，先讓 Drake 編寫代碼，并在此過程中做出技術決策。接下來，将代碼交給 Verity 進行審查，不對的地方就讓 Drake 重做。
在 Transformer 模型推理中，大型的驗證模型即爲 Verity 角色，Drake 則是一個更小的、能更快生成文本的草稿模型。
開發團隊使用草稿模型生成 8 個 token，然後使用驗證模型并行處理，丢棄不匹配的部分。
由此一來，打破了串行依賴，再次提高速度。
值得一提的是，推測性解碼不會改變輸出的質量。隻要使用草稿模型生成 token+ 驗證這些 token 所需的時間少于單獨生成這些 token 所需的時間，這種方法就是有效的。
而且使用原生 PyTorch 實現這種技術實際上非常簡單，整個實現過程隻需要大約 50 行原生 PyTorch 代碼。
由于 AMD 也支持 Triton 和 torch.compile 後端，因此之前在 Nvidia GPU 上應用的所有優化也可以在 AMD GPU 上重新應用。
開發團隊觀察到 int8 量化的加速從 22 tok/s 達到 102 tok/s：
之後開發團隊又用了 int4 量化，進一步提升速度，但模型準确性有所下降。
因此使用了分組量化和 GPTQ 降低權重大小。
最後在保證準确性的前提下，速度提升至 202.1 tok/s：
将以上技術結合使用，達到更高速度 244.7 tok/s：
到目前爲止，研發團隊一直都是在單個 GPU 上提速。但其實很多情況下是可以使用多個 GPU 的。
而使用多個 GPU 可以增加内存帶寬，從而提高模型的整體性能。
在選擇并行處理策略時，需要在多個設備上分割一個 token 的處理過程，所以需要使用張量并行性。
而 PyTorch 也提供了用于張量并行性的底層工具，可以與 torch.compile 結合使用。
開發團隊還透露也正在開發用于表達張量并行性的更高級别的 API。
然而，即使沒有更高級别的 API，添加張量并行性也很容易，150 行代碼即可實現，且不需要對模型進行任何改變。
之前提到的所有優化都可以與張量并行性相結合。将這些優化結合起來，能夠以 55 tokens/s 的速度爲 Llama-70B 提供 int8 量化。
最後總結成果，忽略量化，僅用 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
