Hugging Face Transformers と Habana Gaudi を使用した BERT の事前学習

TL;DR

Hugging Face は、AWS の Habana Gaudi ベースの DL1 インスタンスを使用して BERT-base をスクラッチから事前学習するための技術ガイドを提供しています。optimum-habana ライブラリを活用することで、同等の NVIDIA V100 GPU トレーニング環境と比較して 25% のコスト削減を実現しています。

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

事前学習は AWS DL1 インスタンス上で実行され、8 つの HPU コアを備えています。ソフトウェアスタックは以下の Hugging Face ライブラリを統合しています：

  • Transformers: モデルアーキテクチャのコアライブラリです。
  • Optimum Habana: GaudiTrainer と GaudiTrainingArguments を提供し、これは標準の Hugging Face Trainer のラッパーで、Habana Gaudi ハードウェア向けに特化しています。
  • Datasets: 事前学習コーパスのロードと管理に使用されます。

事前学習ワークフロー

事前学習プロセスは、CPU 集中型のデータ準備と HPU 集中型のモデル訓練に分かれています。

1. データ準備とトークナイゼーション

データ準備は CPU 最適化インスタンス（例: AWS c6i.12xlarge）で実行され、DL1 インスタンスに移行する前に行われます。ワークフローは以下を含みます：

  • Dataset Merging: Wikipedia（20220301.en スプリット）と BookCorpus を単一のデータセットに結合します。
  • Tokenizer Training: 語彙サイズ 32,000 トークンの新しい BertTokenizerFast をスクラッチから訓練します。
  • Preprocessing: テキストをトークナイズし、512 トークンのチャンク（BERT の最大シーケンス長）にまとめ、上限を超える文書は切り捨てます。

2. マスク付き言語モデリング（MLM）によるモデル訓練

BERT はマスク付き言語モデリング（MLM）を使用して事前学習されます。このタスクは、モデルが双方向コンテキストを用いて文中の隠れた単語を予測するものです。

DL1 インスタンスの 8 HPU コアを活用するため、トレーニングはスクリプト（run_mlm.py）と optimum-habana の DistributedRunner（または gaudi_spawn.py）を使用した分散データパラレル訓練として実装されています。

パフォーマンスとコスト分析

訓練時間とコスト

Hugging Face は、グローバルバッチサイズ 256 の 100,000 ステップの実行をベンチマークし、約 12.5 時間で完了しました。これを元の BERT 事前学習要件である 1,000,000 ステップに外挿すると：

  • Total Estimated Time: 125 時間。
  • Total Estimated Cost: AWS 上の Habana Gaudi を使用して約 $1,650。

GPU ベースラインとの比較

DGX-2（16 NVIDIA V100 GPU）での最速 BERT 事前学習に関する DeepSpeed 記録と比較すると、Habana Gaudi アプローチは大幅なコスト削減を提供します：

  • GPU Cost Estimate: 参照として AWS p3dn.24xlarge インスタンス（8x V100 32GB）を使用した場合、同様の構成で約 $2,075 のコストがかかります。
  • Cost Reduction: Habana Gaudi は V100 ベースの構成に対して 25% のコスト削減を実現します。

特筆すべきは、DeepSpeed 最適化がない場合、GPU 訓練コストはさらに約 $3,000〜$4,000 に増加する可能性があると著者が指摘している点です。

ドメイン適応への示唆

スクラッチからの事前学習により、組織はモデルを特定の言語や専門領域に適応させることができます。情報源によれば、このプロセスは汎用 BERT モデルを使用する場合と比較して、モデル精度を最大 10% 向上させることができます。

Sources