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 で使用されるデフォルトの malloc は tcmalloc になります。しかし、ワークロードによっては（たとえば、埋め込みテーブルへの割り当てが非常に大きい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 でデータのダウンロードを行います。

Cloud Shell で次のコマンドを実行して、gcloud の最新バージョンを実行していることを確認します。
```
$ sudo /opt/google-cloud-sdk/bin/gcloud components update
```

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

$ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=zone
$ export RUNTIME_VERSION=v2-alpha

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)$ で始まる各コマンドを実行します。

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

注: TPU Pod では、最初の TPU VM はデフォルトで ssh になります。TPU Pod に関連付けられている他の TPU VM に ssh するには、ssh コマンドに --worker ${WORKER_NUMBER} を追加します。WORKER_NUMBER は 0 から始まるインデックスです。詳細については、TPU VM ユーザーガイドをご覧ください。
```
(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 サーバーの設定は、サーバーを再起動するまで取得されません。

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

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

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

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

$ export TPU_NAME=tpu-name
$ export ZONE=zone
$ export PROJECT_ID=project-id

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'

リモートコーディネーター VM の作成に必要な次の環境変数をエクスポートします。
```
$ export VM_NAME=vm-name
$ export ZONE=zone
```
次のコマンドを実行してリモートコーディネーター 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 にログインしていることになります。
ssh をリモートコーディネーターインスタンスに接続します。
```
(vm)$ gcloud compute ssh vm-name --zone=zone
```

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%

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