ビジョン言語モデルの高速化: Habana Gaudi2 上の BridgeTower

BridgeTower モデルの概要

BridgeTower は、uni‑modal エンコーダーのトップレイヤーをクロス‑モーダル エンコーダーの各レイヤーに接続する複数のブリッジ レイヤーを導入し、異なるセマンティック レベルでの効果的なボトムアップ クロス‑モーダル アラインメントと融合を可能にします。Conceptual Captions、SBU Captions、MSCOCO Captions、Visual Genome からの 4M 画像のみで事前学習を行い、VQAv2 test‑std での 78.73% の精度など、ビジョン言語タスクにおいて最先端のパフォーマンスを達成し、追加のパラメータと計算量がほとんどない状態で METER を 1.09% 上回ります。

ハードウェア プラットフォーム

Habana Gaudi2 は、各 96 GB のメモリを備えた 8 つの HPU を提供し、Optimum Habana によって Transformers スクリプトの移植が容易になります。Nvidia H100 は、fp8 ミックスドプリシジョン専用の Transformer Engine と 80 GB のメモリを備えた最新の GPU 世代です。Nvidia A100（80 GB バリアント）は、第3世代の Tensor Core テクノロジーを使用し、クラウド プロバイダーで広く利用可能な最速の GPU のままです。

ベンチマーク セットアップ

私たちは、New Yorker Caption Contest データセット上で BridgeTower Large チェックポイント（866 M パラメータ）をファインチューニングし、すべてのアクセラレータでハイパーパラメータを同じに保ち、サンプル/秒単位でトレーニング スループットを測定しました。

dataloader_num_workers の影響

データローディングのためのサブプロセスを増やすと、すべてのデバイスでスループットが向上します。デバイスごとにバッチサイズ 48 で以下の結果が得られました：

  • Gaudi2 HPU: 601.5 samples/s (workers=0), 747.4 samples/s (workers=1), 768.7 samples/s (workers=2).
  • H100 GPU: 336.5 samples/s (workers=0), 580.1 samples/s (workers=1), 602.1 samples/s (workers=2).
  • A100 GPU: 227.5 samples/s (workers=0), 339.7 samples/s (workers=1), 345.4 samples/s (workers=2). これらの結果から、2 つのワーカーを使用した場合、Gaudi2 は H100 より x1.28 倍、A100 より x2.23 倍 高速であることが示されており、ワーカー数を増やすと Gaudi2 で x1.28 倍、H100 で x1.79 倍、A100 で x1.52 倍 のスピードアップが得られます。

Optimum Habana を使用したハードウェア アクセラレートされたデータ ローディングの影響

--mediapipe_dataloader フラグ（Gaudi2 でのみ利用可能）を介して画像のデコードとオーグメンテーションをアクセラレータにオフロードすると、スループットがさらに向上します。2 つのデータローダー ワーカーを使用した同じ実験では以下の結果が得られました：

  • Gaudi2 HPU: 847.7 samples/s (workers=2 + mediapipe_dataloader).
  • H100 と A100: 適用不可（このフラグはこれらの GPU で動作しません）。 ベースラン（workers=0）と比較すると、これは Gaudi2 で x1.41 倍のスピードアップを示し、2 つのワーカーとメディア パイプラインを使用した場合、Gaudi2 は H100 より x1.41 倍、A100 より x2.45 倍 高速になります。

このベンチマークの再現

結果を再現するには、Intel Developer Cloud を通じて Gaudi2 にアクセスし、最新の Optimum Habana をインストールし、リポジトリをクローンし、依存関係をインストールして、適切な引数で提供された run_bridgetower.py スクリプトを実行します。Gaudi2 の場合、ベース コマンドには --use_habana --use_lazy_mode --use_hpu_graphs_for_inference --gaudi_config_name Habana/clip が含まれます。--dataloader_num_workers N と --mediapipe_dataloader を追加すると、他の構成をテストできます。A100 と H100 の場合、Habana 固有のクラスを標準の Transformers Trainer と TrainingArguments に置き換え、Gaudi の設定参照を削除し、Transformers から set_seed をインポートします。

結論

ビジョン言語モデルのトレーニングでは、データローダー ワーカーを増やし、画像のデコード/オーグメンテーションをアクセラレータにオフロードするという 2 つのシンプルな手法により、大幅なスピードアップが得られます。Habana Gaudi2 上の Optimum Habana を使用すると、BridgeTower のファインチューニング時に Nvidia H100 の約 1.4×、Nvidia A100 80GB の約 2.5× の速度を達成でき、これにはわずかな追加のトレーニング引数のみが必要です。

Sources