Hugging Face BLOOM 推論最佳化

Hugging Face 透過一系列迭代最佳化，使 BLOOM 模型的推論延遲降低 5 倍，吞吐量提升 50 倍。最終架構結合了 PyTorch 與張量平行化 (TP)、用於注意力的自訂 CUDA 核心，以及 torch.jit.script 進行核心融合。

從流水線平行化轉向張量平行化

BLOOM（176 億參數，352GB bf16）最初的推論是透過 accelerate 套件的 device_map="auto" 使用流水線平行化 (PP) 實作的。在 PP 中，每張 GPU 擁有特定的層，依序處理資料並將結果傳遞給下一張 GPU。

為了降低延遲，Hugging Face 轉而使用張量平行化 (TP)，在此模式下每張 GPU 持有每層權重的一部分，使所有 GPU 能同時工作。此轉變帶來了顯著的效能提升：

  • 延遲: 從 300ms/token 降至 91ms/token。
  • 吞吐量: 提升至每秒 10 個請求 (RPS)。

雖然 TP 透過 ncclAllReduce 帶來通訊開銷，但批次處理的能力（批次大小 1 與 32 常有相似的延遲）大幅提升了整體吞吐量。

基於 PyTorch 的最佳化

除了平行化策略外，還實作了多項低階 PyTorch 最佳化，以消除透過 TensorBoard 觀測到的瓶頸。

使用 torch.jit.script 進行核心融合

Gelu 運算子原本會啟動多個逐元素核心，導致過多的張量拷貝。透過在 bloom_gelu_forward 函式上套用 @torch.jit.script，Hugging Face 將其融合為單一核心操作，將延遲從 91ms/token 降至 81ms/token。

高效的 PyTorch 實作

  • ALiBi 最佳化: 位置嵌入先前在過多位置被計算。將此計算集中化後，使該操作加速了 10 倍。
  • 減少張量拷貝: 觀測顯示注意力路徑受到 reshape 與 transpose 操作的沉重負擔。重新設計權重與 KV 快取（即 "past"）後，移除了這些不必要的拷貝。

自訂 CUDA 核心與硬體加速

為了進一步最佳化 torch.jit.script 無法覆蓋的熱點路徑，Hugging Face 開發了自訂 CUDA 核心，將 masked fill 與 softmax 操作融合。

具體而言，該核心優化了以下序列：

  1. 使用注意力遮罩對注意力分數執行 masked_fill_。
  2. 在 float32 上計算 softmax 以確保穩定性。

透過僅在核心內對必要的求和與累加進行升階，將延遲進一步從 81ms/token 降至 71ms/token。

網路伺服器架構與請求處理

為了服務多樣化的使用者請求，包含不同參數與長度，Hugging Face 實作了彈性的批次系統：

  • 跨程序通訊: 由於 torch.distributed 需要獨立的程序，伺服器使用 Redis pub/sub 將原始字串分發給所有程序。
  • 自訂生成迴圈: 標準的 generate 函式被自訂實作取代，該實作會對批次中的每個成員套用不同參數（例如抽樣、top-p）。
  • 動態批次抽取: 為避免短請求被同批次的長請求拖延，伺服器會在請求達到其 token 限制時即時抽取並回傳完成的請求，而不是等整個批次結束。

評估過但被棄用的方法

在最佳化過程中，探索了其他多條路徑：

  • JAX/Flax 在 TPU 上: 雖然平行化較易實作，但團隊遭遇 Ray 與 TPU 工作者通訊的重大穩定性問題，且缺乏細緻的編譯控制。
  • DeepSpeed: 雖然提供了與最終版本相似的驚人結果，但在壓力測試下會出現穩定性問題，包括頻繁的核心崩潰（CUDA illegal access）。
  • Rust 實作: 使用 tch-rs 用 Rust 撰寫了一個版本，以提升併發控制。然而發現所謂的效能提升其實是因為在 PyTorch 基準測試中遺留了啟用的分析器。
  • ONNX/TensorRT: 這些方案被認為過於僵硬，無法滿足文字生成迴圈所需的彈性，以及在 GPU 上保留張量以計算 logits 的需求。

SUMMARY: Hugging Face 透過從流水線平行化轉向張量平行化並實作自訂 CUDA 核心，使 BLOOM 模型的延遲降低 5 倍、吞吐量提升 50 倍。

TITLE: Hugging Face BLOOM 推論最佳化

Sources