Cloud TPU v5e トレーニング

Cloud TPU v5e は、Google Cloud の最新の生成 AI アクセラレータです。Pod あたり 256 チップのフットプリントが小さい v5e は、トランスフォーマー、text-to-image、および、畳み込みニューラル ネットワーク（CNN）のトレーニング、微調整、サービス提供に適した最適なプロダクトとなるように最適化されています。 Cloud TPU v5e をサービングに使用する方法については、v5e を使用した推論をご覧ください。

Cloud TPU v5e TPU のハードウェアと構成の詳細については、TPU v5e をご覧ください。

始める

以降のセクションでは、TPU v5e の使用を開始する方法について説明します。

リクエストの割り当て

トレーニングに TPU v5e を使用するには、割り当てが必要です。オンデマンド TPU、予約 TPU、TPU Spot VM には、さまざまな割り当てタイプがあります。推論に TPU v5e を使用する場合は、個別の割り当てが必要です。割り当ての詳細については、割り当てをご覧ください。TPU v5e の割り当てをリクエストするには、Cloud セールスにお問い合わせください。

Google Cloud アカウントとプロジェクトを作成する

Cloud TPU を使用するには、Google Cloud アカウントとプロジェクトが必要です。詳細については、Cloud TPU 環境を設定するをご覧ください。

注: TPU v5e を使用するには、gcloud のアルファ版コンポーネント をインストールする必要があります。

Cloud TPU の作成

queued-resource create コマンドを使用して、Cloud TPU v5es をキューに格納されたリソースとしてプロビジョニングすることをおすすめします。詳細については、キューに入れられたリソースを管理するをご覧ください。

Create Node API（gcloud alpha compute tpus tpu-vm create）を使用して Cloud TPU v5es をプロビジョニングすることもできます。詳細については、TPU リソースの管理をご覧ください。

トレーニングに使用できる v5e 構成の詳細については、トレーニング用の Cloud TPU v5e タイプをご覧ください。

フレームワークの設定

このセクションでは、TPU v5e で JAX または PyTorch を使用したカスタムモデルのトレーニングの一般的な設定プロセスについて説明します。TensorFlow のサポートは、tpu-vm-tf-2.17.0-pjrt と tpu-vm-tf-2.17.0-pod-pjrt の TPU ランタイム バージョンで利用できます。

推論の設定手順については、v5e 推論の概要をご覧ください。

JAX を設定する

スライス形状が 8 チップを超える場合、1 つのスライスに複数の VM があります。この場合、SSH を使用して個別にログインすることなく、--worker=all フラグを使用してすべての TPU VM に 1 つのステップでインストールを実行する必要があります。

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

次のコマンドを実行して、デバイスの数を確認できます（ここに表示されている出力は、v5litepod-16 スライスで生成されたものです）。このコードは、JAX が Cloud TPU TensorCore を認識し、基本オペレーションを実行できることを確認することで、すべてが正しくインストールされていることをテストします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

出力は次のようになります。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() は、指定されたスライス内のチップの合計数を示します。jax.local_device_count() は、このスライス内の単一の VM からアクセス可能なチップの数を示します。

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

出力は次のようになります。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

このドキュメントの JAX チュートリアルを試して、JAX を使用した v5e トレーニングを開始します。

PyTorch を設定する

v5e は PJRT ランタイムのみをサポートしているのでご注意ください。PyTorch 2.1 以降では、すべての TPU バージョンのデフォルト ランタイムとして PJRT が使用されます。

このセクションでは、すべてのワーカー用のコマンドで PyTorch/XLA を使用して v5e 上で PJRT の使用を開始する方法について説明します。

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

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

torch、torch_xla、torchvision のホイール（pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222 など）を取り付ける際にエラーが発生した場合は、次のコマンドを使用してバージョンをダウングレードします。

pip3 install setuptools==62.1.0

PJRT でスクリプトを実行する

注: 大きい割り当てが頻繁に発生するモデルの場合、tcmallocを使用すると、デフォルトの malloc 実装よりもトレーニング時間が大幅に短縮されるため、TPU VM で使用されるデフォルトの malloc は tcmalloc です。しかし、ワークロードによっては（例えば、埋め込みテーブルへの割り当てが非常に大きい DLRM など）tcmalloc では速度が低下する可能性があります。その場合は、代わりにデフォルトの malloc を使用して次の変数を解除してみてください。

unset LD_PRELOAD

Python スクリプトを使用して v5e VM で計算を行う例を次に示します。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

これにより、次のような出力が生成されます。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

このドキュメントの PyTorch チュートリアルを試して、PyTorch を使用した v5e トレーニングを開始します。

セッションの終了時に TPU とキューに格納されたリソースを削除します。キューに格納されたリソースを削除するには、スライスを削除してから、キューに格納されたリソースを削除する、2 つのステップで行います。

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

これら 2 つのステップは、FAILED 状態にある、キューに格納されたリソース リクエストを削除するためにも使用できます。

JAX/FLAX の例

以降のセクションでは、TPU v5e で JAX モデルと FLAX モデルをトレーニングする方法の例について説明します。

v5e で ImageNet をトレーニングする

このチュートリアルでは、架空の入力データを使用して v5e で ImageNet をトレーニングする方法について説明します。実際のデータを使用する場合は、GitHub の README ファイルをご覧ください。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   QueuedResource が ACTIVE 状態の場合、出力は次のようになります。

    state: ACTIVE
   

3. 最新バージョンの JAX と jaxlib をインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   注: jax>=0.4.8 を使用することを示すエラーが表示された場合は、このメッセージは無視してかまいません。

4. ImageNet モデルのクローンを作成し、対応する要件をインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'
   

5. 架空のデータを生成するには、データセットのサイズに関する情報が必要です。これは、ImageNet データセットのメタデータから収集できます。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} --zone=${ZONE} --worker=all --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'
   

モデルをトレーニングする

前の手順をすべて完了したら、モデルをトレーニングできます。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Hugging Face FLAX モデル

FLAX で実装された Hugging Face モデルは Cloud TPU v5e ですぐに使用できます。このセクションでは、一般的なモデルの実行手順について説明します。

Imagenette で ViT をトレーニングする

このチュートリアルでは、Cloud TPU v5e で Fast AI Imagenette データセットを使用して、HuggingFace から Vision Transformer（ViT）モデルをトレーニングする方法について説明します。

ViT モデルは、畳み込みネットワークと比較して優れた結果で ImageNet で Transformer エンコーダをトレーニングした最初のモデルです。詳しくは、次のリソースをご覧ください。

• ViT の概要
• ViT コードリソース

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。

    state: ACTIVE
   

3. JAX とそのライブラリをインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. Hugging　Face のリポジトリをダウンロードし、要件をインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.17.0 && pip install -r examples/flax/vision/requirements.txt'
   

5. Imagenette データセットをダウンロードします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
   

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

ViT のベンチマーク結果

トレーニング スクリプトは、v5litepod-4、v5litepod-16、v5litepod-64 で実行されました。次の表に、異なるアクセラレータ タイプでのスループットを示します。

アクセラレータ タイプ	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
グローバル バッチサイズ	32	128	512
スループット（例/秒）	263.40	429.34	470.71

Pokémon 上で Diffusion をトレーニングする

このチュートリアルでは、Cloud TPU v5e で Pokémon データセットを使用して、HuggingFace から Stable Diffusion モデルをトレーニングする方法について説明します。

Stable Diffusion モデルは、テキスト入力からフォトリアリスティックな画像を生成する、潜在的 text-to-image モデルです。詳しくは、次のリソースをご覧ください。

• Stable Diffusion の概要
• Stable Dufffusion のコードソース

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。

    state: ACTIVE
   

3. JAX とそのライブラリをインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. HuggingFace のリポジトリをダウンロードし、要件をインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.17.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
   

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

拡散のベンチマーク結果

トレーニング スクリプトは、v5litepod-4、v5litepod-16、v5litepod-64 で実行されました。 次の表では、スループットを示します。

アクセラレータ タイプ	v5litepod-4	v5litepod-16	v5litepod-64
トレーニング ステップ	1500	1500	1500
グローバル バッチサイズ	32	64	128
スループット（例/秒）	36.53	43.71	49.36

OSCAR データセットで GPT2 をトレーニングする

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

GPT2 は、ヒューマン ラベリングのない未加工のテキストで事前にトレーニングされた Transformer モデルです。文の次の単語を予測するようにトレーニングされました。詳しくは、次のリソースをご覧ください。

• GPT2 の概要
• GPT2 のコードソース

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。

    state: ACTIVE
   

3. JAX とそのライブラリをインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

4. HuggingFace のリポジトリをダウンロードし、要件をインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/language-modeling/requirements.txt'
   

5. モデルをトレーニングするための構成ファイルをダウンロードします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'
   

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

GPT2 のベンチマーク結果

トレーニング スクリプトは、v5litepod-4、v5litepod-16、v5litepod-64 で実行されました。 次の表では、スループットを示します。

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
グローバル バッチサイズ	64	64	64
スループット（例/秒）	74.60	72.97	72.62

PyTorch/XLA

以下のセクションでは、TPU v5e で PyTorch/XLA モデルをトレーニングする方法の例について説明します。

PJRT ランタイムを使用して ResNet をトレーニングする

PyTorch/XLA は、PyTorch 2.0 以降、XRT から PjRt に移行しています。PyTorch/XLA トレーニング ワークロード用に v5e を設定する手順が更新されています。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --{QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。

    state: ACTIVE
   

3. Torch/XLA 固有の依存関係をインストールする

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         sudo apt-get update -y
         sudo apt-get install libomp5 -y
         pip3 install mkl mkl-include
         pip3 install tf-nightly tb-nightly tbp-nightly
         pip3 install numpy
         sudo apt-get install libopenblas-dev -y
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

ResNet モデルをトレーニングする

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

ベンチマークの結果

次の表に、ベンチマークのスループットを示します。

アクセラレータ タイプ	スループット（例/秒）
v5litepod-4	4240 ex/s
v5litepod-16	10,810 ex/s
v5litepod-64	46,154 ex/s

v5e で GPT2 をトレーニングする

このチュートリアルでは、wikitext データセットを使用し、PyTorch/XLA で HuggingFace のリポジトリを使用して、v5e で GPT2 を実行する方法を説明します。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。

   state: ACTIVE
   

3. PyTorch/XLA の依存関係をインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         sudo apt-get -y update
         sudo apt install -y libopenblas-base
         pip3 install torchvision
         pip3 uninstall -y torch
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

4. HuggingFace のリポジトリをダウンロードし、要件をインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         git clone https://github.com/pytorch/xla.git
         pip install --upgrade accelerate
         git clone https://github.com/huggingface/transformers.git
         cd transformers
         git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
         pip install .
         git log -1
         pip install datasets evaluate scikit-learn
         '
   

5. 事前トレーニング済みモデルの構成をダウンロードします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command='
         gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/
         gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
   

モデルのトレーニング

バッチサイズ 16 を使用して 2B モデルをトレーニングします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

gcloud compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

ベンチマークの結果

トレーニング スクリプトは、v5litepod-4、v5litepod-16、v5litepod-64 で実行されました。次の表に、さまざまなアクセラレータ タイプのベンチマーク スループットを示します。

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
構成	6 億	20 億	16B
グローバル バッチサイズ	64	128	256
スループット（例/秒）	66	77	31

v5e で ViT をトレーニングする

このチュートリアルでは、cifar10 データセット上の PyTorch/XLA で HuggingFace のリポジトリを使用して、v5e で VIT を実行する方法について説明します。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=tpu-name
   export QUEUED_RESOURCE_ID=queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_UNTIL_DURATION=1d
   

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --valid-until-duration=${VALID_UNTIL_DURATION} \
      --service-account=${SERVICE_ACCOUNT} \
      --${QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。

    state: ACTIVE
   

3. PyTorch/XLA の依存関係をインストールする

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all
      --command='
         sudo apt-get update -y
         sudo apt-get install libomp5 -y
         pip3 install mkl mkl-include
         pip3 install tf-nightly tb-nightly tbp-nightly
         pip3 install numpy
         sudo apt-get install libopenblas-dev -y
         pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
         pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
   

4. HuggingFace のリポジトリをダウンロードし、要件をインストールします。

      gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --worker=all \
      --command="
         git clone https://github.com/suexu1025/transformers.git vittransformers; \
         cd vittransformers; \
         pip3 install .; \
         pip3 install datasets; \
         wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"
   

モデルのトレーニング

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

gcloud compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

ベンチマークの結果

次の表に、さまざまなアクセラレータ タイプのベンチマーク スループットを示します。

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
グローバル バッチサイズ	32	128	512
スループット（例/秒）	201	657	2,844

TensorFlow 2.x

以降のセクションでは、TPU v5e で TensorFlow 2.x モデルをトレーニングする方法の例について説明します。

単一ホスト v5e で Resnet をトレーニングする

このチュートリアルでは、架空のデータセットを使用して v5litepod-4 または v5litepod-8 で ImageNet をトレーニングする方法について説明します。別のデータセットを使用する場合は、データセットの準備をご覧ください。

設定

1. 環境変数を作成します。

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5litepod-4
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type
   

   ACCELERATOR_TYPE は、v5litepod-4 または v5litepod-8 のいずれかになります。

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。キューに格納されたリソースの状態を確認するには、次のコマンドを使用します。

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

3. SSH を使用して TPU に接続します

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

4. いくつかの環境変数を設定します

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
   export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
   

5. モデル リポジトリのディレクトリに移動し、要件をインストールします。

   cd ${MODELS_REPO} && git checkout r2.15.0
   pip install -r official/requirements.txt
   

モデルのトレーニング

トレーニング スクリプトを実行します。

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

TPU とキューに格納されたリソースを削除する

1. TPU を削除します

   gcloud compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

2. キューに格納されたリソース リクエストを削除します

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

マルチホスト v5e で Resnet をトレーニングする

このチュートリアルでは、架空のデータセットを使用して v5litepod-16 以上での ImageNet をトレーニングする方法について説明します。別のデータセットを使用する場合は、データセットの準備をご覧ください。

1. 環境変数を作成します。

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5litepod-16
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type
   

   ACCELERATOR_TYPE は v5litepod-16 かそれ以上にすることができます。

2. TPU リソースを作成します。

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   注: QUOTA_TYPE　フラグは、reserved または　best-effort　のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。

   キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。キューに格納されたリソースの状態を確認するには、次のコマンドを使用します。

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

3. SSH を使用して TPU（ワーカーゼロ）に接続します

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

4. いくつかの環境変数を設定します

   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export TPU_LOAD_LIBRARY=0
   export TPU_NAME=your_tpu_name
   

5. モデル リポジトリのディレクトリに移動し、要件をインストールします。

    cd $MODELS_REPO && git checkout r2.15.0
    pip install -r official/requirements.txt
   

モデルのトレーニング

トレーニング スクリプトを実行します。

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

TPU とキューに格納されたリソースを削除する

1. TPU を削除します

   gcloud compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

2. キューに格納されたリソース リクエストを削除します

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet