研究完llama.cpp，我發(fā)現(xiàn)手機(jī)跑大模型竟這么簡(jiǎn)單（超大手機(jī)模型）

機(jī)器之心報(bào)道

機(jī)器之心編輯部

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

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

通過(guò)一些優(yōu)化和量化權(quán)重，它能讓我們?cè)诟鞣N以前無(wú)法想象的硬件上本地運(yùn)行 LLaMa 模型。其中：

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

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

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

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

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

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

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

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

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

首先有關(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 有兩個(gè)權(quán)重矩陣，形狀為 [d_model, 4 * d_model] 和 [4 * d_model，d_model]
• 嵌入矩陣的大小為 [d_vocab, d_model]。

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

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

這里使用 n_bytes 來(lái)表示每個(gè)參數(shù)的字節(jié)數(shù)；對(duì)于 float32 是 4，對(duì)于 float16 是 2，以此類(lèi)推。中間的 2 是因?yàn)槲覀儽仨殲?K 值存儲(chǔ)一組權(quán)重，為 V 存儲(chǔ)一組權(quán)重。

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

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

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

現(xiàn)在，當(dāng)談到實(shí)際運(yùn)行推理時(shí)，每個(gè) token 大約需要 2P FLOPS，因?yàn)槲覀冋谑褂每偣?P 個(gè)參數(shù)進(jìn)行一系列矩陣乘法，與之相乘的矩陣尺寸是 (m, n) 向量 ( n,)，成本為 200 mn。

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

• 除了所有參數(shù)之外，還將 KV 緩存保留在內(nèi)存中。
• 將 HBM 中的所有權(quán)重讀入片上存儲(chǔ)。因?yàn)槲覀兪亲曰貧w采樣，所以我們必須對(duì)采樣的每個(gè) token 重復(fù)此操作。
• 進(jìn)行實(shí)際的矩陣乘法來(lái)計(jì)算我們網(wǎng)絡(luò)的輸出。

延遲是計(jì)算延遲或內(nèi)存延遲的最大值，因?yàn)樵谒鞋F(xiàn)代張量編程庫(kù)中將參數(shù)讀取到片上內(nèi)存中都是異步發(fā)生的。因此，我們寫(xiě)道：

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

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

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

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

請(qǐng)注意，F(xiàn)LOPS/token 的數(shù)量與所需的內(nèi)存帶寬相同，因?yàn)槲覀儽仨?1) 將所有參數(shù)加載到片上內(nèi)存中，然后 2) 使用這些參數(shù)來(lái)計(jì)算結(jié)果。這些都是同時(shí)發(fā)生的，因?yàn)樗鞋F(xiàn)代張量編程框架都能夠異步處理「加載到內(nèi)存」位，因此所需的總時(shí)間是 max（compute time, memory time）。

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

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

在 MacBook 上運(yùn)行 LLaMa

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

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

在樹(shù)莓派 4 上運(yùn)行 LLaMa

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

總結(jié)

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

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

他表示：除了并行推理和訓(xùn)練之外，提示編碼即使在 batch_size = 1 時(shí)也是可并行的，因?yàn)樘崾緲?biāo)記可以由 LLM 并行編碼，而不是一一串行解碼。隨著提示越來(lái)越長(zhǎng)，MacBook 的推理性能就會(huì)越落后于 A100。

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

陳天奇也對(duì)這種觀點(diǎn)表示贊同。

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

通過(guò)對(duì)于各種性能參數(shù)的權(quán)衡，或許我們很快就會(huì)真正擁有更加「智能」的設(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