機器之心報道

機器之心編輯部

在一些大模型的推理任務(wù)上，瓶頸不是算力 FLOPS。

最近在開源社區(qū)，很多人都在探索大模型的優(yōu)化方法。有一個叫 llama.cpp 的項目用原始 C 重寫了 LLaMa 的推理代碼，效果極好，獲得了人們的廣泛關(guān)注。

通過一些優(yōu)化和量化權(quán)重，它能讓我們在各種以前無法想象的硬件上本地運行 LLaMa 模型。其中：

• 在谷歌 Pixel5 手機上，它能以 1 token/s 的速度運行 7B 參數(shù)模型。
• 在 M2 芯片的 MacBook Pro 上，使用 7B 參數(shù)模型的速度約為 16 token/s
• 我們甚至于可以在 4GB RAM 的樹莓派上運行 7B 模型，盡管速度只有 0.1 token/s

GitHub 鏈接：https://github.com/ggerganov/llama.cpp

我們知道，除了通用化能力，大模型落地的關(guān)鍵在于推理性能的優(yōu)化，然而如今這個優(yōu)化程度超出了我們的預(yù)料。llama.cpp 至今在 GitHub 上已經(jīng)收獲了 3.8 萬個 Star，幾乎和 LLaMa 模型本身一樣多。以至于到了 6 月份，llama.cpp 的作者 Georgi Gerganov 干脆開始創(chuàng)業(yè)，宣布創(chuàng)立一家新公司 ggml.ai，旨在用純 C 語言框架降低大模型運行成本。

很多人看到這里都會發(fā)問：這怎么可能？大語言模型不是需要英偉達(dá) H100 之類的 GPU 才能跑的嗎？為了解決這個疑惑，最近有人深入研究了圍繞大模型推理的數(shù)學(xué)，并試圖進行解答。

讓我們從「為什么 AI 訓(xùn)練都需要用 GPU？」開始，GPU 對深度學(xué)習(xí)有兩個主要好處：

• 它們具有很大的內(nèi)存帶寬（如 A100：1935 GB/s，RTX 4090：1008 GB/s）
• 它們具有很大的算力（A100：FP16 有 312 TFLOPS，RTX 4090：FP16 有 82.6 TFLOPS）

內(nèi)存帶寬之所以重要，是因為它關(guān)系到數(shù)據(jù)從 HBM 內(nèi)存（即 RAM）移動到片上內(nèi)存需要花費的時間。在實際使用 GPU 進行數(shù)學(xué)計算時，我們需要將相關(guān)矩陣移至片上內(nèi)存，該內(nèi)存相當(dāng)?。ˋ100 上為 40MB，而 RAM 為 40-80GB）。內(nèi)存帶寬比計算性能小約 2 個數(shù)量級 —— 這稍后會很重要，因為內(nèi)存帶寬往往是推理的瓶頸。

從計算機體系結(jié)構(gòu)的角度而言，我們需要把不同速度和容量的 memory 分出層級，以追求效率和成本之間的平衡。需要頻繁訪問的數(shù)據(jù)放在速度最快，但又容量最小的寄存器和 L1 cache 里，訪問量最少的數(shù)據(jù)放在最慢最大的內(nèi)存條里。

這在 LLaMa 推理任務(wù)上意味著什么？讓我們從一些推理數(shù)學(xué)計算開始。我們可以使用 Kipply 的文章（https://kipp.ly/transformer-param-count/）對 LLM 的推理性能進行一些粗略的計算。

首先有關(guān)模型尺寸：

• Q、K 和 V 權(quán)重矩陣的形狀都是 [ d_model, d_head]，每層有 n_heads；注意力輸出矩陣具有相同的形狀，總共 4 * [ d_model, n_heads * d_head]。按照慣例，GPT 風(fēng)格的網(wǎng)絡(luò)具有 d_head * n_heads = d_model。
• MLP 有兩個權(quán)重矩陣，形狀為 [d_model, 4 * d_model] 和 [4 * d_model，d_model]
• 嵌入矩陣的大小為 [d_vocab, d_model]。

這為我們提供了一個方便的類 GPT 模型參數(shù)數(shù)量方程：

在這里，我們將重點討論在本地運行類 ChatGPT 服務(wù)的情況，這就是 llama.cpp 所做的事情，讓我們假設(shè) batch size 為 1。為了高效推理，KV 緩存必須存儲在內(nèi)存中；KV 緩存需要存儲每一層的 KV 值，這相當(dāng)于存儲：

這里使用 n_bytes 來表示每個參數(shù)的字節(jié)數(shù)；對于 float32 是 4，對于 float16 是 2，以此類推。中間的 2 是因為我們必須為 K 值存儲一組權(quán)重，為 V 存儲一組權(quán)重。

給定一個 n 層模型，KV 緩存的總內(nèi)存為：

除了將 KV 緩存存儲在內(nèi)存中之外，我們還需要將權(quán)重本身存儲在內(nèi)存中；這需要 n_bytes * P 字節(jié)。

這是量化的主要優(yōu)點之一。通過使用較低的精度，我們可以從根本上減少存儲模型所需的內(nèi)存量。請注意，在 int4 精度下，所有這些模型都適合英偉達(dá)的 A100（也是目前數(shù)據(jù)中心里常見的 GPU）上的內(nèi)存，并且除了最大的模型之外，所有這些模型都適合高端消費級 GPU（如 RTX 3090/4090，具有 24GB RAM）。

現(xiàn)在，當(dāng)談到實際運行推理時，每個 token 大約需要 2P FLOPS，因為我們正在使用總共 P 個參數(shù)進行一系列矩陣乘法，與之相乘的矩陣尺寸是 (m, n) 向量 ( n,)，成本為 200 mn。

完成所有數(shù)學(xué)計算后，讓我們計算一下使用 LLaMa 運行推理的要求。sampling 的主要要求是：

• 除了所有參數(shù)之外，還將 KV 緩存保留在內(nèi)存中。
• 將 HBM 中的所有權(quán)重讀入片上存儲。因為我們是自回歸采樣，所以我們必須對采樣的每個 token 重復(fù)此操作。
• 進行實際的矩陣乘法來計算我們網(wǎng)絡(luò)的輸出。

延遲是計算延遲或內(nèi)存延遲的最大值，因為在所有現(xiàn)代張量編程庫中將參數(shù)讀取到片上內(nèi)存中都是異步發(fā)生的。因此，我們寫道：

其中 B 是 batch size。由于內(nèi)存帶寬約為 1.935e12，需要的 FLOPS 量約為 3.12e14，所以只要 batch size 小于 161，模型就會受到內(nèi)存限制。

當(dāng) batch size 為 1，即在計算機上僅生成單個預(yù)測流時，這是相同的等式，就像在大多數(shù)硬件（如英偉達(dá)的 GPU）上一樣，當(dāng)你降低精度時，會出現(xiàn)線性加速：使用 fp16 代替 fp32 時，F(xiàn)LOPS 會翻倍，轉(zhuǎn)到 int 8，F(xiàn)LOPS 會再增加一倍，用 int4 時再次加倍。

由于 llama.cpp 使用目前深度學(xué)習(xí)推理中較為激進的 int4 格式，因此 KV 緩存的 RAM 需求減少到 1.33GB，模型參數(shù)的 VRAM 減少到 16.25GB。這看起來很不錯

由于內(nèi)存帶寬幾乎總是遠(yuǎn)小于 FLOPS 數(shù)，因此內(nèi)存帶寬是瓶頸所在。

請注意，F(xiàn)LOPS/token 的數(shù)量與所需的內(nèi)存帶寬相同，因為我們必須 1) 將所有參數(shù)加載到片上內(nèi)存中，然后 2) 使用這些參數(shù)來計算結(jié)果。這些都是同時發(fā)生的，因為所有現(xiàn)代張量編程框架都能夠異步處理「加載到內(nèi)存」位，因此所需的總時間是 max（compute time, memory time）。

在英偉達(dá) A100 上運行 LLaMa

在 A100 (80GB PCIe) 上，內(nèi)存帶寬為 1935GB/s。int4 計算量為 1248 TOPS。因此，該模型較嚴(yán)重地受到內(nèi)存的限制。我們預(yù)計 65B 模型的速度約為 30 token/s，7B 模型的速度約為 277 token/s。

在 MacBook 上運行 LLaMa

接下來是正片了，蘋果 MacBook 上常見的 M1 芯片，其 GPU 的帶寬為 68.25 GB/s，而 M1 GPU 可執(zhí)行高達(dá) 5.5 TFLOPS 的 fp16 計算。因此，我們預(yù)計使用 int4 的 65B 模型采樣的上限為大約 1 token/s，使用 7B 模型的采樣上限為 10 token/s。

由于 M2 Pro 芯片具有 200 GB/s 的帶寬，而 M2 Max 具有 400 GB/s 的帶寬，因此我們應(yīng)該期待它們在這里可以獲得巨大的性能提升，使用 65B 版模型時 M2 Max 可以達(dá)到 6 token/s。這對于筆記本電腦來說已經(jīng)很不錯了。

在樹莓派 4 上運行 LLaMa

Raspberry Pi 4 具有 13.5 GFLOPS 的計算能力和約 4GB/s 的內(nèi)存帶寬。鑒于此，如果 7B 模型受內(nèi)存限制，我們預(yù)計會看到大約 2 token/s 的推理速度。然而我們目前看到的是約 0.1 token/s，有理由懷疑這實際上是因為算力受限導(dǎo)致的。這個嘗試是在不知道硬件性能的條件下進行的 —— 我們無法找到有關(guān) Raspberry Pi 低精度運算規(guī)格的足夠信息來確定這一點。

總結(jié)

內(nèi)存帶寬幾乎是與 transformer 采樣相關(guān)的最大限制因素。任何降低這些模型內(nèi)存需求的方法都會使它們更容易提供服務(wù) —— 比如量化！這是蒸餾（或者只是長時間訓(xùn)練較小的模型）非常重要的另一個原因。

OpenAI 科學(xué)家 Andrej Karpathy 對于這個觀察進行了進一步解釋。

他表示：除了并行推理和訓(xùn)練之外，提示編碼即使在 batch_size = 1 時也是可并行的，因為提示標(biāo)記可以由 LLM 并行編碼，而不是一一串行解碼。隨著提示越來越長，MacBook 的推理性能就會越落后于 A100。

但另一方面，蘋果的 M2 芯片看起來在大模型的推理任務(wù)上展示了強大的實力?！敢虼耍琈2 Ultra 是當(dāng)今體量最小、最漂亮、開箱即用、最簡單、最強大的個人 LLM 節(jié)點?！?/span>

陳天奇也對這種觀點表示贊同。

當(dāng)然這一切并不是免費的午餐。從本質(zhì)上講，使用低精度會損失一些準(zhǔn)確性，并且可能會出現(xiàn)一些奇怪的答案，讓大模型的回應(yīng)偏離軌道或產(chǎn)生幻覺。不過隨著模型參數(shù)越多，質(zhì)量損失就越低。因此，對于非常大的模型體量，差異或許可以忽略不計。此外，這只是推理成本。訓(xùn)練就完全是另一回事了。

通過對于各種性能參數(shù)的權(quán)衡，或許我們很快就會真正擁有更加「智能」的設(shè)備。

參考內(nèi)容：

https://finbarrtimbers.substack.com/p/how-is-llamacpp-possible

https://news.ycombinator.com/item?id=37140013

https://twitter.com/karpathy/status/1691571869051445433