PyTorch を使用して Cloud TPU VM で計算を実行する

このクイックスタートでは、Cloud TPU の作成、PyTorch のインストール、Cloud TPU で簡単な計算を行う方法を説明します。Cloud TPU でモデルをトレーニングする方法の詳細については、Cloud TPU PyTorch のチュートリアルをご覧ください。

準備

このクイックスタートを使用する前に、Google Cloud Platform アカウントを作成し、Google Cloud CLI をインストールして gcloud コマンドを構成する必要があります。詳細しくは、アカウントと Cloud TPU プロジェクトを設定するをご覧ください。

gcloud を使用して Cloud TPU を作成する

デフォルトのユーザー プロジェクトで TPU VM を作成するには、ネットワークと zone を実行します。

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-a \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

コマンドフラグの説明

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

TPU を作成するときに、デフォルトのネットワークとサブネットワークを指定する場合は、追加の --network フラグと --subnetwork フラグを渡します。 デフォルト ネットワークを使用しない場合は、--network フラグを渡す必要があります。--subnetwork フラグは省略可能で、使用するネットワーク（デフォルトまたはユーザー指定）のデフォルトのサブネットワークを指定するために使用します。これらのフラグの詳細については、gcloud API リファレンス ページをご覧ください。

Cloud TPU VM に接続する

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

TPU VM に PyTorch/XLA をインストールする

   (vm)$ pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html

TPU ランタイム構成を設定する

PyTorch/XLA ランタイムが TPU を使用していることを確認します。

   (vm) $ export PJRT_DEVICE=TPU

簡単な計算を実行する

  1. 現在のディレクトリに tpu-test.py という名前のファイルを作成し、次のスクリプトをコピーして貼り付けます。

    import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

  2. スクリプトを実行します。

      (vm)$ python3 tpu-test.py

    スクリプトからの出力には、計算の結果が表示されます。

    tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

クリーンアップ

このページで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、次の手順を行います。

  1. Compute Engine インスタンスとの接続を切断していない場合は切断します。

    (vm)$ exit

    プロンプトが username@projectname に変わります。これは、現在、Cloud Shell 内にいることを示しています。

  2. Cloud TPU を削除します。

    $ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central1-a

このコマンドの出力で、TPU が削除されたことがわかります。

次のステップ

Cloud TPU VM の詳細を確認する。