TPU v5e を使用してモデルをトレーニングする

Pod あたり 256 チップのフットプリントが小さい TPU 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 を使用したカスタムモデルのトレーニングの一般的な設定プロセスについて説明します。

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

環境変数をいくつか定義します。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

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 プロジェクトを設定するの説明に従って新しいプロジェクトを作成します。
ZONE	サポートされているゾーンについては、TPU のリージョンとゾーンのドキュメントをご覧ください。
worker	基盤となる 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
      pip install mkl mkl-include
      pip install tf-nightly tb-nightly tbp-nightly
      pip install numpy
      sudo apt-get install libopenblas-dev -y
      pip install torch~=PYTORCH_VERSION torchvision torch_xla[tpu]~=PYTORCH_VERSION -f https://storage.googleapis.com/libtpu-releases/index.html -f https://storage.googleapis.com/libtpu-wheels/index.html'

PYTORCH_VERSION は、使用する PyTorch のバージョンに置き換えます。PYTORCH_VERSION は、PyTorch/XLA に同じバージョンを指定するために使用されます。2.6.0 が推奨です。

PyTorch と PyTorch/XLA のバージョンの詳細については、PyTorch - スタートガイドと PyTorch/XLA リリースをご覧ください。

PyTorch/XLA のインストールの詳細については、PyTorch/XLA のインストールをご覧ください。

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
      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