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


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

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

目標

  • 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. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  3. Make sure that billing is enabled for your Google Cloud project.

  4. In the Google Cloud console, on the project selector page, select or create a Google Cloud project.

    Go to project selector

  5. Make sure that billing is enabled for your Google Cloud project.

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

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 サービス アカウントを次の形式で返します。

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

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.17.0-pjrt
    

    コマンドフラグの説明

    zone
    Cloud TPU を作成するゾーン
    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
    version
    Cloud TPU ソフトウェアのバージョン
  2. 次の ssh コマンドを実行して Cloud TPU VM に接続します。

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

依存関係のインストール

  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
    

次のステップ

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