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

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 を作成する

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

TPU VM

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

コマンドフラグの説明

zone
Cloud TPU を作成するゾーン
accelerator-type
作成する Cloud TPU のタイプ
version
Cloud TPU ソフトウェアのバージョン

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
    GCP プロジェクト 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 のタイプ

Cloud TPU VM に接続する

TPU VM

$ gcloud alpha 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')
    
    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 alpha 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 の詳細を確認する。