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

(資料圖片僅供參考)

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

一般來(lái)講，使用單節(jié)點(diǎn)會(huì)帶來(lái)最快的吞吐，因?yàn)榇蠖鄶?shù)時(shí)候節(jié)點(diǎn)內(nèi)的 GPU 互聯(lián)硬件比節(jié)點(diǎn)間的快，但未必總是如此。

如果你沒(méi)有這么高端的硬件或沒(méi)有這么多硬件，你仍可能通過(guò) CPU 卸載 (CPU offload) 或是 NVMe 卸載 (NVMe offload) 的方式在更小的 GPU 上對(duì) BLOOM 進(jìn)行推理。當(dāng)然，生成時(shí)間會(huì)慢很多。

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

測(cè)試基準(zhǔn)

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

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

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

輸入提示詞僅包含幾個(gè)詞。我們會(huì)緩存先前見(jiàn)到的詞，因?yàn)槊看沃匦掠?jì)算它們相當(dāng)慢。

首先，讓我們快速看一下從開始到準(zhǔn)備好花了多長(zhǎng)時(shí)間， 即模型加載和準(zhǔn)備花了多長(zhǎng)時(shí)間：

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

加載時(shí)間有可能重要也可能并不重要，因?yàn)橐坏┘虞d成功你可以一遍遍持續(xù)不斷地生成詞而不再需要額外地加載開銷。

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

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

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

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

Accelerate 也超級(jí)快。它使用了非常簡(jiǎn)單的管線并行 (Pipeline Parallelism，PP)。因?yàn)樗浅：?jiǎn)單，所以它應(yīng)該對(duì)任何模型都是開箱即用的。

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

讓我們?cè)僦匦驴匆幌逻@些數(shù)字是怎么計(jì)算出來(lái)的。舉個(gè)例子，使用 Deepspeed-Inference fp16 模式實(shí)時(shí)生成 batch size 為 128、長(zhǎng)度為 100 個(gè)新詞的文本花了 8832 毫秒，因此我們可以這樣計(jì)算吞吐：鐘面時(shí)間 / ( batch size * 新詞數(shù) ) 或?8821/(128*100) = 0.69。

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

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

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

方案

首先獲取最新的演示代碼倉(cāng)庫(kù)：

本文我們準(zhǔn)備使用?bloom-inference-scripts/文件夾下的 3 個(gè)腳本。

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

HuggingFace Accelerate

Accelerate 按如下步驟進(jìn)行大模型推理：

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

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

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

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

在有多個(gè) GPU 且有足夠空間放下整個(gè)模型的情形下，該方案在 GPU 間逐個(gè)切換直至所有層運(yùn)行完畢。每個(gè)給定的時(shí)間只有一個(gè) GPU 工作，這聽起來(lái)很沒(méi)效率。但盡管該方案 GPU 存在空閑，它的吞吐卻相當(dāng)不錯(cuò)。

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

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

設(shè)置

運(yùn)行

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

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

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

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

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

DeepSpeed-Inference

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

設(shè)置

運(yùn)行

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

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

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

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

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

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

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

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

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

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

Deepspeed ZeRO-Inference

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

設(shè)置

運(yùn)行

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

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

你還可以在一個(gè)小型 GPU 上試試卸載方案，運(yùn)行的時(shí)間會(huì)很長(zhǎng)，但是如果你沒(méi)有 8 個(gè)巨型 GPU 的話這也是一個(gè)聊甚于無(wú)的方案。

CPU 卸載 (1x GPUs):

NVMe 卸載 (1x GPUs):

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

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

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

這里我們提供一些預(yù)覽。

服務(wù)器方案：

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

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

更多的客戶端方案：

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

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

因?yàn)槿绻阍诒静┪陌l(fā)布幾個(gè)月后讀到它，很有可能它已經(jīng)不能反映最新的狀態(tài)了，你可以去?transformers-bloom-inference 的 GitHub 倉(cāng)庫(kù)[20]?找到最新的方案。

致謝

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

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

譯者: Matrix Yao (姚偉峰)

文內(nèi)鏈接

[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]?法國(guó) 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 代碼倉(cāng)庫(kù):?https://github.com/huggingface/transformers-bloom-inference[15]?GitHub 用戶: Nicolas Patry:?https://github.com/Narsil[16]?基于 Rust 的網(wǎng)絡(luò)服務(wù)方案:?https://github.com/Narsil/bloomserver[17]?GitHub 用戶: Thomas Wang:?https://github.com/thomasw21[18]?定制 CUDA 核函數(shù)的 BLOOM 模型:?https://github.com/huggingface/transformers_bloom_parallel[19]?基于 JAX 的方案:?https://github.com/huggingface/bloom-jax-inference[20]?transformers-bloom-inference 的 GitHub 倉(cāng)庫(kù):?https://github.com/huggingface/transformers-bloom-inference

免責(zé)聲明：本文不構(gòu)成任何商業(yè)建議，投資有風(fēng)險(xiǎn)，選擇需謹(jǐn)慎！本站發(fā)布的圖文一切為分享交流，傳播正能量，此文不保證數(shù)據(jù)的準(zhǔn)確性，內(nèi)容僅供參考

關(guān)鍵詞： 網(wǎng)絡(luò)服務(wù) 多長(zhǎng)時(shí)間