DeepSpeed と Accelerate を用いた驚異的に高速な BLOOM 推論

TL;DR

Hugging Face は、176 billion パラメータの BLOOM モデルが DeepSpeed‑Inference のテンソル並列（または Accelerate のパイプライン並列）を使用して、8 × 80 GB A100 GPU の単一ノード上で 1 ms 未満でトークンを生成できることを示し、メモリ使用量を半減する 8‑bit 量子化オプションも提供しています。

ハードウェア要件とセットアップ

  • 最適構成: 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 を実行できますが、生成レイテンシは大幅に増加します。
  • 量子化推論: BitsAndBytes による 8‑bit モデルは、GPU メモリを約半分に抑えつつ、わずかなスループット低下で動作します。

ベンチマーク環境

  • Node: Jean Zay HPC、8 × 80 GB A100、512 GB CPU RAM、GPFS ストレージ（約 3 GB/s 読み取り速度）。
  • Task: 短いプロンプトと KV‑cache 有効で、100 トークンの貪欲生成（max_length=100, do_sample=False）。
  • Metrics: モデルロード時間（秒）とトークンあたりスループット（ms per token = 壁時計時間 ÷ (バッチ × トークン数)）。

モデルロード時間

ソリューション ロード時間 (秒)
Accelerate 121
DeepSpeed‑Inference (shard‑int8) 61
DeepSpeed‑Inference (shard‑fp16) 60
DeepSpeed‑Inference (unsharded) 662
DeepSpeed‑ZeRO 462

事前にシャードされた DeepSpeed チェックポイントは約 1 分でロードされますが、シャードされていないチェックポイントは 10 分以上かかることがあります。

トークン生成スループット (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 はテンソル並列 (TP) とカスタム融合 CUDA カーネルにより、バッチ 128 でトークンあたり 1 ms 未満を実現します。
  • Accelerate は単純なパイプライン並列 (PP) を使用し、バッチ 32 でトークンあたり約 10 ms に達しますが、バッチ 64 以上では GPU メモリ制限を超えることができません。
  • 量子化 int8 の実行はメモリ消費を半減させます。DeepSpeed‑Inference int8 はバッチ 128 でトークンあたり約 1 ms を維持しますが、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

  • Approach: モデルの重みをレイヤーサイズと利用可能メモリに基づいてデバイスに遅延ロードし、各レイヤーで 1 つの GPU のみがアクティブになるシンプルなパイプライン並列を使用します。
  • Pros: 任意のハードウェア構成で即座に動作し、GPU メモリが不足している場合は CPU やディスクにオフロードできます。
  • Cons: フォワードパスの大部分で GPU がアイドル状態になるため、ピークスループットが制限されます。バッチサイズが大きくなるとすぐに OOM になります。
  • Usage:
    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

  • Approach: テンソル並列で各レイヤーを GPU 間で分割し、カスタム融合カーネルでメモリコピーとカーネル起動回数を削減します。
  • Performance drivers:
    1. TP vs PP – すべての GPU が同時に計算を行い、利用率が向上します。
    2. Fused kernels – メモリオーバーヘッドが低く、カーネル起動回数が減ります。
  • Pre‑sharded checkpoints (microsoft/bloom-deepspeed-inference-fp16) は約 1 分でロードされ、非シャードのチェックポイントは 10‑20 分かかりますが、ロード後の速度は同じです。
  • Quantized int8 (microsoft/bloom-deepspeed-inference-int8) はメモリ要件を半減し、8 × 80 GB A100 でバッチ 128 時にトークンあたり 1 ms 未満を達成、4 × 80 GB A100 でも動作します。
  • Usage:
    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

  • Approach: モデル状態を GPU 間でシャード（ZeRO‑3 スタイル）し、複数の独立した生成ストリームを並列に実行でき、GPU 数に比例した速度向上が得られます。
  • Limitations: 提供されたスクリプトは同じ入力をすべての GPU で実行します。理論的な 8× または 16× の速度向上を得るには、GPU ごとにカスタムストリームが必要です。
  • Off‑loading options: CPU‑offload または NVMe‑offload により、単一 GPU での推論が可能ですが、レイテンシは大きくなります。
  • Usage:
    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 はデモスクリプトをすぐに実行できるウェブサーバーにパッケージ化しました。
    • Nicolas Patry は高性能な Rust ベースのサーバーを作成しました。
  • Client‑side projects:
    • Thomas Wang はカスタム CUDA カーネルで加速された BLOOM モデルを開発しています。
    • HuggingFace JAX チームは BLOOM 用の JAX 推論バックエンドをリリースしました。
  • Keeping up‑to‑date:
    • transformers-bloom-inference リポジトリは最新のスクリプトとサーバー実装を集約しています。

実践的なポイント

  • 単一ノードで最大のスループット を得るには、事前にシャードされた fp16 チェックポイントを使用した DeepSpeed‑Inference を利用してください。バッチ 128 でトークンあたり 1 ms 未満が期待できます。
  • GPU メモリが制限されている場合 は、BitsAndBytes（DeepSpeed または Accelerate）による 8 ビット量子化に切り替えて、メモリ要件を半減し、速度低下は控えめです。
  • Accelerate は異種ハードウェアに最も柔軟なソリューション ですが、パイプライン並列により DeepSpeed のテンソル並列よりもレイテンシが高くなります。
  • ZeRO 推論は多数の GPU にスケールするか、オフロードを使用して単一 GPU で実行する道を提供 し、大規模 GPU クラスタが利用できない場合に有用です。

「このケースで OOM にならずに実行できた最大バッチサイズは 128 でした。」 – Hugging Face ブログ、DeepSpeed‑Inference int8（8 × 80 GB A100）について。

この投稿は 2022 年 9 月時点の BLOOM 推論性能を反映しています。最新の最適化については transformers-bloom-inference GitHub リポジトリをご参照ください。

Sources