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 は、C/C++ ランタイム関数 malloc と比較してパフォーマンスが大幅に向上します。TPU VM で使用されるデフォルトの malloc は tcmalloc です。LD_PRELOAD 環境変数を設定解除すると、TPU VM ランタイムで標準の malloc を強制的に使用できます。

   (vm)$ unset LD_PRELOAD

PyTorch バージョンの変更

TPU VM にプリインストールされた PyTorch バージョンを使用しない場合は、使用するバージョンをインストールします。次に例を示します。

(vm)$ sudo bash /var/scripts/docker-login.sh
(vm)$ sudo docker rm libtpu || true
(vm)$ sudo docker create --name libtpu gcr.io/cloud-tpu-v2-images/libtpu:pytorch-1.9 "/bin/bash"
(vm)$ sudo docker cp libtpu:libtpu.so /lib
(vm)$ sudo pip3 uninstall --yes torch torch_xla torchvision
(vm)$ sudo pip3 install torch==1.9.0
(vm)$ sudo pip3 install torchvision==0.10.0
(vm)$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.9-cp38-cp38-linux_x86_64.whl

簡単な計算を実行する

TPU VM で Python インタープリタを起動します。
```
(vm)$ python3
```
次の PyTorch パッケージをインポートします。
```
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)

次の出力が表示されます。

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 alpha compute tpus list --zone=${ZONE} を実行して、リソースが削除されたことを確認します。

高度な設定

上の例（単純な計算と ResNet50）では、PyTorch/XLA プログラムは、Python インタープリタと同じプロセスでローカル XRT サーバーを起動します。XRT ローカルサービスを別のプロセスで開始することもできます。

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

このアプローチの利点は、コンパイルのキャッシュがトレーニングの実行後も保持されることです。XLA サーバーを別のプロセスで実行する場合、サーバー側のロギング情報は /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
```

次の環境変数をエクスポートします。

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

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

ssh を任意の TPU ワーカー（ワーカー 0 など）に接続します。データ/モデルのダウンロードが完了したかどうかを確認します。ダウンロードは通常非常に高速ですが、ファイル処理に数分かかることがあります。ファイル処理の完了を確認するには、ダウンロードディレクトリで ls -al を実行します。ファイルは、起動スクリプトで指定した順序でダウンロードされます。次に例を示します。
```
download file1
download file2
download file3
```
file3 が起動スクリプトの最後のファイルであり、ls コマンドにより表示された場合、ファイル処理は完了しています。

ファイル処理が完了したら、Pod の VM ワーカー間で ssh に ssh-keys を生成します。その後、トレーニングを開始します。

注: 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 alpha compute tpus list --zone=${ZONE} を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。

リモートコーディネーターを含む Pod

TPU のメンテナンスイベントの発生時に Pod で自動的に復元できるように、リモートコーディネーターを使用することをおすすめします。リモートコーディネーターは、TPU VM ではなく、標準の Compute Engine VM です。このコマンドは、Pod 内の TPU VM にコマンドを発行します。