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

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

Optimum Habana v1.6 on Habana Gaudi2 では、最新のビジョン言語モデルである BridgeTower のファインチューニングにおいて、A100 と比較してほぼ3倍の高速化を実現しています。ハードウェアアクセラレーションによるデータの読み込みと高速な DDP 実装の2つの新機能がパフォーマンス向上に寄与しています。

これらの技術は、データの読み込みに制約がある他のワークロードにも適用できます。これは、さまざまなタイプのビジョンモデルに頻繁に起こるケースです。この投稿では、BridgeTower のファインチューニングを Habana Gaudi2 と Nvidia A100 80GB で比較するために使用したプロセスとベンチマークを紹介します。また、トランスフォーマーベースのモデルでこれらの機能を簡単に活用する方法も示します。

BridgeTower

最近のビジョン言語（VL）モデルは、さまざまなVLタスクで非常に重要であり、優位性を示しています。最も一般的なアプローチは、それぞれのモダリティから表現を抽出するためにユニモーダルエンコーダを利用することです。その後、これらの表現は融合されるか、クロスモーダルエンコーダに供給されます。VL表現学習のパフォーマンス制約と制限を効果的に扱うために、BridgeTower は複数のブリッジ層を導入し、ユニモーダルエンコーダのトップ層とクロスモーダルエンコーダの各層との間に接続を構築します。これにより、クロスモーダルエンコーダ内の異なる意味レベルで視覚とテキストの表現の効果的なボトムアップのクロスモーダルの整合性と融合が可能になります。

BridgeTower は、4M枚の画像のみで事前学習を行い（詳細は以下を参照）、さまざまな下流のビジョン言語タスクで最先端の性能を実現しています。特に、VQAv2テストセットで78.73%の精度を達成し、前の最先端モデル（METER）を1.09%上回り、同じ事前学習データとほとんど無視できる追加パラメータと計算コストを使用しています。特筆すべきは、モデルをさらにスケーリングすると、BridgeTower は桁違いに大きなデータセットで事前学習されたモデルを凌駕する81.15%の精度を達成します。

• チューリングテスト、中国の部屋、そして大規模言語モデル
• Benfordの法則が機械学習と出会って、偽のTwitterフォロワーを検出する
• ドレスコードの解読👗 自動ファッションアイテム検出のためのディープラーニング

ハードウェア

Nvidia A100 Tensor Core GPU には、Tensor Core テクノロジーの第3世代が搭載されています。最近新しい世代（H100）がリリースされましたが、これはクラウドプロバイダのほとんどで見つけることができる最速のGPUです。ここでは、メモリ帯域幅が40GBのものよりも速い80GBメモリバリアントを使用しています。

Habana Gaudi2 は、Habana Labs によって設計された第2世代のAIハードウェアアクセラレータです。1つのサーバには、96GBのメモリを持つ8つのアクセラレータデバイス（HPU）が搭載されています。より詳細な紹介と、Intel Developer Cloud を介してアクセスする方法については、以前のブログ投稿をご覧ください。市場にある多くのAIアクセラレータとは異なり、Gaudi2 と Optimum Habana を最大限に活用するための高度な機能は非常に簡単に適用できます。これにより、Transformers互換のスクリプトをGaudiに移植するだけで済みます。

ベンチマーク

ベンチマークのために、866Mのパラメータで構成される BridgeTower Large のチェックポイントをファインチューニングします。このチェックポイントは、Conceptual Captions、SBU Captions、MSCOCO Captions、Visual Genome でマスクされた言語モデリング、画像テキストマッチング、画像テキスト対照的損失を使用して英語で事前学習されました。

さらに、このチェックポイントを The New Yorker のカートゥーンと最も投票されたキャプションから成る New Yorker Caption Contest データセットでファインチューニングします。

ハイパーパラメータは、バッチサイズを除いて両方のアクセラレータで同じです：Gaudi2 では40サンプルを、A100 では32サンプルを収めることができました。Gaudi2 用のハイパーパラメータはこちらで、A100 用のハイパーパラメータはこちらで確認できます。

画像を含むデータセットを扱う場合、データの読み込みは頻繁にボトルネックになります。なぜなら、コストの高い操作（画像のデコード、画像の拡張）がCPUで計算され、その後フルサイズの画像がトレーニングデバイスに送信されるからです。理想的には、デバイスには生のバイトのみを送信し、デコードやさまざまな画像変換をデバイス上で行いたいと考えています。しかし、まずは実行を加速するためにデータの読み込みにより多くのリソースを割り当てる方法を見てみましょう。

dataloader_num_workersの活用

画像の読み込みがCPUで行われる場合、データの読み込みを高速化するためには、より多くのサブプロセスをデータの読み込みに割り当てる方法があります。Transformersの TrainingArguments（またはその Optimum Habana 版 GaudiTrainingArguments ）を使用すると、dataloader_num_workers=N 引数を使用して、データの読み込みに割り当てられるCPUのサブプロセスの数（N）を設定できます。

デフォルト値は0であり、これはデータがメインプロセスでロードされることを意味します。これは最適ではないかもしれません、なぜならメインプロセスは多くの処理を管理する必要があるからです。データのロードに完全に専用のサブプロセスを1つ割り当てるために、値を1に設定することができます。複数のサブプロセスが割り当てられる場合、それぞれのサブプロセスがバッチの準備を担当します。これは、ワーカーの数とともにRAMの消費量が増加することを意味します。推奨される設定はCPUコアの数に設定することですが、これらのコアが完全に空いているわけではないため、最適な設定を見つけるために試す必要があります。

以下の2つの実験を実行しましょう:

• データロードがすべてのプロセスと同じプロセスで実行される8つのデバイス間で分散された混合精度（ bfloat16 / float ）ラン（つまり、dataloader_num_workers=0 ）
• データロードに1つの専用のサブプロセスがある8つのデバイス間で分散された混合精度（ bfloat16 / float ）ラン（つまり、dataloader_num_workers=1 ）

次に、Gaudi2とA100で得られたスループットを示します:

まず、Gaudi2はdataloader_num_workers=1ではA100よりもx2.16高速であり、dataloader_num_workers=0ではx2.53高速です。これは、以前に報告した高速化と同等のものです！

次に、データロードにより多くのリソースを割り当てることで簡単に高速化できることがわかります : Gaudi2ではx1.20、A100ではx1.41です。

Gaudi2とA100の両方で複数の専用のサブプロセスでデータロードを実行しましたが、パフォーマンスはdataloader_num_workers=1と同等またはそれ以下でした。したがって、通常はイメージを含む実行を加速するための最初の方法としてdataloader_num_workers=1を使用することがおすすめです！

Gaudi2のTensorboardログはこちらで、A100のTensorboardログはこちらで表示できます。

Optimum Habanaの高速DDP

ハードウェアアクセラレーションされたデータロードの方法を説明する前に、Gaudiで分散実行を高速化する非常に簡単な方法を見てみましょう。Optimum Habanaの新しいリリース、バージョン1.6.0では、ユーザーが使用する分散戦略を選択できる新機能が導入されました:

• distribution_strategy="ddp" はPyTorchのDistributedDataParallel（DDP）を使用します
• distribution_strategy="fast_ddp" はより軽量で通常より高速な実装を使用します

Optimum Habanaの高速DDPは、DDPとは異なり、パラメータ勾配をバケットに分割しません。また、HPUグラフを使用してすべてのプロセスで勾配を収集し、その後（all_reduce操作が実行された後）最小限のホストオーバーヘッドでそれらを更新します。この実装はこちらで確認できます。

単にdistribution_strategy="fast_ddp"（dataloader_num_workers=1 のまま）を使用するだけで、Gaudi2では705.9サンプル/秒を実現できます。 これはDDPよりもx1.10高速であり、A100よりもx2.38高速です！

したがって、2つのトレーニング引数（dataloader_num_workers=1 と distribution_strategy="fast_ddp" ）を追加するだけで、Gaudi2ではx1.33の高速化、A100ではdataloader_num_workers=1の場合にはx2.38の高速化が実現されます。

Optimum Habanaを使用したハードウェアアクセラレーションされたデータロード

さらなる高速化のために、CPUからアクセラレータデバイス（Gaudi2の場合はHPUs、A100の場合はGPU）へのデータロード操作を可能な限り移動します。これは、Gaudi2でHabanaのメディアパイプラインを使用して行うことができます。

データセットが与えられた場合、ほとんどのデータローダーは次の手順に従います:

1. データを取得します（たとえば、JPEG画像がディスク上のどこに保存されているかなど）
2. CPUがエンコードされた画像を読み込みます
3. CPUが画像をデコードします
4. CPUが画像を拡張するための画像変換を適用します
5. 最後に、画像はデバイスに送信されます（通常、これはデータローダー自体では行われません）

デバイスに対して全体のプロセスをCPU上で行い、トレーニングに使用するデータを準備してから送信する代わりに、より効率的なワークフローは、まずエンコードされた画像をデバイスに送信してから画像のデコードと拡張を行うことです：

1. 前と同じ
2. 前と同じ
3. エンコードされた画像がデバイスに送信されます
4. デバイスが画像をデコードします
5. デバイスが画像変換を適用して画像を増幅します

この方法で、デバイスの計算能力を利用して画像のデコードと変換を高速化することができます。ただし、これを行う際には2つの注意点があります：

• デバイスのメモリ使用量が増加するため、十分な空きメモリがない場合はバッチサイズを減らす必要があります。これにより、このアプローチによる高速化が緩和される可能性があります。
• デバイスがデータのロードをCPU上で行う際に使用率が高い（100％またはそれに近い）場合、デバイスで行う際の高速化は期待できません。

Gaudi2でこれを実装するには、Optimum Habanaのコントラスティブな画像テキストの例に、テキストと画像を含むCOCOのようなデータセットで使用できる使用準備のできたメディアパイプラインが提供されています！使用するには、コマンドに--mediapipe_dataloaderを追加するだけです。

興味のある読者のために、Gaudiのドキュメントでより詳細な概要がここに、すべてのサポートされているオペレータのリストがそこに利用可能です。

dataloader_num_workers=1、distribution_strategy="fast_ddp"、およびmediapipe_dataloaderを使用した実行をベンチマークします。これらの最適化はすべて互換性があります：

dataloader_num_workers=1およびdistribution_strategy="fast_ddp"と比較して、前回の実行と比較して追加のx1.14の高速化を得ました。この最終実行は、Gaudi2でのベース実行に比べてx1.51高速です。3つの使用準備のできたトレーニング引数を追加するだけで、A100よりもx2.70高速です！

このベンチマークの再現

このベンチマークを再現するには、まずIntel Developer Cloudを介してGaudi2へのアクセスを取得する必要があります（詳細については、このガイドを参照してください）。

次に、最新バージョンのOptimum Habanaをインストールし、run_bridgetower.pyを実行する必要があります。ここで見つけることができます。以下はその手順です：

pip install optimum[habana]
git clone https://github.com/huggingface/optimum-habana.git
cd optimum-habana/examples/contrastive-image-text
pip install -r requirements.txt

スクリプトを実行するためのベースのコマンドラインは次のとおりです：

python ../gaudi_spawn.py --use_mpi --world_size 8 run_bridgetower.py \
--output_dir /tmp/bridgetower-test \
--model_name_or_path BridgeTower/bridgetower-large-itm-mlm-itc \
--dataset_name jmhessel/newyorker_caption_contest --dataset_config_name matching \
--image_column image --caption_column image_description \
--remove_unused_columns=False \
--do_train --do_eval --do_predict \
--per_device_train_batch_size="40" --per_device_eval_batch_size="16" \
--num_train_epochs 5 \
--learning_rate="1e-5" \
--push_to_hub --report_to tensorboard --hub_model_id bridgetower\
--overwrite_output_dir \
--use_habana --use_lazy_mode --use_hpu_graphs_for_inference --gaudi_config_name Habana/clip \
--throughput_warmup_steps 3 \
--logging_steps 10

これは--dataloader_num_workers 0の場合に対応します。他の設定をテストするには、--dataloader_num_workers 1、--distribution_strategy fast_ddp、および--mediapipe_dataloaderを追加できます。

モデルとTensorboardログをHugging Face Hubにプッシュするには、まずアカウントにログインする必要があります：

huggingface-cli login

A100では、いくつかの小さな変更を加えて同じrun_bridgetower.pyスクリプトを使用できます：

• GaudiTrainerとGaudiTrainingArgumentsをTrainerとTrainingArgumentsに変更する（Transformersから）
• GaudiConfig、gaudi_config、およびHabanaDataloaderTrainerへの参照を削除する
• set_seedをTransformersから直接インポートする：from transformers import set_seed

このベンチマークに表示される結果は、8つのGPUを搭載したNvidia A100 80GB GCPインスタンスで得られました。

--distribution_strategy fast_ddpおよび--mediapipe_dataloaderは、Gaudi2にのみ対応しており、A100では動作しません。

結論

画像を扱う際に、トレーニングワークフローの高速化のために2つの解決策を示しました：データローダーへのリソースの割り当ての増加、およびCPUではなくアクセラレータデバイス上での画像のデコードと拡張。Habana Gaudi2 with Optimum Habanaは、BridgeTowerのようなSOTAビジョン言語モデルのトレーニングで、Nvidia A100 80GB with Transformersと比べてほぼ3倍高速です！さらに、追加のトレーニング引数を提供するだけで簡単に使用できます。

さらに進むために、HPUグラフを使用してさらに高速なモデルトレーニングを行い、Gaudi2上でDeepSpeed ZeRO-3を使用してLLMのトレーニングを加速する方法を紹介することを楽しみにしています。お楽しみに！

最新のAIハードウェアアクセラレータとソフトウェアライブラリを使用して、機械学習のトレーニングと推論ワークフローを高速化することに興味がある場合は、専門家向けアクセラレーションプログラムをご覧ください。Habanaソリューションについて詳しくは、パートナーシップについて読んでこちらからお問い合わせください。AIハードウェアアクセラレータの使いやすさを向上させるためのHugging Faceの取り組みについては、ハードウェアパートナープログラムをご覧ください。

関連トピック

• 高速なトレーニングと推論：Habana Gaudi-2 vs Nvidia A100 80GB
• 巨大言語モデルでの高速推論：Habana Gaudi2アクセラレータ上のBLOOMZ

We will continue to update VoAGI; if you have any questions or suggestions, please contact us!