計(jì)算成本是人們打造 ChatGPT 等大模型面臨的重大挑戰(zhàn)之一。

據(jù)統(tǒng)計(jì)，從 GPT 進(jìn)化到 GPT-3 的過程也是模型體量增長(zhǎng)的過程 —— 參數(shù)量從 1.17 億增加到了 1750 億，預(yù)訓(xùn)練數(shù)據(jù)量從 5GB 增加到 45TB，其中 GPT-3 訓(xùn)練一次的費(fèi)用是 460 萬(wàn)美元，總訓(xùn)練成本達(dá) 1200 萬(wàn)美元。

除了訓(xùn)練，推理也很花錢。有人估算，現(xiàn)在 OpenAI 運(yùn)行 ChatGPT 的算力費(fèi)用每天就有 10 萬(wàn)美元。

在發(fā)展技術(shù)，讓大模型掌握更多能力的同時(shí)，也有人在嘗試降低 AI 所需的算力資源。最近，一種名為 FlexGen 的技術(shù)因?yàn)椤敢粔K RTX 3090 跑 ChatGPT 體量模型」而獲得了人們的關(guān)注。

雖然 FlexGen 加速后的大模型看起來仍然很慢 —— 跑 1750 億參數(shù)的語(yǔ)言模型時(shí)每秒 1 個(gè) token，但令人印象深刻的是，它已經(jīng)把不可能變成了可能。

傳統(tǒng)上，大語(yǔ)言模型（LLM）推理的高計(jì)算和內(nèi)存要求使人們必須使用多個(gè)高端 AI 加速器進(jìn)行訓(xùn)練。本研究探索了如何將 LLM 推理的要求降低到一個(gè)消費(fèi)級(jí) GPU 并實(shí)現(xiàn)實(shí)用性能。

近日，來自斯坦福大學(xué)、UC Berkeley、蘇黎世聯(lián)邦理工學(xué)院、Yandex、莫斯科國(guó)立高等經(jīng)濟(jì)學(xué)院、Meta、卡耐基梅隆大學(xué)等機(jī)構(gòu)的新研究提出了 FlexGen，這是一種用于運(yùn)行有限 GPU 內(nèi)存的 LLM 的高吞吐量生成引擎。

通過聚合來自 GPU、CPU 和磁盤的內(nèi)存和計(jì)算，F(xiàn)lexGen 可以在各種硬件資源限制下靈活配置。通過線性規(guī)劃優(yōu)化器，它搜索存儲(chǔ)和訪問張量的最佳模式，包括權(quán)重、激活和注意力鍵 / 值（KV）緩存。FlexGen 將權(quán)重和 KV 緩存進(jìn)一步壓縮到 4 位，精度損失低到可以忽略不計(jì)。與最先進(jìn)的 offloading 系統(tǒng)相比，F(xiàn)lexGen 在單個(gè) 16GB GPU 上運(yùn)行 OPT-175B 的速度提高了 100 倍，并首次實(shí)現(xiàn)了 1 token/s 的實(shí)際生成吞吐量。如果提供了更多的分布式 GPU，F(xiàn)lexGen 還帶有流水線并行 runtime，以允許在解碼時(shí)進(jìn)行超線性擴(kuò)展。

目前，該技術(shù)已經(jīng)放出代碼，獲得了幾千 Star 量：??https://github.com/FMInference/FlexGen??

簡(jiǎn)介

近年來，大語(yǔ)言模型在廣泛的任務(wù)中表現(xiàn)出卓越的性能。LLM 在展現(xiàn)出前所未有的通用智能的同時(shí)，也讓人們?cè)跇?gòu)建時(shí)面臨著前所未有的挑戰(zhàn)。這些模型可能有數(shù)十億甚至數(shù)萬(wàn)億個(gè)參數(shù)，這導(dǎo)致運(yùn)行它們需要極高的計(jì)算和內(nèi)存要求。例如，GPT-175B（GPT-3）僅用于存儲(chǔ)模型權(quán)重就需要 325GB 的內(nèi)存。要讓此模型進(jìn)行推理，至少需要五塊英偉達(dá) A100（80GB）和復(fù)雜的并行策略。

降低 LLM 推理資源需求的方法是最近人們經(jīng)常討論的內(nèi)容。這些努力分為三個(gè)方向：

（1）模型壓縮以減少總內(nèi)存占用量；

（2）協(xié)同推理，通過去中心化分?jǐn)偝杀荆?/span>

（3）Offloading 以利用 CPU 和磁盤的內(nèi)存。

這些技術(shù)顯著降低了使用 LLM 的計(jì)算資源需求。然而，人們通常假設(shè)模型適合 GPU 內(nèi)存，而現(xiàn)有的基于 offloading 的系統(tǒng)仍然難以使用單塊 GPU 以可接受的吞吐量運(yùn)行 1750 億參數(shù)規(guī)模的模型。

在新研究中，作者專注于高吞吐量生成推理的有效 offloading 策略。當(dāng) GPU 顯存不夠用時(shí)，我們需要將其卸載到二級(jí)存儲(chǔ)，通過部分加載的方式，逐段進(jìn)行計(jì)算。在典型的機(jī)器上，內(nèi)存層次結(jié)構(gòu)分為三級(jí)，如下圖所示。高級(jí)內(nèi)存速度快但稀缺，低級(jí)內(nèi)存速度慢但充裕。

在 FlexGen 中，作者不追求低延遲，而是瞄準(zhǔn)面向吞吐量的場(chǎng)景，這些場(chǎng)景在基準(zhǔn)測(cè)試、信息提取、數(shù)據(jù)整理等應(yīng)用中很受歡迎。實(shí)現(xiàn)低延遲對(duì)于 offloading 來說本質(zhì)上是一個(gè)挑戰(zhàn)，但是對(duì)于面向吞吐量的場(chǎng)景，可以大大提高 offloading 的效率。圖 1 說明了三個(gè)具有 offloading 的推理系統(tǒng)的延遲吞吐量權(quán)衡。通過仔細(xì)的調(diào)度，I/O 成本可以通過大量輸入分?jǐn)偛⑴c計(jì)算重疊。在研究中，作者展示了就單位算力成本而言，單塊消費(fèi)級(jí) GPU 吞吐量?jī)?yōu)化的 T4 GPU 效率要比云上延遲優(yōu)化的 8 塊 A100 GPU 的效率高 4 倍。

圖 1. OPT-175B（左）和 OPT-30B（右）上三個(gè)基于 offloading 的系統(tǒng)的延遲和吞吐量權(quán)衡。FlexGen 實(shí)現(xiàn)了新的帕累托最優(yōu)邊界，OPT-175B 的最大吞吐量提高了 100 倍。由于內(nèi)存不足，其他系統(tǒng)無(wú)法進(jìn)一步提高吞吐量。

盡管已有研究在訓(xùn)練的背景下討論了 offloading 的延遲 - 吞吐量權(quán)衡，但尚未有人將其用于生成 LLM 推理，這是一個(gè)截然不同的過程。由于 LLM 的自回歸性質(zhì)，生成推理提出了獨(dú)特的挑戰(zhàn)。除了存儲(chǔ)所有參數(shù)外，它還需要順序解碼并維護(hù)一個(gè)大的注意力鍵 / 值緩存（KV 緩存）。現(xiàn)有的 offload 系統(tǒng)都無(wú)法應(yīng)對(duì)這些挑戰(zhàn)，因此它們執(zhí)行過多的 I/O，只能實(shí)現(xiàn)遠(yuǎn)低于硬件能力的吞吐量。

為生成推理設(shè)計(jì)良好的 offloading 策略具有一定挑戰(zhàn)性。首先，這個(gè)過程中存在三種張量：權(quán)重、激活和 KV 緩存。該策略應(yīng)指定在三級(jí)層次結(jié)構(gòu)上的卸載內(nèi)容、位置以及卸載時(shí)機(jī)。其次，逐個(gè) batch、逐個(gè) token 和逐個(gè) layer 計(jì)算的結(jié)構(gòu)形成了一個(gè)復(fù)雜的依賴圖，可以通過多種方式進(jìn)行計(jì)算。該策略應(yīng)該選擇一個(gè)可以最小化執(zhí)行時(shí)間的時(shí)間表。這些選擇共同構(gòu)成了一個(gè)復(fù)雜的設(shè)計(jì)空間。

為此，在新方法 FlexGen 上，人們提出了一種用于 LLM 推理的 offloading 框架。FlexGen 聚合來自 GPU、CPU 和磁盤的內(nèi)存，并能有效地調(diào)度 I/O 操作，作者也討論了可能的壓縮方法和分布式管道并行性。

該研究的主要貢獻(xiàn)如下：

1、作者正式定義了可能的 offloading 策略的搜索空間，并使用成本模型和線性規(guī)劃求解器搜索最佳策略。值得關(guān)注的是，研究人員證明了搜索空間捕獲了一個(gè)幾乎 I/O 最優(yōu)的計(jì)算順序，其 I/O 復(fù)雜度在最優(yōu)計(jì)算順序的 2 倍以內(nèi)。搜索算法可以針對(duì)各種硬件規(guī)格和延遲 / 吞吐量限制進(jìn)行配置，從而提供一種平滑導(dǎo)航權(quán)衡空間的方法。與現(xiàn)有策略相比，F(xiàn)lexGen 解決方案統(tǒng)一了權(quán)重、激活和 KV 緩存的放置，從而實(shí)現(xiàn)了更大的 batch size。

2、研究表明，可以將 OPT-175B 等 LLM 的權(quán)重和 KV 緩存壓縮到 4 位，而無(wú)需重新訓(xùn)練 / 校準(zhǔn)，精度損失可忽略不計(jì)。這是通過細(xì)粒度分組量化實(shí)現(xiàn)的，可以顯著降低 I/O 成本。

3、通過在英偉達(dá) T4 GPU (16GB) 上運(yùn)行 OPT-175B 來展示 FlexGen 的效率。在單塊 GPU 上，給定相同的延遲要求，與 DeepSpeed Zero-Inference (Aminabadi et al., 2022) 和 Hugging Face Accelerate (HuggingFace, 2022) 相比，不壓縮的 FlexGen 可以實(shí)現(xiàn)高出 65 倍的吞吐量，后者是目前業(yè)內(nèi)最先進(jìn)的基于 offloading 的推理系統(tǒng)。如果允許更高的延遲和壓縮，F(xiàn)lexGen 可以進(jìn)一步提高吞吐量并達(dá)到 100 倍的改進(jìn)。FlexGen 是第一個(gè)可以使用單塊 T4 GPU 為 OPT-175B 實(shí)現(xiàn) 1 token/s 速度吞吐量的系統(tǒng)。如果給定多塊分布式 GPU，具有流水線并行性的 FlexGen 可在解碼時(shí)實(shí)現(xiàn)超線性擴(kuò)展。

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

運(yùn)行機(jī)制

通過聚合來自 GPU、CPU 和磁盤的內(nèi)存和計(jì)算，F(xiàn)lexGen 可以在各種硬件資源限制下靈活配置。通過線性規(guī)劃優(yōu)化器，它搜索存儲(chǔ)和訪問張量的最佳模式，包括權(quán)重、激活和注意力鍵 / 值 (KV) 緩存。FlexGen 將權(quán)重和 KV 緩存進(jìn)一步壓縮到 4 位，精度損失可以忽略不計(jì)。

FlexGen 的一個(gè)關(guān)鍵思想是進(jìn)行延遲 - 吞吐量權(quán)衡。實(shí)現(xiàn)低延遲對(duì)于卸載方法來說本來就具有挑戰(zhàn)性，但對(duì)于面向吞吐量的場(chǎng)景，可以極大地提升卸載效率（見下圖）。FlexGen 利用塊調(diào)度來重用權(quán)重并將 I/O 與計(jì)算重疊，如下圖 (b) 所示，而其他基線系統(tǒng)使用低效的逐行調(diào)度，如下圖 (a) 所示。

目前，該研究作者的下一步計(jì)劃包括對(duì)蘋果 M1、M2 芯片的支持和 Colab 部署的支持。

FlexGen 自發(fā)布后在 GitHub 上的 Star 量很快上千，在社交網(wǎng)絡(luò)上熱度也很高。人們紛紛表示這個(gè)項(xiàng)目很有前途，似乎運(yùn)行高性能大型語(yǔ)言模型的障礙正在被逐漸克服，希望在今年之內(nèi)，單機(jī)就能搞定 ChatGPT。

有人用這種方法訓(xùn)練了一個(gè)語(yǔ)言模型，結(jié)果如下：

雖然沒有經(jīng)過大量數(shù)據(jù)的投喂，AI 不知道具體知識(shí)，但回答問題的邏輯似乎比較清晰，或許未來的游戲中，我們能看見這樣的 NPC？