本文展示了如何使用 1760 億 (176B) 參數(shù)的?BLOOM 模型[1]?生成文本時如何獲得超快的詞吞吐 (per token throughput)。

因為在使用 bf16 (bfloat16) 權(quán)重時該模型內(nèi)存占用為 352 GB (176*2)，所以最高效的硬件配置是使用 8x80GB 的 A100 GPU。也可使用 2x8x40GB 的 A100 或者 2x8x48GB 的 A6000。使用這些 GPU 的主要原因是截至本文成稿時為止它們是能提供最大顯存的 GPU，但你也可以使用其他 GPU。比如，可以使用 24x32GB V100。

一般來講，使用單節(jié)點會帶來最快的吞吐，因為大多數(shù)時候節(jié)點內(nèi)的 GPU 互聯(lián)硬件比節(jié)點間的快，但未必總是如此。

如果你沒有這么高端的硬件或沒有這么多硬件，你仍可能通過 CPU 卸載 (CPU offload) 或是 NVMe 卸載 (NVMe offload) 的方式在更小的 GPU 上對 BLOOM 進行推理。當然，生成時間會慢很多。

我們計劃涉及?8 比特量化方案[2]，該方案以稍慢的吞吐為代價將顯存需求減少到一半。我們還會討論?BitsAndBytes[3]?和?Deepspeed-Inference[4]?庫。

測試基準

事不宜遲，我們先展示一些數(shù)據(jù)吧。

為了保持一致性，除非另有說明，本文的測試基準都是在相同的配有 512GB CPU 內(nèi)存的 8x80GB A100 節(jié)點上完成的，該節(jié)點來自?法國 Jean Zay 超算中心[5]

所有的測試基準都是使用貪心搜索完成最多 100 個詞的生成任務(wù):

輸入提示詞僅包含幾個詞。我們會緩存先前見到的詞，因為每次重新計算它們相當慢。

首先，讓我們快速看一下從開始到準備好花了多長時間， 即模型加載和準備花了多長時間：

Deepspeed-Inference 使用了預分片的權(quán)重倉庫，整個加載時間大約在 1 分鐘。Accelerrate 的加載時間也很優(yōu)秀，只有大約 2 分鐘。其他方案就慢得多。

加載時間有可能重要也可能并不重要，因為一旦加載成功你可以一遍遍持續(xù)不斷地生成詞而不再需要額外地加載開銷。

接著是最重要的測試基準指標：詞生成吞吐 (token generation throughput)。這個吞吐的度量比較簡單，即：生成 100 個新詞的時間除以 100 和 batch size (也就是除以生成的總詞數(shù))。

下面列出了 8x80GB GPU 的吞吐，單位為毫秒：

這里， 當內(nèi)存耗盡 (Out Of Memory，OOM) 時即表明 batch size 太大 GPU 顯存放不下了。

使用 Deepspeed-Inference 的張量并行 (Tensor Parallelism，TP) 和定制化融合 CUDA 核函數(shù)可以得到小于 1 毫秒的吞吐！太棒了！盡管使用這個方案去推理那些尚未被驗證過的模型時，你可能會需要花一些時間去開發(fā)從而讓它工作起來。

Accelerate 也超級快。它使用了非常簡單的管線并行 (Pipeline Parallelism，PP)。因為它非常簡單，所以它應該對任何模型都是開箱即用的。

因為 Deepspeed-ZeRO 可以并行處理多路生成流，其吞吐可以再除以 8 或者 16，具體數(shù)值取決于在調(diào)用?generate?時用了 8 個 GPU 還是 16 個 GPU。當然，這也意味著在 8x80GB A100 的情況下 (見上表) ，可以處理的 batch size 為 64 且吞吐可至大約 4 毫秒。因此，這 3 種方案的性能是接近的。

讓我們再重新看一下這些數(shù)字是怎么計算出來的。舉個例子，使用 Deepspeed-Inference fp16 模式實時生成 batch size 為 128、長度為 100 個新詞的文本花了 8832 毫秒，因此我們可以這樣計算吞吐：鐘面時間 / ( batch size * 新詞數(shù) ) 或?8821/(128*100) = 0.69。

現(xiàn)在我們一起看看 Deepspeed-Inference 和 BitsAndBytes 提供的 int8 量化模型的威力，它僅需占用 bfloat16 或 float16 推理所需顯存的一半。

以下為 4x80GB GPU 的吞吐，單位為毫秒：

你只需在下述 3 個腳本里添加?--benchmark?即可重現(xiàn)這些測試基準的結(jié)果。

方案

首先獲取最新的演示代碼倉庫：

本文我們準備使用?bloom-inference-scripts/?文件夾下的 3 個腳本。

下面我們按框架的字母序逐一展示相關(guān)方案。

HuggingFace Accelerate

Accelerate 按如下步驟進行大模型推理：

1. 用空的權(quán)重實例化模型。

2. 分析每層的大小以及每個設(shè)備 (CPU, CPU) 的可用空間，并決定每層應該在哪個設(shè)備上推理。

3. 逐比特加載模型 checkpoint 并把權(quán)重加載到相應的設(shè)備。

然后，它會使用鉤子代碼 (hook) 來確保模型正確運行，鉤子代碼被用于在正確的設(shè)備間傳輸輸入和輸出，并在前向輪運行前加載那些卸載到 CPU (甚至硬盤) 上的權(quán)重到 GPU，然后在前向輪結(jié)束后再次把權(quán)重卸載。

在有多個 GPU 且有足夠空間放下整個模型的情形下，該方案在 GPU 間逐個切換直至所有層運行完畢。每個給定的時間只有一個 GPU 工作，這聽起來很沒效率。但盡管該方案 GPU 存在空閑，它的吞吐卻相當不錯。

因為相同的代碼可以運行在任意給定的設(shè)置中，所以本方案非常靈活。Accelerate 首先使用所有可用的 GPU，當顯存已滿時會卸載到 CPU 內(nèi)存直至卸載到硬盤。卸載到 CPU 或硬盤會讓推理變慢。舉個例子，與 8x80 A100 上的 10 毫秒相比，已有用戶報告，不作任何代碼改動，在 2 個 A100 上運行 BLOOM 吞吐是每詞 15 秒。

你可以你從?Accelerate 文檔[6]?中獲取本方案的更多信息。

設(shè)置

運行

簡單執(zhí)行如下命令。

如需使用?BitsAndBytes[7]?的 8 比特量化方案，首先要安裝?bitsandbytes：

然后在前述命令行中增加?--dtype int8。

如果你有 4 個以上 GPU，你可以通過如下命令限制腳本只使用其中 4 個 GPU：

在這個例子中，不 OOM 的最大 batch size 是 40。如果你深入研究腳本，你會看到我們需要調(diào)整顯存分配映射從而把第一個 GPU 解放出來去僅處理激活和先前詞的緩存。

DeepSpeed-Inference

DeepSpeed-Inference[8]?使用張量并行 (Tensor Parallelism) 以及高效的融合 CUDA 核函數(shù)在 128 這個大 batch size 下達到了每詞 1 毫秒的超快推理性能。

設(shè)置

運行

1.最快的方法是使用 TP 預分片 (TP = Tensor Parallel) 的 checkpoint，與非預分片的 bloom checkpoint 相比，它僅需大約 1 分鐘即可加載：

1a. 如果你想要運行原始 bloom checkpoint，這個 checkpoint 一旦加載就會跟之前的方案跑到相同的吞吐，但加載需要花 10 到 20 分鐘。

2a. 8 比特量化版本與一般的半精度版本相比僅需一半 GPU 顯存。

這里我們使用?microsoft/bloom-deepspeed-inference-int8?checkpoint 并告訴腳本跑在?int8?模式。

當然，現(xiàn)在僅需 4x80GB A100 GPU 就夠了：

這種情況下，不 OOM 的最大 batch size 是 128。

可以看到，本方案中有兩個因素在獲得更好的性能上起到了主導作用。

1. 本方案的吞吐提高主要來自于張量并行 (Tensor Parallelism，TP) 而不是 Acclerate 的管線并行 (Pipeline Parallelism，PP)。因為 Accelerate 旨在成為非常通用的方案，因此也非常不幸地很難最大化 GPU 使用率。它首先在 GPU 0 上完成所有計算，然后是 GPU 1，等等，一直到 GPU 8，這意味著任何時刻都有 7 個 GPU 是空閑的。而另一方面，DeepSpeed-Inference 使用了 TP，意味著它會向所有 GPU 發(fā)送張量，在每個 GPU 上計算部分生成結(jié)果，然后在所有的 GPU 間通信計算結(jié)果，并繼續(xù)做下一層。這就是說 TP 所有的 GPU 都同時是活躍的，但它需要比 PP 多得多的通信。

2. DeepSpeed-Inference 還使用了定制的 CUDA 核函數(shù)以避免分配太多內(nèi)存以及太多進出 GPU 的張量拷貝。這么做會減少顯存需求及核函數(shù)啟動次數(shù)從而提高吞吐，另外還可以支持更大的 batch size 從而進一步增加總吞吐。

如果你對更多的例子感興趣，可以看看?在 GPU 上使用 DeepSpeed-Inference 加速 GPT-J 推理[9]?或?在 GPU 上使用 DeepSpeed-Inference 加速 BERT 推理[10]。

Deepspeed ZeRO-Inference

Deepspeed ZeRO[11]?使用一個魔術(shù)般的分片方法，使得它可以輸入幾乎任何模型并將它擴展到少至幾個多至上百個 GPU，進行訓練或推理。

設(shè)置

運行

注意到現(xiàn)在為止的腳本都是所有 GPU 都處理相同的輸入，但你其實可以在每個 GPU 上運行不同的流，從而得到?n_gpu?倍的吞吐。你不能用 Deepspeed-Inference 達到這個目的。

請記住用戶可以用 ZeRO 同時創(chuàng)建多個不同的流，因此總性能應該是每秒每詞的吞吐除以參與計算的 GPU 的數(shù)目，因此根據(jù)你是使用 16 個 GPU 還是 8 個 GPU，可以獲得 8 倍或者 16 倍的更快性能。

你還可以在一個小型 GPU 上試試卸載方案，運行的時間會很長，但是如果你沒有 8 個巨型 GPU 的話這也是一個聊甚于無的方案。

CPU 卸載 (1x GPUs):

NVMe 卸載 (1x GPUs):

請確保在你的高速 NVMe 盤上預留約 400GB 的空間，并把?/path/to/nvme_offload?設(shè)成它。

更多客戶端及服務(wù)器方案

你可以從?transformers-bloom-inference[12]?找到更多非常高效的方案，包括服務(wù)器方案。

這里我們提供一些預覽。

服務(wù)器方案：

• Mayank Mishra[13]?拿著本博文中討論的所有演示代碼并把它們變成了一個網(wǎng)絡(luò)服務(wù)包，你可以從?這兒[14]?下載。

• Nicolas Patry[15]?開發(fā)了一個超高效的?基于 Rust 的網(wǎng)絡(luò)服務(wù)方案[16]。

更多的客戶端方案：

• Thomas Wang[17]?正在開發(fā)一個很快的?定制 CUDA 核函數(shù)的 BLOOM 模型[18]。

• HuggingFace 的 JAX 組已開發(fā)了一個?基于 JAX 的方案[19]

因為如果你在本博文發(fā)布幾個月后讀到它，很有可能它已經(jīng)不能反映最新的狀態(tài)了，你可以去?transformers-bloom-inference 的 GitHub 倉庫[20]?找到最新的方案。

致謝

萬分感謝如下這些人，他們提出了好的問題并幫助提高了本文的可讀性：Olatunji Ruwase 和 Philipp Schmid。

英文原文:https://huggingface.co/blog/bloom-inference-pytorch-scripts

譯者: Matrix Yao (姚偉峰)

文內(nèi)鏈接

[1]?BLOOM 模型:?https://huggingface.co/bigscience/bloom