Habana Gaudi2アクセラレータ上でのBLOOMZ推論

TL;DR

Hugging Faceは、Habana Gaudi2アクセラレータが現在利用可能なGPUと比較して、BLOOMZのような大規模言語モデルに対して高速な推論を提供することを実証しました。1760億パラメータのBLOOMZモデルにおいて、Gaudi2はNvidia A100 80GBよりも1.42倍速いです。

BLOOMZモデルの機能

BLOOMZは、BLOOMのファインチューニング版で、46言語と13プログラミング言語にわたるテキストシーケンスを完了するよう設計された1760億パラメータの自己回帰モデルです。BLOOMZは、元のBLOOMアーキテクチャを改良し、汎化性能とゼロショット能力—事前の学習例なしで未見の入力データに対してタスクを完了できる能力—を強化しています。

BLOOMZのデプロイは計算負荷が高く、16ビット精度では単一インスタンスに352 GBのメモリが必要です。そのため、低レイテンシを実現するには高性能ハードウェアと最適化ライブラリが不可欠です。

Habana Gaudi2ハードウェアとソフトウェアスタック

Gaudi2はHabana Labsの第2世代AIハードウェアアクセラレータです。そのアーキテクチャはGeneral Matrix Multiplication（GeMM）やその他の演算を並列に実行するよう設計されており、トレーニングと推論の両方のディープラーニングワークフローを最適化します。

ハードウェア仕様

1台のGaudi2サーバーには8つのHabana Processing Units（HPU）が搭載されており、各HPUは96 GBのメモリを備えているため、非常に大規模なモデルをホストするのに十分な容量を提供します。

ソフトウェア統合

  • SynapseAI™ SDK: PyTorch と DeepSpeed をサポートします。オペレータ融合、データレイアウト管理、並列化、パイプライン化、メモリ管理を通じて実行を最適化するグラフコンパイラが含まれています。
  • HPU Graphs: 最近SynapseAIに導入され、レイテンシに敏感なアプリケーションをサポートします。
  • Optimum Habana: Gaudi2ハードウェアとTransformersライブラリをつなぐHugging Faceのライブラリで、デプロイプロセスを簡素化します。

推論ベンチマーク: Gaudi2 vs. A100

BLOOMZのメモリ要件に対応するため、Hugging FaceはDeepSpeed-inference（HabanaのDeepSpeedフォーク経由）を利用し、8台のデバイスにわたるモデルとパイプラインの並列化を実装しました。

レイテンシ結果

ベンチマークは16ビット精度、100トークンの貪欲生成、キー・バリューキャッシュを使用して実施しました。結果は、Gaudi2がレイテンシにおいてNvidia A100 80GBを一貫して上回っていることを示しています：

モデル デバイス数 Gaudi2 レイテンシ（秒） A100-80GB レイテンシ（秒） 第1世代 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はモデル並列から大きく恩恵を受けますが、A100は小規模な7Bモデルにおいて単一デバイスでより高速です。

第1世代Gaudiの価格性能

BLOOMZ-7Bモデルに対して、第1世代GaudiはA100と比較して優れた価格性能比を提供します。A100は1時間あたり30ドル超ですが、第1世代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キャッシュのフラグを付与して生成を最適化します。

Sources