BLOOMZ 在 Habana Gaudi2 加速器上的推論
TL;DR
Hugging Face 已示範，Habana Gaudi2 加速器相較於目前可用的 GPU，能為像 BLOOMZ 這樣的大型語言模型提供更快的推論。對於 1760 億參數的 BLOOMZ 模型，Gaudi2 的速度比 Nvidia A100 80GB 快 1.42 倍。
BLOOMZ 模型能力
BLOOMZ 是 BLOOM 的微調版本，BLOOM 是一個 1760 億參數的自回歸模型，旨在完成 46 種語言和 13 種程式語言的文字序列。BLOOMZ 透過提升泛化能力與零樣本能力（即在未見過的輸入資料上完成任務，無需先前的訓練範例）來改進原始的 BLOOM 架構。
部署 BLOOMZ 計算需求極高；在 16 位元精度下，單一實例需要 352 GB 記憶體，必須使用高效能硬體與最佳化函式庫才能達到低延遲。
Habana Gaudi2 硬體與軟體堆疊
Gaudi2 是 Habana Labs 的第二代 AI 硬體加速器。其架構設計用於平行執行一般矩陣乘法 (GeMM) 及其他運算，從而優化深度學習工作流程，無論是訓練或推論。
硬體規格
單一 Gaudi2 伺服器內含八顆 Habana Processing Units (HPU)，每顆配備 96 GB 記憶體，提供足夠容量以容納極大型模型。
軟體整合
- SynapseAI™ SDK：支援 PyTorch 與 DeepSpeed。它包含一個圖形編譯器，透過算子融合、資料布局管理、平行化、流水線化與記憶體管理來最佳化執行。
- HPU Graphs：最近在 SynapseAI 中推出，以支援對延遲敏感的應用程式。
- Optimum Habana：Hugging Face 的函式庫，充當 Gaudi2 硬體與 Transformers 函式庫之間的橋樑，簡化部署流程。
推論基準：Gaudi2 與 A100
為了滿足 BLOOMZ 的記憶體需求，Hugging Face 使用 DeepSpeed‑inference（透過 Habana 的 DeepSpeed 分支）在八個裝置上實作模型與管線平行化。
延遲結果
基準測試使用 16 位元精度、貪婪生成 100 個 token，並使用鍵值快取。結果顯示 Gaudi2 在延遲上始終優於 Nvidia A100 80GB：
| 模型 | 裝置數量 | Gaudi2 延遲（秒） | A100‑80GB 延遲（秒） | 第一代 Gaudi 延遲（秒） |
|---|---|---|---|---|
| BLOOMZ (176B) | 8 | 3.103 | 4.402 | / |
| BLOOMZ-7B | 8 | 0.734 | 2.417 | 3.321 |
| BLOOMZ-7B | 1 | 0.772 | 2.119 | 2.387 |
主要發現：
- 176B 模型：Gaudi2 的速度比 A100 80GB 快 1.42 倍。
- 7B 模型：Gaudi2 的速度比 A100 80GB 快 2.89 倍。
- 模型平行化：Gaudi2 從模型平行化中獲益顯著，而對於較小的 7B 模型，A100 在單一裝置上更快。
第一代 Gaudi 的性價比
對於 BLOOMZ‑7B 模型，第一代 Gaudi 相較於 A100 提供更佳的性價比。A100 每小時成本超過 $30，而第一代 Gaudi（AWS 上的 DL1 實例）每小時約 $13，且仍保持具競爭力的延遲（2.387 秒）。
實作與再現
可使用 optimum-habana 倉庫中提供的腳本，在完整資料集上執行推論。基準環境使用 Transformers v4.28.1、SynapseAI v1.9.0 與 Optimum Habana v1.5.0。
若要再現結果，使用者必須安裝最新的 SynapseAI 與 Gaudi 驅動程式，接著安裝 optimum-habana 函式庫與 Habana 專屬的 DeepSpeed 分支 (v1.9.0)。執行指令使用 gaudi_spawn.py，並加入 DeepSpeed、HPU 圖形與 KV 快取等旗標，以最佳化生成。