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 環境を設定するをご覧ください。
Cloud TPU の作成
queued-resource create
コマンドを使用して、Cloud TPU v5es をキューに格納されたリソースとしてプロビジョニングすることをおすすめします。詳細については、キューに入れられたリソースを管理するをご覧ください。
Create Node API(gcloud compute tpus tpu-vm create
)を使用して Cloud TPU v5es をプロビジョニングすることもできます。詳細については、TPU リソースの管理をご覧ください。
トレーニングに使用できる v5e 構成の詳細については、トレーニング用の Cloud TPU v5e タイプをご覧ください。
フレームワークの設定
このセクションでは、TPU v5e で JAX または PyTorch を使用したカスタムモデルのトレーニングの一般的な設定プロセスについて説明します。TensorFlow のサポートは、tpu-vm-tf-2.18.0-pjrt
と tpu-vm-tf-2.18.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 -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
コマンドフラグの説明
変数 | 説明 |
TPU_NAME | キューに入れられたリソース リクエストの割り当て時に作成される TPU のユーザー割り当てテキスト ID。 |
PROJECT_ID | Google Cloud プロジェクト名。既存のプロジェクトを使用するか、Google Cloud プロジェクトを設定するで新しいプロジェクトを作成します。 |
ゾーン | サポートされているゾーンについては、TPU のリージョンとゾーンのドキュメントをご覧ください。 |
ワーカー | 基盤となる TPU にアクセスできる TPU VM。 |
次のコマンドを実行して、デバイスの数を確認できます(ここに表示されている出力は、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 でスクリプトを実行する
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 ファイルをご覧ください。
設定
環境変数を作成します。
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 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 -U "jax[tpu]" -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 データセットを使用して、HuggingFace から 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 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 -U "jax[tpu]" -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.18.0 && 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 |
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 モデルです。詳しくは、次のリソースをご覧ください。
設定
モデル出力用のストレージ バケットを設定します。
gcloud storage buckets create gs://your_bucket
--project=your_project
--location=us-west1
export GCS_BUCKET_NAME=your_bucket環境変数を作成する
export GCS_BUCKET_NAME=your_bucket export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west1-c 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
コマンドフラグの説明
変数 説明 GCS_BUCKET_NAME Google Cloud コンソール -> Cloud Storage -> バケットに表示されます。 PROJECT_ID Google Cloud プロジェクト名。既存のプロジェクトを使用するか、Google Cloud プロジェクトを設定するで新しいプロジェクトを作成します。 ACCELERATOR_TYPE TPU のバージョンについては、TPU のバージョンをご覧ください。 ZONE サポートされているゾーンについては、TPU のリージョンとゾーンのドキュメントをご覧ください。 RUNTIME_VERSION RUNTIME_VERSION に v2-alpha-tpuv5 を使用します。 SERVICE_ACCOUNT これは、Google Cloud コンソール -> IAM] -> サービス アカウント で確認できるサービス アカウントのアドレスです。 例: tpu-service-account@myprojectID。iam.gserviceaccount.com TPU_NAME キューに格納されたリソース リクエストの割り当て時に作成される TPU のユーザー割り当て ID。 QUEUED_RESOURCE_ID キューに格納されたリソース リクエストのユーザー割り当て ID。キューに格納されたリソースについては、キューに格納されたリソースのドキュメントをご覧ください。 QUOTA_TYPE reserved
、spot
のいずれかを設定できます。どちらも指定されていない場合、QUOTA_TYPE はデフォルトでon-demand
になります。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。VALID_UNTIL_DURATION リクエストが有効である期間。さまざまな有効期間の詳細については、キューに入れられたリソースをご覧ください。 -
gcloud 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.18.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
モデルのトレーニング
事前にマッピングされたバッファ(4GB)を使用してモデルをトレーニングします。
gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --project ${PROJECT_ID} --worker=all --command="
git clone https://github.com/google/maxdiffusion
cd maxdiffusion
git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
pip3 install -r requirements.txt
pip3 install .
export LIBTPU_INIT_ARGS=""
python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"
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 モデルです。文の次の単語を予測するようにトレーニングされました。詳しくは、次のリソースをご覧ください。
設定
環境変数を作成します。
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 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 -U "jax[tpu]" -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 && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'
モデルのトレーニング
事前にマッピングされたバッファ(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 を設定する手順が更新されています。
設定
環境変数を作成します。
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 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 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
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'
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=' gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive gcloud storage 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 |
config | 6 億 | 20 億 | 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 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
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'
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 をトレーニングする方法について説明します。別のデータセットを使用する場合は、データセットの準備をご覧ください。
設定
環境変数を作成します。
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 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 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