1750 億引數，只需要一塊 RTX 3090，ChatGPT 終於不再是大廠專屬的遊戲？

計算成本是人們打造 ChatGPT 等大模型面臨的重大挑戰之一。

據統計，從 GPT 進化到 GPT-3 的過程也是模型體量增長的過程 —— 引數量從 1.17 億增加到了 1750 億，預訓練資料量從 5GB 增加到 45TB，其中 GPT-3 訓練一次的費用是 460 萬美元，總訓練成本達 1200 萬美元。

除了訓練，推理也很花錢。有人估算，現在 OpenAI 執行 ChatGPT 的算力費用每天就有 10 萬美元。

在發展技術，讓大模型掌握更多能力的同時，也有人在嘗試降低 AI 所需的算力資源。最近，一種名為 FlexGen 的技術因為「一塊 RTX 3090 跑 ChatGPT 體量模型」而獲得了人們的關注。

雖然 FlexGen 加速後的大模型看起來仍然很慢 —— 跑 1750 億引數的語言模型時每秒 1 個 token，但令人印象深刻的是，它已經把不可能變成了可能。

傳統上，大語言模型（LLM）推理的高計算和記憶體要求使人們必須使用多個高階 AI 加速器進行訓練。本研究探索瞭如何將 LLM 推理的要求降低到一個消費級 GPU 並實現實用效能。

近日，來自史丹佛大學、UC Berkeley、蘇黎世聯邦理工學院、Yandex、莫斯科國立高等經濟學院、Meta、卡耐基梅隆大學等機構的新研究提出了 FlexGen，這是一種用於執行有限 GPU 記憶體的 LLM 的高吞吐量生成引擎。

透過聚合來自 GPU、CPU 和磁碟的記憶體和計算，FlexGen 可以在各種硬體資源限制下靈活配置。透過線性規劃最佳化器，它搜尋儲存和訪問張量的最佳模式，包括權重、啟用和注意力鍵 / 值（KV）快取。FlexGen 將權重和 KV 快取進一步壓縮到 4 位，精度損失低到可以忽略不計。與最先進的 offloading 系統相比，FlexGen 在單個 16GB GPU 上執行 OPT-175B 的速度提高了 100 倍，並首次實現了 1 token/s 的實際生成吞吐量。如果提供了更多的分散式 GPU，FlexGen 還帶有流水線並行 runtime，以允許在解碼時進行超線性擴充套件。

目前，該技術已經放出程式碼，獲得了幾千 Star 量：https://github.com/FMInference/FlexGen

簡介

近年來，大語言模型在廣泛的任務中表現出卓越的效能。LLM 在展現出前所未有的通用智慧的同時，也讓人們在構建時面臨著前所未有的挑戰。這些模型可能有數十億甚至數萬億個引數，這導致執行它們需要極高的計算和記憶體要求。例如，GPT-175B（GPT-3）僅用於儲存模型權重就需要 325GB 的記憶體。要讓此模型進行推理，至少需要五塊英偉達 A100（80GB）和複雜的並行策略。

降低 LLM 推理資源需求的方法是最近人們經常討論的內容。這些努力分為三個方向：

（1）模型壓縮以減少總記憶體佔用量；

（2）協同推理，透過去中心化分攤成本；

（3）Offloading 以利用 CPU 和磁碟的記憶體。

這些技術顯著降低了使用 LLM 的計算資源需求。然而，人們通常假設模型適合 GPU 記憶體，而現有的基於 offloading 的系統仍然難以使用單塊 GPU 以可接受的吞吐量執行 1750 億引數規模的模型。

在新研究中，作者專注於高吞吐量生成推理的有效 offloading 策略。當 GPU 視訊記憶體不夠用時，我們需要將其解除安裝到二級儲存，透過部分載入的方式，逐段進行計算。在典型的機器上，記憶體層次結構分為三級，如下圖所示。高階記憶體速度快但稀缺，低階記憶體速度慢但充裕。

在 FlexGen 中，作者不追求低延遲，而是瞄準面向吞吐量的場景，這些場景在基準測試、資訊提取、資料整理等應用中很受歡迎。實現低延遲對於 offloading 來說本質上是一個挑戰，但是對於面向吞吐量的場景，可以大大提高 offloading 的效率。圖 1 說明了三個具有 offloading 的推理系統的延遲吞吐量權衡。透過仔細的排程，I/O 成本可以透過大量輸入分攤並與計算重疊。在研究中，作者展示了就單位算力成本而言，單塊消費級 GPU 吞吐量最佳化的 T4 GPU 效率要比雲上延遲最佳化的 8 塊 A100 GPU 的效率高 4 倍。

圖 1. OPT-175B（左）和 OPT-30B（右）上三個基於 offloading 的系統的延遲和吞吐量權衡。FlexGen 實現了新的帕累托最優邊界，OPT-175B 的最大吞吐量提高了 100 倍。由於記憶體不足，其他系統無法進一步提高吞吐量。

儘管已有研究在訓練的背景下討論了 offloading 的延遲 – 吞吐量權衡，但尚未有人將其用於生成 LLM 推理，這是一個截然不同的過程。由於 LLM 的自迴歸性質，生成推理提出了獨特的挑戰。除了儲存所有引數外，它還需要順序解碼並維護一個大的注意力鍵 / 值快取（KV 快取）。現有的 offload 系統都無法應對這些挑戰，因此它們執行過多的 I/O，只能實現遠低於硬體能力的吞吐量。

為生成推理設計良好的 offloading 策略具有一定挑戰性。首先，這個過程中存在三種張量：權重、啟用和 KV 快取。該策略應指定在三級層次結構上的解除安裝內容、位置以及解除安裝時機。其次，逐個 batch、逐個 token 和逐個 layer 計算的結構形成了一個複雜的依賴圖，可以透過多種方式進行計算。該策略應該選擇一個可以最小化執行時間的時間表。這些選擇共同構成了一個複雜的設計空間。

為此，在新方法 FlexGen 上，人們提出了一種用於 LLM 推理的 offloading 框架。FlexGen 聚合來自 GPU、CPU 和磁碟的記憶體，並能有效地排程 I/O 操作，作者也討論了可能的壓縮方法和分散式管道並行性。

該研究的主要貢獻如下：

1、作者正式定義了可能的 offloading 策略的搜尋空間，並使用成本模型和線性規劃求解器搜尋最佳策略。值得關注的是，研究人員證明了搜尋空間捕獲了一個幾乎 I/O 最優的計算順序，其 I/O 複雜度在最優計算順序的 2 倍以內。搜尋演算法可以針對各種硬體規格和延遲 / 吞吐量限制進行配置，從而提供一種平滑導航權衡空間的方法。與現有策略相比，FlexGen 解決方案統一了權重、啟用和 KV 快取的放置，從而實現了更大的 batch size。

2、研究表明，可以將 OPT-175B 等 LLM 的權重和 KV 快取壓縮到 4 位，而無需重新訓練 / 校準，精度損失可忽略不計。這是透過細粒度分組量化實現的，可以顯著降低 I/O 成本。

3、透過在英偉達 T4 GPU (16GB) 上執行 OPT-175B 來展示 FlexGen 的效率。在單塊 GPU 上，給定相同的延遲要求，與 DeepSpeed Zero-Inference (Aminabadi et al., 2022) 和 Hugging Face Accelerate (HuggingFace, 2022) 相比，不壓縮的 FlexGen 可以實現高出 65 倍的吞吐量，後者是目前業內最先進的基於 offloading 的推理系統。如果允許更高的延遲和壓縮，FlexGen 可以進一步提高吞吐量並達到 100 倍的改進。FlexGen 是第一個可以使用單塊 T4 GPU 為 OPT-175B 實現 1 token/s 速度吞吐量的系統。如果給定多塊分散式 GPU，具有流水線並行性的 FlexGen 可在解碼時實現超線性擴充套件。

在研究中，作者還將 FlexGen 和 Petals 作為 offloading 和去中心化集合推理方法的代表進行了比較。結果表明，具有單塊 T4 GPU 的 FlexGen 在吞吐量方面勝過具有 12 塊 T4 GPU 的分散式 Petal 叢集，並且在某些情況下甚至可以實現更低的延遲。

執行機制

透過聚合來自 GPU、CPU 和磁碟的記憶體和計算，FlexGen 可以在各種硬體資源限制下靈活配置。透過線性規劃最佳化器，它搜尋儲存和訪問張量的最佳模式，包括權重、啟用和注意力鍵 / 值 (KV) 快取。FlexGen 將權重和 KV 快取進一步壓縮到 4 位，精度損失可以忽略不計。

FlexGen 的一個關鍵思想是進行延遲 – 吞吐量權衡。實現低延遲對於解除安裝方法來說本來就具有挑戰性，但對於面向吞吐量的場景，可以極大地提升解除安裝效率（見下圖）。FlexGen 利用塊排程來重用權重並將 I/O 與計算重疊，如下圖 (b) 所示，而其他基線系統使用低效的逐行排程，如下圖 (a) 所示。

目前，該研究作者的下一步計劃包括對蘋果 M1、M2 晶片的支援和 Colab 部署的支援。

FlexGen 自發布後在 GitHub 上的 Star 量很快上千，在社交網路上熱度也很高。人們紛紛表示這個專案很有前途，似乎執行高效能大型語言模型的障礙正在被逐漸克服，希望在今年之內，單機就能搞定 ChatGPT。

有人用這種方法訓練了一個語言模型，結果如下：

雖然沒有經過大量資料的投餵，AI 不知道具體知識，但回答問題的邏輯似乎比較清晰，或許未來的遊戲中，我們能看見這樣的 NPC？

參考內容：https://news.ycombinator.com/item?id=34869960

來自：機器學習研究組訂閱