Cloud TPU v5e トレーニング
Cloud TPU v5e は、Google Cloud の最新の生成 AI アクセラレータです。Pod あたり 256 チップのフットプリントが小さい v5e は、トランスフォーマー、text-to-image、および、畳み込みニューラル ネットワーク（CNN）のトレーニング、微調整、サービス提供に適した最適なプロダクトとなるように最適化されています。
Cloud TPU v5e のコンセプト、システム アーキテクチャ、構成
Cloud TPU を初めて使用する場合は、TPU ドキュメント ホームをご覧ください。
Cloud TPU の一般的なコンセプト（スライス、ホスト、チップ、TensorCore など）と Cloud TPU システム アーキテクチャについては、Cloud TPU システム アーキテクチャのページをご覧ください。
各 Cloud TPU のバージョンでは、トレーニングと推論のために特定のアクセラレータ タイプが必要です。これらのアクセラレータ タイプについては、v5e 構成のセクションをご覧ください。
推論
推論とは、トレーニング済みモデルを使用して新しいデータについて予測を行うプロセスです。サービス提供プロセスで使用されます。
スライス
スライスは、高速チップ間相互接続（ICI）で接続された同じ Pod 内にあるチップの集合体を表します。v5e は 2D スライス形状です。サポートされているスライス形状については、v5e 構成セクションの表をご覧ください。
チップ形状とチップ トポロジもスライス形状を指します。
サービス提供
サービス提供とは、トレーニング済み機械学習モデルを本番環境にデプロイするプロセスであり、本番環境を使用して予測や決定を行うために使用できます。提供には、レイテンシとサービスレベルの可用性が重要です。
単一ホストとマルチホスト
ホストは、VM を実行する物理コンピュータ（CPU）です。1 つのホストが同時に複数の VM を実行できます。
8 チップ未満のスライスでは、最大で 1 つのホストを使用します。8 チップを超えるスライスは、複数のホストにアクセスでき、複数のホストを使用して分散トレーニングを実行できます。スライスとチップの詳細については、TPU システム アーキテクチャのページをご覧ください。
v5e は、マルチホスト トレーニングとマルチホスト推論をサポートしています（SAX を使用）。
TPU VM
基盤となる TPU にアクセスできる Linux を実行している仮想マシン。v5e TPU の場合、各 TPU VM はユーザー指定のアクセラレータ タイプに応じて、1、4、8 チップに直接アクセスできます。TPU VM は、ワーカーとも呼ばれます。
ワーカー
TPU VM をご覧ください。
始める
v5e TPU ハードウェアの詳細については、システム アーキテクチャをご覧ください。
容量の確保
Cloud セールスに連絡して、AI ワークロードに対する Cloud TPU v5e の使用を開始しましょう。
Google Cloud プロジェクトを準備する
- Google アカウントにログインします。Google アカウントをまだ持っていない場合は、新しいアカウントを作成します。
- Google Cloud コンソールで、[プロジェクト セレクタ] ページから Google Cloud プロジェクトを選択するか作成します。
すべての Google Cloud の使用にはお支払い情報の設定が必要なため、プロジェクトで課金が有効になっていることを確認します。
gcloud alpha components をインストールします。
既存の
gcloud alphaコンポーネントを再利用する TPU ユーザーの場合、コンポーネントを更新して関連するコマンドとフラグがサポートされていることを確認します。gcloud components updateCloud Shell で、次の
gcloudコマンドを使用して TPU API を有効にします。 （Google Cloud コンソールから有効にすることもできます）。gcloud services enable tpu.googleapis.comTPU サービス アカウントを有効にします。
サービス アカウントにより、Cloud TPU サービスが他の Google Cloud サービスにアクセスできるようになります。ユーザー管理のサービス アカウントは、Google Cloud で推奨される方法です。ガイドに沿って作成し、ロールを付与します。 次のロールが必要です。
- TPU 管理者
- ストレージ管理者: Cloud Storage へのアクセスに必要
- ログ書き込み: Logging API でログを書き込むために必要
- モニタリング指標の書き込み: Cloud Monitoring に指標を書き込むために必要
プロジェクトとゾーンを構成します。
プロジェクト ID は、Cloud コンソールに表示されるプロジェクトの名前です。
export PROJECT_ID=your-project-id export ZONE=us-west4-a gcloud compute tpus tpu-vm service-identity create --zone=${ZONE} gcloud auth login gcloud config set project ${PROJECT_ID} gcloud config set compute/zone ${ZONE}
Cloud TPU 環境をプロビジョニングする
queued-resource create コマンドを使用して、Cloud TPU v5es をキューに格納されたリソースとしてプロビジョニングすることをおすすめします。ただし、Create Node API（gcloud alpha compute tpus tpu-vm create）を使用して Cloud TPU v5es をプロビジョニングすることもできます。
環境変数を作成する
TPU の作成に必要な環境変数を設定します。
次のリストの変数（赤色）を、トレーニング ジョブまたは推論ジョブに使用する値に置き換えます。
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
| 変数 | Description |
| PROJECT_ID | Google Cloud プロジェクト名 |
| ACCELERATOR_TYPE | サポートされているアクセラレータ タイプについては、アクセラレータ タイプのセクションをご覧ください。 |
| ZONE | 容量はすべて us-west4-a です。 |
| RUNTIME_VERSION | v2-alpha-tpuv5-lite を使用します。 |
| SERVICE_ACCOUNT | これは、Google Cloud コンソール -> IAM] -> サービス アカウント で確認できるサービス アカウントのアドレスです。 例: tpu-service-account@myprojectID.iam.gserviceaccount.com |
| TPU_NAME | キューに格納されたリソース リクエストの割り当て時に作成される TPU のユーザー割り当て ID。 |
| QUEUED_RESOURCE_ID | キューに格納されたリソース リクエストのユーザー割り当て ID。キューに格納されたリソースについては、キューに格納されたリソースのドキュメントをご覧ください。 |
| QUOTA_TYPE | reserved、best-effort のいずれかを設定できます。どちらも指定されていない場合、QUOTA_TYPE はデフォルトで on-demand になります。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。 |
| VALID_UNTIL_DURATION | リクエストが有効である期間。さまざまな有効期間の詳細については、キューに格納されたリソースをご覧ください。 |
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}
キューに格納されたリソースが正常に作成されると、response フィールド内の状態は WAITING_FOR_RESOURCES または FAILED のいずれかになります。キューに格納されたリソースが WAITING_FOR_RESOURCES 状態の場合は、キューに格納されたリソースが予備検証に合格し、容量を待っていることを意味します。容量が利用可能になると、リクエストは PROVISIONING に移行します。WAITING_FOR_RESOURCES の状態であることは、割り当てが保証されることを意味するわけではなく、WAITING_FOR_RESOURCES ステータスから ACTIVE ステータスに変更されるまでには時間がかかることがあります。キューに格納されたリソースが FAILED 状態の場合、失敗の理由が出力に表示されます。--valid-until-duration にリクエストが入力されない場合、リクエストは期限切れになり、状態は「FAILED」になります。
キューに格納されたリソースが ACTIVE 状態になると、SSH を使用して TPU VM にアクセスできるようになります。
[list](/tpu/docs/managing-tpus-tpu-vm) コマンドまたは [describe](/tpu/docs/managing-tpus-tpu-vm) コマンドを使用して、キューに格納されたリソースのステータスをクエリします。
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE}
状態は、キューに入れられたリソースのステータスを表します。状態は次のように定義されます。
| 状態 | Description |
WAITING_FOR_RESOURCES |
キューに格納されたリソース create コマンドを受信しました。容量が使用可能になり次第、プロビジョニングを開始します。
|
PROVISIONING |
TPU スライスがプロビジョニング中です。 |
ACTIVE |
すべての TPU がプロビジョニングされ、使用できる状態になります。起動スクリプトを指定した場合、キューに格納されたリソース状態が ACTIVE に遷移すると、すべての TPU でスクリプトの実行が開始されます。
|
FAILED |
スライスをプロビジョニングできませんでした。 |
SUSPENDING |
1 つ以上のスライスを削除中です。 |
SUSPENDED |
基盤となるスライスはすべて削除されますが、キューに格納されたリソースは明示的に削除されるまでそのまま残ります。現時点では、キューに格納された一時停止中のリソースを再開することはできません。削除する必要があります。 |
DELETING |
キューに格納されたリソースを削除中です。 |
SSH を使用して TPU VM に接続する
次のセクションでは、TPU スライスの各 TPU VM にバイナリをインストールしてコードを実行する方法について説明します。このコンテキストで、TPU VM はワーカーとも呼ばれます。
スライスに含まれる VM の数を決定するには、VM のタイプのセクションをご覧ください。
バイナリをインストールするかコードを実行するには、tpu-vm ssh コマンドを使用して TPU VM に接続します。
gcloud compute tpus tpu-vm ssh ${TPU_NAME}
SSH を使用して特定の TPU VM にアクセスするには、0 ベースのインデックスに続けて --worker フラグを使用します。
gcloud compute tpus tpu-vm ssh ${TPU_NAME} --worker=1
スライス形状が 8 チップを超える場合、1 つのスライスに複数の VM があります。この場合、--worker=all フラグを使用して、すべての TPU VM にインストールを実行します。個別に接続する必要はありません。次に例を示します。
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'
管理
TPU VM の管理に使用できるすべてのコマンドについては、TPU の管理をご覧ください。
フレームワークの設定
このセクションでは、TPU v5e で JAX または PyTorch を使用したカスタムモデル トレーニングの一般的な設定プロセスについて説明します。TensorFlow のサポートは、tpu-vm-tf-2.15.0-pjrt と tpu-vm-tf-2.15.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 を使用して v5e トレーニングを開始するには、このドキュメントの JAX チュートリアルをお試しください。
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 を使用してスクリプトを実行します。
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 を使用して v5e トレーニングを開始するには、このドキュメントの PyTorch チュートリアルをお試しください。
セッションの終了時に 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 状態にある、キューに格納されたリソース リクエストを削除するためにも使用できます。
モニタリングとプロファイル
Cloud TPU v5e では、前世代の Cloud TPU と同じ方法を使用したモニタリングとプロファイリングがサポートされています。プロファイリングの詳細については、Cloud TPU ツールでモデルをプロファイリングするをご覧ください。モニタリングの詳細については、Cloud TPU VM のモニタリングをご覧ください。
JAX/FLAX の例
v5e で ImageNet をトレーニングする
このチュートリアルでは、架空の入力データを使用して v5e で ImageNet をトレーニングする方法について説明します。実際のデータを使用する場合は、GitHub の README ファイルをご覧ください。
設定
環境変数を作成します。
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
-
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}キューに格納されたリソースが
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}QueuedResource が
ACTIVE状態の場合、出力は次のようになります。state: ACTIVE 最新バージョンの 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'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'架空のデータを生成するには、データセットのサイズに関する情報が必要です。これは、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 データセットを使用して、HugingFace から Vision Transformer（ViT）モデルをトレーニングする方法について説明します。
ViT モデルは、畳み込みネットワークと比較して優れた結果を持ち、ImageNet で Transformer エンコーダを正常にトレーニングした最初のモデルでした。詳しくは、次のリソースをご覧ください。
設定
環境変数を作成します。
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
-
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}キューに格納されたリソースが
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}キューに格納されたリソースが
ACTIVE状態の場合、出力は次のようになります。state: ACTIVE 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'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.16.1 && pip install -r examples/flax/vision/requirements.txt'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 |
| エポック | 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 モデルです。詳しくは、次のリソースをご覧ください。
設定
環境変数を作成します。
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
-
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}キューに格納されたリソースが
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}キューに格納されたリソースが
ACTIVE状態の場合、出力は次のようになります。state: ACTIVE 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'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.16.1 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 は、ヒューマン ラベリングなしで、未加工のテキストで事前にトレーニングされた変形モデルです。センテンス内の次の単語を予測するようにトレーニングされています。 詳しくは、次のリソースをご覧ください。
設定
環境変数を作成します。
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
-
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}キューに格納されたリソースが
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}キューに格納されたリソースが
ACTIVE状態の場合、出力は次のようになります。state: ACTIVE 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'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'構成をダウンロードしてモデルをトレーニングします。
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 | |
| エポック | 3 | 3 | 3 |
| グローバル バッチサイズ | 64 | 64 | 64 |
| スループット（例/秒） | 74.60 | 72.97 | 72.62 |
PyTorch/XLA
PJRT ランタイムを使用して ResNet をトレーニングする
PyTorch/XLA は、PyTorch 2.0 以降、XRT から PjRt に移行しています。PyTorch/XLA トレーニング ワークロード用に v5e を設定する手順が更新されています。
設定
環境変数を作成します。
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
-
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}QueuedResource が
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}キューに格納されたリソースが
ACTIVE状態の場合、出力は次のようになります。state: ACTIVE 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 を実行する方法を説明します。
設定
環境変数を作成します。
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
-
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}QueuedResource が
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}キューに格納されたリソースが
ACTIVE状態の場合、出力は次のようになります。state: ACTIVE torch/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'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 '事前トレーニング済みモデルの構成ファイルをダウンロードします。
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 | |
| エポック | 3 | 3 | 3 |
| 構成 | 600M | 2B | 16B |
| グローバル バッチサイズ | 64 | 128 | 256 |
| スループット（例/秒） | 66 | 77 | 31 |
v5e で ViT をトレーニングする
このチュートリアルでは、cifar10 データセット上の PyTorch/XLA で HuggingFace のリポジトリを使用して、v5e で VIT を実行する方法について説明します。
設定
環境変数を作成します。
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
-
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}QueuedResource が
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}キューに格納されたリソースが
ACTIVE状態の場合、出力は次のようになります。state: ACTIVE 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'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 | |
| エポック | 3 | 3 | 3 |
| グローバル バッチサイズ | 32 | 128 | 512 |
| スループット（例/秒） | 201 | 657 | 2,844 |
TensorFlow 2.x
単一ホスト v5e での Resnet のトレーニング
このチュートリアルでは、架空のデータセットを使用して v5litepod-4 または v5litepod-8 で ImageNet をトレーニングする方法について説明します。別のデータセットを使用する場合は、データセットの準備をご覧ください。
設定
環境変数を作成します。
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のいずれかになります。 /-
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}キューに格納されたリソースが
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。 キューに格納されたリソースの状態を確認するには、次のコマンドを使用します。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} SSH を使用して TPU に接続します
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}いくつかの環境変数を設定します
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モデル リポジトリのディレクトリに移動し、要件をインストールします。
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 とキューに格納されたリソースを削除する
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
マルチホスト v5e での Resnet のトレーニング
このチュートリアルでは、架空のデータセットを使用して v5litepod-16 以上での ImageNet をトレーニングする方法について説明します。別のデータセットを使用する場合は、データセットの準備をご覧ください。
環境変数を作成します。
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以上に設定できます。-
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}キューに格納されたリソースが
ACTIVE状態になると、TPU VM に SSH 接続できるようになります。 キューに格納されたリソースの状態を確認するには、次のコマンドを使用します。gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} SSH を使用して TPU（ワーカーゼロ）に接続します
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}いくつかの環境変数を設定します
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モデル リポジトリのディレクトリに移動し、要件をインストールします。
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 とキューに格納されたリソースを削除する
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