使用 DeepSpeed 和 Accelerate 實現快得驚人的 BLOOM 推理

本文展示了在使用擁有 176B 引數的 BLOOM 模型進行文字生成時，如何獲得快得驚人的單 token 吞吐量。

由於該模型在 bf16 (bfloat16) 精度下需要 352GB 的權重 (176*2)，最高效的配置是 8x80GB 的 A100 GPU。此外，2x8x40GB 的 A100s 或 2x8x48GB 的 A6000 也可以使用。使用這些 GPU 的主要原因是，在撰寫本文時，它們提供了最大的 GPU 視訊記憶體，但其他 GPU 也可以使用。例如，24x32GB 的 V100s 也是可以的。

使用單個節點通常能提供最快的吞吐量，因為大多數情況下，節點內 GPU 之間的連線硬體比節點間連線更快，但這並非總是如此。

如果你沒有那麼多硬體，仍然可以在較小的 GPU 上執行 BLOOM 推理，方法是使用 CPU 或 NVMe 解除安裝，但當然，生成時間會慢得多。

我們還將介紹 8 位元量化解決方案，這些方案需要的 GPU 視訊記憶體減半，但代價是吞吐量略有降低。我們將討論 BitsAndBytes 和 Deepspeed-Inference 這兩個庫。

基準測試

話不多說，我們先來看一些資料。

為了保持一致性，除非另有說明，本文中的所有基準測試都是在 Jean Zay HPC 上的同一個 8x80GB A100 節點上完成的，該節點配有 512GB 的 CPU 記憶體。JeanZay HPC 使用者享有約 3GB/s 的快速 IO 讀取速度（GPFS）。這對於檢查點載入時間非常重要。慢速磁碟會導致載入時間變慢。尤其是在我們併發地在多個程序中進行 IO 操作時。

所有基準測試都在進行 貪心生成 100 個 token 的輸出。

Generate args {'max_length': 100, 'do_sample': False}

輸入提示僅包含幾個 token。同時也開啟了前一個 token 的快取，因為每次都重新計算會相當慢。

首先，我們快速看一下準備生成需要多長時間——即載入和準備模型需要多長時間

DeepSpeed-Inference 附帶了預分片的權重倉庫，載入時間大約為 1 分鐘。Accelerate 的載入時間也非常出色，僅需約 2 分鐘。其他解決方案在這裡則慢得多。

載入時間可能重要，也可能不重要，因為一旦載入完成，你就可以一次又一次地持續生成 token，而無需額外的載入開銷。

接下來是最重要的 token 生成吞吐量基準測試。這裡的吞吐量指標很簡單——生成 100 個新 token 所需的時間除以 100 和批處理大小（即除以生成的 token 總數）。

以下是在 8x80GB GPU 上的吞吐量（毫秒）

其中 OOM == Out of Memory（記憶體不足），表示批處理大小太大而無法裝入 GPU 記憶體。

使用 Deepspeed-Inference 的張量並行（Tensor Parallelism，TP）和定製的融合 CUDA 核心，實現了低於 1 毫秒的吞吐量！這絕對是驚人的！不過，將此解決方案用於其他未經測試的模型可能需要一些開發時間才能使其工作。

Accelerate 也超級快。它使用一種非常簡單的樸素流水線並行（Pipeline Parallelism，PP）方法，由於其簡單性，它應該可以與任何模型開箱即用。

由於 Deepspeed-ZeRO 可以並行處理多個生成流，其吞吐量可以進一步除以 8 或 16，具體取決於在 `generate` 呼叫期間使用了 8 個還是 16 個 GPU。當然，這意味著在 8x80 A100 的情況下，它可以處理大小為 64 的批次（見上表），因此吞吐量約為 4 毫秒——所以這三種解決方案非常接近。

讓我們再回顧一下這些數字是如何計算的。在使用 Deepspeed-Inference 的 fp16 模式時，為批大小為 128 的批次生成 100 個新 token 的實際時間為 8832 毫秒。因此，為了計算吞吐量，我們做了：walltime/(batch_size*new_tokens) 或 `8832/(128*100) = 0.69`。

現在讓我們看看 Deepspeed-Inference 和 BitsAndBytes 提供的基於 int8 量化模型的強大功能，因為它只需要 bfloat16 或 float16 推理所需 GPU 視訊記憶體的一半。

4x80GB A100 上的吞吐量（毫秒）

要重現基準測試結果，只需在下面討論的這 3 個指令碼中的任何一箇中新增 `--benchmark` 即可。

解決方案

首先檢出演示倉庫

git clone https://github.com/huggingface/transformers-bloom-inference
cd transformers-bloom-inference

在本文中，我們將使用位於 `bloom-inference-scripts/` 下的 3 個指令碼。

特定於框架的解決方案按字母順序呈現

HuggingFace Accelerate

Accelerate

Accelerate 以下列方式處理大模型的推理

1. 用空權重例項化模型。
2. 分析每一層的大小以及每個裝置（GPU、CPU）上的可用空間，以決定每一層應該放在哪裡。
3. 逐位載入模型檢查點，並將每個權重放在其裝置上

然後，它透過鉤子確保模型正常執行，這些鉤子將輸入和輸出傳輸到正確的裝置上，並將解除安裝到 CPU（甚至磁碟）上的模型權重在正向傳播之前載入到 GPU 上，然後在正向傳播結束後再次解除安裝。

在有多個 GPU 且空間足以容納整個模型的情況下，它會從一個 GPU 切換到下一個 GPU，直到所有層都執行完畢。任何給定時間只有一個 GPU 在工作，這聽起來效率很低，但儘管 GPU 處於空閒狀態，它仍能產生不錯的吞吐量。

它也非常靈活，因為相同的程式碼可以在任何給定的設定上執行。Accelerate 會首先使用所有可用的 GPU，然後在 RAM 滿時解除安裝到 CPU，最後解除安裝到磁碟。解除安裝到 CPU 或磁碟會使速度變慢。例如，有使用者報告稱，在僅有 2 個 A100 的情況下，無需更改程式碼即可執行 BLOOM，吞吐量為每 token 15 秒，而在 8x80 A100 上則為 10 毫秒。

你可以在 Accelerate 文件中瞭解更多關於此解決方案的資訊。

設定

pip install transformers>=4.21.3 accelerate>=0.12.0

執行

簡單的執行方式是

python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --batch_size 1 --benchmark

要啟用來自 BitsAndBytes 的 8 位元量化解決方案，首先安裝 `bitsandbytes`

pip install bitsandbytes

然後在之前的命令列中新增 `--dtype int8`

python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

如果你有超過 4 個 GPU，你可以告訴它只使用 4 個，使用

CUDA_VISIBLE_DEVICES=0,1,2,3 python bloom-inference-scripts/bloom-accelerate-inference.py --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

在這種情況下，我們能夠執行的最高批處理大小是 40，沒有出現 OOM。如果你檢視指令碼內部，我們不得不調整記憶體分配對映，以釋放第一個 GPU 來專門處理啟用和前一個 token 的快取。

DeepSpeed-Inference

DeepSpeed-Inference 使用張量並行和高效的融合 CUDA 核心，在大批處理大小為 128 的情況下，提供了超快的每 token 不到 1 毫秒的推理速度。

設定

pip install deepspeed>=0.7.3

執行

1. 最快的方法是使用一個經過 TP-pre-sharded（TP = 張量並行）的檢查點，載入只需約 1 分鐘，而未經預分片的 bloom 檢查點則需要 10 分鐘

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-fp16

1a. 如果你想執行原始的 bloom 檢查點，載入後其吞吐量與前述方案相同，但載入過程將需要 10-20 分鐘

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name bigscience/bloom

2a. 8 位元量化版本只需要正常半精度版本一半的 GPU 視訊記憶體

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

這裡我們使用了 `microsoft/bloom-deepspeed-inference-int8`，並告訴指令碼以 `int8` 模式執行。

當然，現在只需要 4x80GB 的 A100 GPU 就足夠了

deepspeed --num_gpus 4 bloom-inference-scripts/bloom-ds-inference.py --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

在這種情況下，我們能夠執行的最高批處理大小是 128，沒有出現 OOM。

你可以看到這裡有兩個因素共同作用，帶來了更好的效能。

1. 這裡的吞吐量透過使用張量並行（TP）而不是 Accelerate 的流水線並行（PP）得到了提升。因為 Accelerate 的目標是具有很高的通用性，不幸的是，這也使得最大化 GPU 使用率變得困難。所有計算首先在 GPU 0 上完成，然後在 GPU 1 上，依此類推，直到 GPU 8，這意味著 7 個 GPU 一直處於空閒狀態。另一方面，DeepSpeed-Inference 使用 TP，這意味著它會將張量傳送到所有 GPU，在每個 GPU 上計算生成的一部分，然後所有 GPU 相互通訊結果，再進行下一層。這意味著所有 GPU 同時處於活動狀態，但它們需要更多的通訊。

2. DeepSpeed-Inference 還使用定製的 CUDA 核心來避免分配過多記憶體以及在 GPU 之間進行張量複製。這樣做的效果是減少了記憶體需求和核心啟動次數，從而提高了吞吐量，並允許更大的批處理大小，最終帶來更高的整體吞吐量。

如果你對更多示例感興趣，可以檢視 在 GPU 上使用 DeepSpeed-Inference 加速 GPT-J 推理 或 在 GPU 上使用 DeepSpeed-Inference 加速 BERT 推理。

Deepspeed ZeRO-Inference

Deepspeed ZeRO 使用一種神奇的分片方法，幾乎可以處理任何模型，並將其擴充套件到幾個或數百個 GPU 上進行訓練或推理。

設定

pip install deepspeed

執行

請注意，該腳本當前在所有 GPU 上執行相同的輸入，但你可以在每個 GPU 上執行不同的流，從而獲得 `n_gpu` 倍的吞吐量。而 Deepspeed-Inference 無法做到這一點。

deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 1 --benchmark

請記住，使用 ZeRO，使用者可以同時生成多個獨立的流，因此整體效能應該是以秒/token 計算的吞吐量除以參與的 GPU 數量——因此，根據使用的是 8 個還是 16 個 GPU，速度會快 8 到 16 倍！

你也可以只用一個較小的 GPU 嘗試解除安裝解決方案，這會花費很長時間，但如果你沒有 8 個大型 GPU，這已經是最好的選擇了。

CPU-Offload (1x GPUs)

deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 8 --cpu_offload --benchmark

NVMe-Offload (1x GPUs)

deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py --name bigscience/bloom --batch_size 8 --nvme_offload_path=/path/to/nvme_offload --benchmark

請確保將 `/path/to/nvme_offload` 調整到一個你有約 400GB 可用空間的快速 NVMe 驅動器上。

額外的客戶端和伺服器解決方案

在 transformers-bloom-inference，你會發現更多非常高效的解決方案，包括伺服器解決方案。

以下是一些預覽。

伺服器解決方案

• Mayank Mishra 將本部落格文章中討論的所有演示指令碼轉化為一個網路伺服器包，你可以從這裡下載

• Nicolas Patry 開發了一個超高效的基於 Rust 的網路伺服器解決方案。

更多客戶端解決方案

• Thomas Wang 正在開發一個非常快的定製 CUDA 核心 BLOOM 模型。

• HuggingFace 的 JAX 團隊開發了一個基於 JAX 的解決方案

由於這篇部落格文章在你閱讀時可能已經過時，如果是在釋出幾個月後閱讀，請使用 transformers-bloom-inference 尋找最新的解決方案。

部落格致謝

非常感謝以下友好人士，他們提出了很好的問題，並幫助提高了文章的可讀性：Olatunji Ruwase 和 Philipp Schmid。