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

(資料圖片僅供參考)

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

一般來講，使用單節點會帶來最快的吞吐，因為大多數時候節點內的 GPU 互聯硬件比節點間的快，但未必總是如此。

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

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

測試基準

事不宜遲，我們先展示一些數據吧。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

方案

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

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

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

HuggingFace Accelerate

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

用空的權重實例化模型。

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

逐比特加載模型 checkpoint 并把權重加載到相應的設備。

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

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

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

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

設置

運行

簡單執行如下命令。

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

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

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

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

DeepSpeed-Inference

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

設置

運行

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

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

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

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

當然，現在僅需 4x80GB A100 GPU 就夠了：

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

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

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

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

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

Deepspeed ZeRO-Inference

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

設置

運行

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

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

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

CPU 卸載 (1x GPUs):

NVMe 卸載 (1x GPUs):

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

更多客戶端及服務器方案

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

這里我們提供一些預覽。

服務器方案：

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

Nicolas Patry[15]?開發了一個超高效的?基于 Rust 的網絡服務方案[16]。

更多的客戶端方案：

Thomas Wang[17]?正在開發一個很快的?定制 CUDA 核函數的 BLOOM 模型[18]。

HuggingFace 的 JAX 組已開發了一個?基于 JAX 的方案[19]

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

致謝

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

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

譯者: Matrix Yao (姚偉峰)

文內鏈接

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

[2]?8 比特量化方案:?https://huggingface.co/blog/hf-bitsandbytes-integration[3]?BitsAndBytes:?https://github.com/TimDettmers/bitsandbytes[4]?Deepspeed-Inference:?https://www.deepspeed.ai/tutorials/inference-tutorial/[5]?法國 Jean Zay 超算中心:?http://www.idris.fr/eng/jean-zay/index.html[6]?HuggingFace Accelerate 文檔:?https://huggingface.co/docs/accelerate/big_modeling[7]?BitsAndBytes:?https://github.com/TimDettmers/bitsandbytes[8]?DeepSpeed-Inference:?https://www.deepspeed.ai/tutorials/inference-tutorial/[9]?在 GPU 上使用 DeepSpeed-Inference 加速 GPT-J 推理:?https://www.philschmid.de/gptj-deepspeed-inference[10]?在 GPU 上使用 DeepSpeed-Inference 加速 BERT 推理:?https://www.philschmid.de/bert-deepspeed-inference[11]?Deepspeed ZeRO:?https://www.deepspeed.ai/tutorials/zero/[12]?transformers-bloom-inference:?https://github.com/huggingface/transformers-bloom-inference[13]?GitHub 用戶: Mayank Mishra:?https://github.com/mayank31398[14]?GitHub transformers-bloom-inference 代碼倉庫:?https://github.com/huggingface/transformers-bloom-inference[15]?GitHub 用戶: Nicolas Patry:?https://github.com/Narsil[16]?基于 Rust 的網絡服務方案:?https://github.com/Narsil/bloomserver[17]?GitHub 用戶: Thomas Wang:?https://github.com/thomasw21[18]?定制 CUDA 核函數的 BLOOM 模型:?https://github.com/huggingface/transformers_bloom_parallel[19]?基于 JAX 的方案:?https://github.com/huggingface/bloom-jax-inference[20]?transformers-bloom-inference 的 GitHub 倉庫:?https://github.com/huggingface/transformers-bloom-inference

免責聲明：本文不構成任何商業建議，投資有風險，選擇需謹慎！本站發布的圖文一切為分享交流，傳播正能量，此文不保證數據的準確性，內容僅供參考

關鍵詞： 網絡服務 多長時間