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

クイックスタート: PyTorch を使用して Cloud TPU VM で計算を実行する

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

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

始める前に

重要: このクイックスタートは、TPU VM または TPU ノード構成で使用できます。この 2 つの VM アーキテクチャについては、システム アーキテクチャをご覧ください。使用する gcloud コマンドは、使用する TPU 構成によって異なります。このクイックスタートでは、それぞれの gcloud コマンドをタブ付きセクションに示します。使用する TPU 構成のタブを選択すると、対応する gcloud コマンドがウェブページに表示されます。TPU ノードを使用する必要がある場合を除き、TPU VM を使用することをおすすめします。

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

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

gcloud コマンドを使用して Compute Engine VM と Cloud TPU を起動します。使用するコマンドは、TPU VM を使用するか TPU ノードを使用するかによって異なります。2 つの VM アーキテクチャの詳細については、システム アーキテクチャをご覧ください。gcloud コマンドの詳細については、gcloud リファレンスをご覧ください。

TPU VM

注: TPU VM の起動時に指定する TPU ソフトウェア バージョンを確認するには、Cloud TPU ソフトウェア バージョンをご覧ください。

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

$ gcloud compute tpus tpu-vm create tpu-name \
--zone=us-central1-b \
--accelerator-type=v3-8 \
--version=tpu-vm-pt-1.11

コマンドフラグの説明

zone

Cloud TPU を作成するゾーン。

accelerator-type

作成する Cloud TPU のタイプ。

version

Cloud TPU ソフトウェアのバージョン。

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

TPU ノード

1. PyTorch 用の TPU ノードを作成する場合、まず Compute Engine VM インスタンスを作成します。

   gcloud compute instances create tpu-name \
   --zone=us-central1-b  \
   --machine-type=n1-standard-16  \
   --image-family=torch-xla \
   --image-project=ml-images  \
   --boot-disk-size=200GB \
   --scopes=https://www.googleapis.com/auth/cloud-platform

   コマンドフラグの説明

   project

   実際の Cloud プロジェクトの ID

   name

   作成する Cloud TPU の名前。

   zone

   Cloud TPU を作成するゾーン。

   disk-size

   gcloud コマンドで作成された VM のハードディスクのサイズ（GB）。

   machine-type

   作成する Compute Engine VM のマシンタイプ。

   tf-version

   Tensorflow gcloud のバージョンが VM にインストールされます。

   accelerator-type

   作成する Cloud TPU のタイプ。

2. 次に、TPU インスタンスを作成します。

   gcloud compute tpus create tpu-name \
   --zone=us-central1-b \
   --network=default \
   --version=pytorch-1.11 \
   --accelerator-type=v3-8

   コマンドフラグの説明

   name

   作成する Cloud TPU の名前。

   zone

   Cloud TPU を作成するゾーン。

   network

   TPU が接続されているネットワーク。

   version

   PyTorch gcloud のバージョンが VM にインストールされます。

   accelerator-type

   作成する Cloud TPU のタイプ。

注: 複数のプロジェクトがある場合は、--project フラグでプロジェクト ID を指定する必要があります。

Cloud TPU VM に接続する

TPU VM

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --zone us-central1-b

TPU ノード

gcloud compute ssh tpu-name --zone=us-central1-b

XRT TPU デバイス構成を設定する

TPU VM

Torch-XLA 環境を構成します。

(vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

TPU ノード

1. TPU ノードの IP アドレスを調べます。

   (vm)$ gcloud compute tpus describe tpu-name --zone=us-central1-b

2. Torch-XLA 環境を構成します。your-tpu-ip-address は、TPU の IP アドレスに置き換えてください。

   (vm)$ conda activate torch-xla-1.11
   (vm)$ export TPU_IP_ADDRESS=your-tpu-ip-address
   (vm)$ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"

簡単な計算を実行する

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

   注: TPU VM で実行している場合は、次のエラー メッセージが表示されます。これは既知の問題であるため、無視してかまいません。この問題は、今後の TPU ソフトウェア リリースで修正される予定です。

   OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
   OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
   

クリーンアップ

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

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

   (vm)$ exit
   

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

2. Cloud TPU を削除します。

   TPU VM

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

   TPU ノード

   $ gcloud compute tpus execution-groups delete tpu-name \
   --zone=us-central1-b

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

次のステップ

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

• PyTorch コードを TPU Pod スライスで実行する
• Cloud TPU ユーザーガイド
• Cloud TPU システム アーキテクチャ
• PyTorch / XLA のドキュメント