Cloud TPU（TF 2.x）での HuggingFace GPT2 のトレーニング

---

Cloud TPU に慣れていない場合は、クイックスタートを参照して、TPU VM の作成方法を学習することをおすすめします。

このチュートリアルでは、Cloud TPU で HuggingFace GPT2 モデルをトレーニングする方法を説明します。

警告: このチュートリアルでは、サードパーティのデータセットを使用します。Google は、このデータセットの有効性またはその他の側面について、いかなる表明、保証もいたしません。

目標

• Cloud TPU の作成
• 依存関係をインストールする
• トレーニング ジョブを実行する

費用

このドキュメントでは、Google Cloud の次の課金対象のコンポーネントを使用します。

• Compute Engine
• Cloud TPU

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。 新しい Google Cloud ユーザーは無料トライアルをご利用いただける場合があります。

始める前に

このチュートリアルを開始する前に、Google Cloud プロジェクトが正しく設定されていることを確認します。

1. Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.

2. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

   注: この手順で作成するリソースをそのまま保持する予定でない場合、既存のプロジェクトを選択するのではなく、新しいプロジェクトを作成してください。チュートリアルの終了後にそのプロジェクトを削除すれば、プロジェクトに関連するすべてのリソースを削除できます。

   プロジェクト セレクタに移動

3. Google Cloud プロジェクトで課金が有効になっていることを確認します。

4. Google Cloud Console の [プロジェクト セレクタ] ページで、Google Cloud プロジェクトを選択または作成します。

   注: この手順で作成するリソースをそのまま保持する予定でない場合、既存のプロジェクトを選択するのではなく、新しいプロジェクトを作成してください。チュートリアルの終了後にそのプロジェクトを削除すれば、プロジェクトに関連するすべてのリソースを削除できます。

   プロジェクト セレクタに移動

5. Google Cloud プロジェクトで課金が有効になっていることを確認します。

このチュートリアルでは、Google Cloud の課金対象となるコンポーネントを使用します。費用を見積もるには、Cloud TPU の料金ページを確認してください。不要な課金を回避するために、このチュートリアルを完了したら、作成したリソースを必ずクリーンアップしてください。

重要: すべてのリソース（Cloud TPU VM と Cloud Storage バケット）を同じリージョン / ゾーンに設定して、ネットワーク レイテンシとネットワーク コストを削減します。TPU VM は、リージョン内の下位区分である特定のゾーンに配置されます。

Cloud TPU で HuggingFace GPT2 をトレーニングする

1. Cloud Shell ウィンドウを開きます。

   Cloud Shell を開く

2. プロジェクト ID 用の環境変数を設定します。

   export PROJECT_ID=your-project-id
   

3. Cloud TPU を作成する Google Cloud プロジェクトを使用するように Google Cloud CLI を構成します。

   gcloud config set project ${PROJECT_ID}
   

   このコマンドを新しい Cloud Shell VM で初めて実行すると、Authorize Cloud Shell ページが表示されます。ページの下部にある [Authorize] をクリックして、gcloud に認証情報を使用した Google Cloud API の呼び出しを許可します。

4. Cloud TPU プロジェクトのサービス アカウントを作成します。

   サービス アカウントにより、Cloud TPU サービスが他の Google Cloud サービスにアクセスできるようになります。

   $ gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
   

   このコマンドでは、Cloud TPU サービス アカウントを次の形式で返します。

   [e-mail]
   

Cloud TPU の作成

1. gcloud コマンドを使用して Cloud TPU VM を作成します。次のコマンドは、v4-8 TPU を作成します。--accelerator-type フラグを Pod スライス タイプ（v4-32 など）に設定することで、TPU Pod スライスを作成することもできます。

   $ gcloud compute tpus tpu-vm create hf-gpt2 \
     --zone=us-central2-b \
     --accelerator-type=v4-8 \
     --version=tpu-vm-tf-2.16.1-pjrt
   

   コマンドフラグの説明

   zone

   Cloud TPU を作成するゾーン。

   accelerator-type

   アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。

   version

   Cloud TPU ソフトウェアのバージョン。

   注: Google Cloud プロジェクトが複数ある場合は、--project フラグを使用して、Cloud TPU を作成する Google Cloud プロジェクト ID を指定する必要があります。

2. 次の ssh コマンドを実行して、Cloud TPU VM に接続します。

   gcloud compute tpus tpu-vm ssh hf-gpt2 --zone=us-central2-b
   

   注: Cloud TPU VM に接続すると、シェル プロンプトが username@projectname からusername@vm-name に変わります。
   キーポイント:これ以降、接頭辞 (vm) $ は Compute Engine VM インスタンスでコマンドを実行する必要があることを意味します。

依存関係をインストールする

1. HuggingFace Transformers リポジトリのクローンを作成します。

   (vm)$ cd /tmp
   (vm)$ git clone https://github.com/huggingface/transformers.git
   (vm)$ cd transformers
   

2. 依存関係をインストールします。

   (vm)$ pip install .
   (vm)$ pip install -r examples/tensorflow/_tests_requirements.txt
   (vm)$ cd /tmp/transformers/examples/tensorflow/language-modeling
   (vm)$ pip install -r requirements.txt
   

3. 一時ディレクトリを作成します。

   (vm)$ mkdir /tmp/gpt2-wikitext
   

4. TPU を作成するときに、--version パラメータを -pjrt で終わるバージョンに設定した場合は、次の環境変数を設定して PJRT ランタイムを有効にします。

     (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
     (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
   

トレーニング スクリプトを実行する

(vm)$ python3 run_clm.py \
  --model_name_or_path distilgpt2 \
  --max_train_samples 1000 \
  --max_eval_samples 100 \
  --num_train_epochs 1 \
  --output_dir /tmp/gpt2-wikitext \
  --dataset_name wikitext \
  --dataset_config_name wikitext-103-raw-v1

コマンドフラグの説明

model_name_or_path

トレーニングするモデルの名前。

max_train_samples

トレーニングに使用するサンプルの最大数。

max_eval_samples

評価に使用するサンプルの最大数。

num_train_epochs

モデルをトレーニングするためのエポックの数。

output_dir

トレーニング スクリプトの出力ディレクトリ。

dataset_name

使用するデータセットの名前。

dataset_config_name

データセット構成名

トレーニングが完了すると、次のようなメッセージが表示されます。

  125/125 [============================] - ETA: 0s - loss: 3.61762023-07-07 21:38:17.902850: W tensorflow/core/framework/dataset.cc:956] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
  125/125 [============================] - 763s 6s/step - loss: 3.6176 - val_loss: 3.4233
  Configuration saved in /tmp/gpt2-wikitext/config.json
  Configuration saved in /tmp/gpt2-wikitext/generation_config.json
  Model weights saved in /tmp/gpt2-wikitext/tf_model.h5
  D0707 21:38:45.640973681   12027 init.cc:191]                          grpc_shutdown starts clean-up now
  

クリーンアップ

1. TPU VM インスタンスから接続を切断します。

   (vm)$ exit
   

   プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

2. TPU リソースを削除します。

   $ gcloud compute tpus tpu-vm delete hf-gpt2 \
   --zone=us-central2-b
   

次のステップ

他のサポートされているリファレンス モデルのいずれかを試す。