極速 BLOOM 推理：DeepSpeed 與 Accelerate

TL;DR

Hugging Face 顯示，176 十億參數的 BLOOM 模型在單節點八張 80 GB A100 GPU 上，使用 DeepSpeed‑Inference 張量並行（或使用 Accelerate 流水線並行達到相當速度）時，能在 1 毫秒以下生成 token，且還提供了將記憶體使用量減半的 8 位元量化替代方案。

硬體需求與設定

最佳配置：8 × 80 GB A100 GPU（352 GB bf16 權重）。其他選項包括 2 × 8 × 40 GB A100、2 × 8 × 48 GB A6000，或 24 × 32 GB V100。
單節點優勢：節點內 GPU 互連速度快於節點間連結，通常能提供更高的吞吐量。
低階選項：CPU 或 NVMe 卸載可在較小的 GPU 上執行 BLOOM，但生成延遲會顯著增加。
量化推理：8 位元模型（透過 BitsAndBytes）大約只需一半的 GPU 記憶體，且吞吐量僅有適度的損失。

基準測試環境

節點：Jean Zay HPC，8 × 80 GB A100，512 GB CPU 記憶體，GPFS 儲存（約 3 GB/s 讀取速度）。
任務：使用短提示且啟用 KV‑cache，貪婪生成 100 個 token（max_length=100, do_sample=False）。
指標：模型載入時間（秒）與每 token 吞吐量（ms per token = wall‑time ÷ (batch × tokens)）。

模型載入時間

方案	載入時間（秒）
Accelerate	121
DeepSpeed‑Inference (shard‑int8)	61
DeepSpeed‑Inference (shard‑fp16)	60
DeepSpeed‑Inference (unsharded)	662
DeepSpeed‑ZeRO	462

預先分片的 DeepSpeed 檢查點大約在 1 分鐘內載入，而未分片的檢查點可能需要超過 10 分鐘。

Token 生成吞吐量（8 × 80 GB A100）

方案（dtype）	批次 1	批次 8	批次 16	批次 32	批次 64	批次 128	批次 256	批次 512
Accelerate bf16	230.38 ms	31.78 ms	17.84 ms	10.89 ms	OOM	–	–	–
Accelerate int8	286.56 ms	40.92 ms	22.65 ms	13.27 ms	OOM	–	–	–
DeepSpeed‑Inference fp16	44.02 ms	5.70 ms	3.01 ms	1.68 ms	1.00 ms	0.69 ms	OOM	–
DeepSpeed‑Inference int8	89.09 ms	11.44 ms	5.88 ms	3.09 ms	1.71 ms	1.02 ms	0.71 ms	OOM
DeepSpeed‑ZeRO bf16	283 ms	34.88 ms	OOM	–	–	–	–	–

關鍵觀察

DeepSpeed‑Inference 在批次 128 時達到每 token 低於 1 ms 的速度，得益於張量並行（TP）與自訂融合 CUDA 核心。
Accelerate 採用簡易的流水線並行（PP），在批次 32 時可達約 10 ms 每 token，但在批次 64 以上因 GPU 記憶體限制無法繼續。
量化的 int8 執行可將記憶體使用量減半；DeepSpeed‑Inference int8 在批次 128 時仍可達約 1 ms 每 token，而 Accelerate int8 更早因記憶體不足而失敗。

量化 8 位元吞吐量（4 × 80 GB A100）

方案	批次 1	批次 8	批次 16	批次 32	批次 64	批次 128
Accelerate int8	284.15 ms	40.14 ms	21.97 ms	OOM	–	–
DeepSpeed‑Inference int8	156.51 ms	20.11 ms	10.38 ms	5.50 ms	2.96 ms	OOM

解決方案細節

HuggingFace Accelerate

方法：根據層大小與可用記憶體，延遲載入模型權重至裝置；使用簡單的流水線並行，使每層僅有一個 GPU 處於活躍狀態。
優點：可直接在任何硬體配置上運作；當 GPU 記憶體不足時，可卸載至 CPU 或磁碟。
缺點：在大部分前向傳播過程中 GPU 處於閒置，限制了峰值吞吐量；較大的批次大小會迅速觸發 OOM。

使用方式：

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
# 8‑bit 量化執行
pip install bitsandbytes
python bloom-inference-scripts/bloom-accelerate-inference.py \
  --name bigscience/bloom --dtype int8 --batch_size 1 --benchmark

DeepSpeed‑Inference

方法：張量並行將每層分割至多個 GPU；自訂融合核心減少記憶體拷貝與核心呼叫。
效能驅動因素：
1. TP vs PP – 所有 GPU 同時運算，提升利用率。
2. 融合核心 – 降低記憶體開銷並減少核心呼叫次數。
預先分片的檢查點（microsoft/bloom-deepspeed-inference-fp16）大約在 1 分鐘內載入；未分片的檢查點載入時間為 10‑20 分鐘，但載入後的執行速度相同。
量化 int8（microsoft/bloom-deepspeed-inference-int8）將記憶體需求減半，且在 8 × 80 GB A100 上的批次 128 時仍能達到低於 1 ms 每 token，或在 4 × 80 GB A100 上運作。

使用方式：

pip install deepspeed>=0.7.3
# 快速 TP‑預分片 fp16
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name microsoft/bloom-deepspeed-inference-fp16
# 原始檢查點（載入較慢）
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name bigscience/bloom
# 8‑bit 版本（記憶體減半）
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-inference.py \
  --name microsoft/bloom-deepspeed-inference-int8 --dtype int8

DeepSpeed‑ZeRO Inference

方法：將模型狀態分片至多個 GPU（ZeRO‑3 風格），並可平行執行多條獨立的生成串流，產生與 GPU 數量成比例的效能提升。
限制：提供的腳本在所有 GPU 上執行相同的輸入；要達到理論上的 8 倍或 16 倍加速，需要自訂每個 GPU 的串流。
卸載選項：CPU‑offload 或 NVMe‑offload 可在單一 GPU 上執行推理，但會帶來較高的延遲。

使用方式：

pip install deepspeed
# 多 GPU ZeRO 推理
deepspeed --num_gpus 8 bloom-inference-scripts/bloom-ds-zero-inference.py \
  --name bigscience/bloom --batch_size 1 --benchmark
# CPU 卸載（單 GPU）
deepspeed --num_gpus 1 bloom-inference-scripts/bloom-ds-zero-inference.py \
  --name bigscience/bloom --batch_size 8 --cpu_offload --benchmark
# NVMe 卸載（單 GPU）
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

社群伺服器與客戶端擴充

Server implementations:
- Mayank Mishra 將示範腳本打包成可直接執行的 Web 伺服器。
- Nicolas Patry 建立了一個高效能的 Rust 伺服器。
Client‑side projects:
- Thomas Wang 正在開發自訂 CUDA 核心加速的 BLOOM 模型。
- HuggingFace JAX 團隊發布了針對 BLOOM 的 JAX 推理後端。
保持最新：transformers-bloom-inference 倉庫彙總了最新的腳本與伺服器實作。

實務要點

欲在單節點獲得最高原始吞吐量，請使用 DeepSpeed‑Inference 搭配預先分片的 fp16 檢查點；在批次 128 時預期每 token 低於 1 ms。
當 GPU 記憶體受限 時，改用 BitsAndBytes 的 8 位元量化（DeepSpeed 或 Accelerate），可將記憶體需求減半，僅有輕微的速度損失。
Accelerate 仍是最彈性的 異構硬體解決方案，但其流水線並行的延遲高於 DeepSpeed 的張量並行。
ZeRO 推理 提供了跨多 GPU 擴展或在單 GPU 上透過卸載執行的途徑，適用於缺乏大型 GPU 叢集的情況。

「在此情況下，我們能夠在不發生 OOM 的前提下執行的最高批次大小為 128。」 – Hugging Face 部落格，描述在 8 × 80 GB A100 上的 DeepSpeed‑Inference int8。

本文反映了截至 2022 年 9 月的 BLOOM 推理效能狀況。欲了解更新的最佳化，請參考 transformers-bloom-inference GitHub 倉庫。

極速 BLOOM 推理：DeepSpeed 與 Accelerate

TL;DR

硬體需求與設定

基準測試環境

模型載入時間

Token 生成吞吐量（8 × 80 GB A100）

量化 8 位元吞吐量（4 × 80 GB A100）

解決方案細節

HuggingFace Accelerate

DeepSpeed‑Inference

DeepSpeed‑ZeRO Inference

社群伺服器與客戶端擴充

實務要點

Sources

Token 生成吞吐量（8 × 80 GB A100）

量化 8 位元吞吐量（4 × 80 GB A100）