Cloud TPU PyTorch/XLA ユーザーガイド

PyTorch/XLA を使用して ML ワークロードを実行する

このガイドの手順を開始する前に、GCP プロジェクトを準備するの説明に従って TPU VM と ssh を設定します。これにより、このガイドのコマンドを実行するために必要なリソースが設定されます。

PyTorch 1.8.1 と PyTorch/XLA 1.8.1 は、TPU VM にプリインストールされています。

基本設定

XRT TPU デバイス構成を設定します。

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

サイズ設定に頻繁な割り当てがあるモデルの場合、tcmalloc はデフォルトの malloc と比較してパフォーマンスが大幅に向上するため、TPU VM で使用されるデフォルトの malloctcmalloc になります。しかし、ワークロードによっては(たとえば、埋め込みテーブルへの割り当てが非常に大きいDLRMの場合)、tcmalloc では速度が低下する可能性があります。その場合は、次の変数の設定を解除して、代わりにデフォルトの malloc を使用してください。

   (vm)$ unset LD_PRELOAD

簡単な計算を実行する

Python 3 インタプリタを起動する

(vm)$ python3
   Python 3.6.9 (default, Jan 26 2021, 15:33:00)
   [GCC 8.4.0] on linux
   Type "help", "copyright", "credits" or "license" for more information.

Pytorch 3 インタープリタから、次の PyTorch パッケージをインポートします。

import torch
import torch_xla.core.xla_model as xm

次の計算を行います。

dev = xm.xla_device()
このコマンドの出力は次のようになります。
2021-04-01 23:20:23.268115: E   55362 tensorflow/core/framework/op_kernel.cc:1693] OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
2021-04-01 23:20:23.269345: E   55362 tensorflow/core/framework/op_kernel.cc:1693] OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

コマンドから次の出力が表示されます。

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

単一デバイス TPU での Resnet の実行

この時点で、任意の PyTorch / XLA コードを実行できます。たとえば、架空のデータを使用して ResNet を実行するには、次のコマンドを使用します。

   (vm)$ git clone --recursive https://github.com/pytorch/xla.git
   (vm)$ python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

ResNet サンプルでは 1 エポックのトレーニングを行い、所要時間は約 7 分です。この場合、次のような出力が返されます。

Epoch 1 test end 20:57:52, Accuracy=100.00 Max Accuracy: 100.00%

ResNet トレーニングが終了したら、TPU VM を削除します。

   (vm)$ exit
   $ gcloud alpha compute tpus tpu-vm delete tpu-name \
   --zone=zone

gcloud list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。以下のようなレスポンスは、インスタンスが正常に削除されたことを示します。

   $ gcloud alpha compute tpus list --zone=zone
   

   Listed 0 items.
   

高度な設定

上記の例(単純な計算と ResNet50)では、PyTorch/XL プログラムが、同じプロセスでローカル XRT サーバーを起動します。XRT ローカル サービスは別のプロセスで開始することもできます。

(vm)$ python3 -m torch_xla.core.xrt_run_server --port 51011 --restart

この方法の利点は、コンパイル キャッシュがトレーニング全体を通じて保持されることです。この場合、/tmp/xrt_server_log の下にサーバー側のロギングがあります。

(vm)$ ls /tmp/xrt_server_log/
server_20210401-031010.log

TPU VM パフォーマンス プロファイリング

TPU VM でのモデルのプロファイリングの詳細については、PyTorch XLA パフォーマンス プロファイリングをご覧ください。

ポッド

PyTorch/XLA では、すべての TPU VM がモデルコードとデータにアクセスできる必要があります。これを行う簡単な方法の 1 つは、TPU VM Pod の作成時に次の起動スクリプトを使用することです。すべての TPU VM でデータのダウンロードを行います。

  1. Cloud Shell で次のコマンドを実行して、gcloud の最新バージョンを実行していることを確認します。

    $ sudo /opt/google-cloud-sdk/bin/gcloud components update
    
  2. Pod TPU VM の作成に必要な次の環境変数をエクスポートします。

    $ export PROJECT_ID=project-id
    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export RUNTIME_VERSION=v2-alpha
    
  3. TPU VM を作成する

    $ gcloud alpha compute tpus tpu-vm create ${TPU_NAME} \
    --zone ${ZONE} --project ${PROJECT_ID} --accelerator-type v3-32 \
    --version ${RUNTIME_VERSION} --metadata startup-script='#! /bin/bash
    cd /usr/share/
    git clone --recursive https://github.com/pytorch/pytorch
    cd pytorch/
    git clone --recursive https://github.com/pytorch/xla.git
    EOF'
    

    これらの手順を続行する場合は、VM セッション ウィンドウで、(vm)$ で始まる各コマンドを実行します。

  4. ssh を任意の TPU ワーカー(ワーカー 0 など)に接続して、データモデルのダウンロードが完了したかどうかを確認します(この処理には時間がかかります。また、Pod 上の VM ワーカー間で ssh への SSH 認証鍵を生成した後にトレーニングが開始されます)。

    (vm)$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone ${ZONE} --project ${PROJECT_ID}
    

    プロジェクトの ssh メタデータを更新します。

    (vm)$ gcloud compute config-ssh
    
    (vm)$ export TPU_NAME=tpu-name
    
    (vm)$ python3 -m torch_xla.distributed.xla_dist \
       --tpu=${TPU_NAME} -- python3 /usr/share/pytorch/xla/test/test_train_mp_imagenet.py \
       --fake_data --model=resnet50 --num_epochs=1
    

    トレーニングには 3 分ほどかかります。完了すると、次のようなメッセージが表示されます。

    Epoch 1 test end 23:49:15, Accuracy=100.00
    10.164.0.11 [0] Max Accuracy: 100.00%
    

    各ワーカーの /tmp/xrt_server_log の下にサーバー側のロギングがあります。

       (vm)$ ls /tmp/xrt_server_log/
    
    server_20210401-031010.log
    

    XRT_SERVER を再起動する場合(サーバーが異常な状態になった場合)は、xla_dist を実行するときに --restart-tpuvm-pod-server を渡します。LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 のような環境変数などの新しい XRT サーバーの設定は、サーバーを再起動するまで取得されません。

  5. ResNet トレーニングが終了したら、TPU VM を削除します。

    (vm)$ exit
    
    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
    --zone=zone
    
  6. gcloud list を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。以下のようなレスポンスは、インスタンスが正常に削除されたことを示します。

    $ gcloud alpha compute tpus list --zone=zone
    
    Listed 0 items.
    

リモート コーディネーターを備えた Pod

リモート コーディネーターを設定して、Pod トレーニングを開始することをおすすめします。リモート コーディネーターは、TPU VM ではなく標準の Compute Engine VM です。リモート コーディネーターが TPU VM Pod にコマンドを発行します。リモート コーディネーターを使用する利点は、TPU メンテナンス イベントが発生するとトレーニングが停止し、TPU VM Pod にアクセスできなくなる場合がありますが、リモート コーディネーターにより、TPU VM を自動で復元できることです。

  1. TPU VM の作成に必要な次の環境変数をエクスポートします。

    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export PROJECT_ID=project-id
    
  2. TPU VM Pod スライスの作成

    TPU Pod スライスを作成するには、gcloud コマンドを使用します。たとえば、v2-32 Pod スライスを作成するには、次のコマンドを使用します。

    $ gcloud alpha compute tpus tpu-vm create tpu-name \
    --zone europe-west4-a --project tpu-prod-env-one-vm --accelerator-type v3-32 \
    --version v2-alpha --metadata startup-script='#! /bin/bash
    cd /usr/share/
    git clone --recursive https://github.com/pytorch/pytorch
    cd pytorch/
    git clone --recursive https://github.com/pytorch/xla.git
    EOF'
    
  3. リモート コーディネーター VM の作成に必要な次の環境変数をエクスポートします。

    $ export VM_NAME=vm-name
    $ export ZONE=zone
    
  4. 次のコマンドを実行してリモート コーディネーター VM を作成します。

    (vm)$ gcloud compute --project=project-id instances create vm-name \
      --zone=zone  \
      --machine-type=n1-standard-1  \
      --image-family=torch-xla \
      --image-project=ml-images  \
      --boot-disk-size=200GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    

    gcloud compute コマンドの実行が終了したら、shell プロンプトが username@projectname から username@vm-name に変更されたことを確認します。この変更により、リモート コーディネーター VM にログインしていることになります。

  5. ssh をリモート コーディネーター インスタンスに接続します。

    (vm)$ gcloud compute ssh vm-name --zone=zone
    
  6. torch-xla-1.8.1 環境を有効にして、そこからトレーニングを実行します。

    (vm)$ gcloud compute config-ssh
    
    (vm)$ conda activate torch-xla-1.8.1
    (vm)$ python3 -m torch_xla.distributed.xla_dist --tpu=tpu-name --restart-tpuvm-pod --env LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 -- python3 /usr/share/pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
    

    このトレーニングの実行には約 3 分かかり、次のようなメッセージが生成されます。

      Epoch 1 test end 23:19:53, Accuracy=100.00
      Max Accuracy: 100.00%
      

  7. ResNet トレーニングが終了したら、TPU VM を終了し、リモート コーディネーター VM と TPU VM を削除します。

    (vm)$ exit
    
    $ gcloud compute instances delete vm-name  \
      --zone=zone
    
    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
      --zone zone