Hugging Face BLOOM 推論最適化

Hugging Face は、BLOOM モデルに対して一連の反復的最適化を行うことで、推論レイテンシを 5 倍削減し、スループットを 50 倍増加させました。最終的なアーキテクチャは、Tensor Parallelism (TP) と組み合わせた PyTorch、注意機構用のカスタム CUDA カーネル、そしてカーネル融合のための torch.jit.script を利用しています。

Pipeline から Tensor Parallelism への移行

BLOOM（176B パラメータ、bf16 で 352GB）の初期推論は、accelerate ライブラリの device_map="auto" を使用した Pipeline Parallelism (PP) で実装されました。PP では、各 GPU が特定のレイヤー集合を所有し、データを順次処理して次の GPU に渡します。

レイテンシを削減するため、Hugging Face は Tensor Parallelism (TP) に移行しました。TP では、各 GPU がすべてのレイヤーの重みの一部を所有し、すべての GPU が同時に動作できます。この変更により、パフォーマンスが劇的に向上しました：

  • Latency: 300ms/token から 91ms/token に低下。
  • Throughput: 1 秒あたり 10 リクエスト (RPS) に増加。

ncclAllReduce による通信オーバーヘッドは TP でも発生しますが、バッチ処理（バッチサイズ 1 と 32 がしばしば同様のレイテンシを示す）できることにより、全体のスループットが大幅に向上しました。

PyTorch ベースの最適化

並列化戦略に加えて、TensorBoard でのプロファイリングにより特定されたボトルネックを除去するため、いくつかの低レベル PyTorch 最適化が実装されました。

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 はマスク付きフィルと softmax 操作を融合するカスタム CUDA カーネルを開発しました。

具体的には、カーネルは以下のシーケンスを最適化します：

  1. 注意マスクを使用した注意スコアの masked_fill_。
  2. 安定性のために float32 での softmax 計算。

カーネル内部で必要な合計と蓄積にのみアップキャストを限定することで、レイテンシは 81ms/token から 71ms/token にさらに削減されました。

Webサーバーアーキテクチャとリクエスト処理

パラメータや長さが多様なユーザーリクエストに対応するため、Hugging Face は柔軟なバッチングシステムを実装しました：

  • プロセス間通信: torch.distributed が別プロセスを必要とするため、サーバーは Redis の pub/sub を使用して生文字列をすべてのプロセスに配布します。
  • カスタム生成ループ: 標準の generate 関数は、バッチ内の各メンバーに異なるパラメータ（例：サンプリング、top-p）を適用するカスタム実装に置き換えられました。
  • 動的バッチ抽出: 同じバッチ内の長いリクエストが短いリクエストを遅延させるのを防ぐため、サーバーはトークン上限に達した時点で完了したリクエストを抽出して返却し、バッチ全体の完了を待ちません。

評価したが採用しなかったアプローチ

最適化プロセス全体で、他にもいくつかのアプローチが検討されました：

  • TPU 上の JAX/Flax: 並列化は実装が容易でしたが、Ray と TPU ワーカー間の通信で重大な安定性問題が発生し、コンパイルの細かい制御が欠如していました。
  • DeepSpeed: 最終イテレーションに匹敵する印象的な結果を示しましたが、ストレス下での定期的なカーネルクラッシュ（CUDA illegal access）など、安定性の問題がありました。
  • Rust 実装: tch-rs を使用して Rust でバージョンが作成され、並行制御が向上しました。しかし、実際のパフォーマンス向上は PyTorch ベンチマークでプロファイラが有効なままだったことが原因であることが判明しました。
  • ONNX/TensorRT: テキスト生成ループの柔軟性と、ロジット計算のためにテンソルを GPU 上に保持する必要性に対して、これらは硬直的すぎると判断されました。

Sources