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

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

このガイドの手順を開始する前に、Google Cloud プロジェクトの準備の説明に従って、TPU VM と ssh を設定します。

Pytorch / XLA 用に利用可能な TPU ソフトウェアには、次の 2 つのバージョンがあります。

最新のイメージ: tpu-vm-pt-1.11
tpu-vm-pt-1.10
v2-alpha

tpu-vm-pt-1.11 イメージは、PyTorch 1.11.0 と Pytorch 1.11.0 でプリインストールされています。tpu-vm-pt-1.10 イメージは、PyTorch 1.10.0 と Pytorch 1.10.0 でプリインストールされています。v2-alpha イメージは PyTorch 1.8.1 でプリインストールされています。

基本設定

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

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

サイズが大きく頻繁な割り当てがあるモデルの場合、tcmalloc は、C/C++ ランタイム関数 malloc と比較してパフォーマンスが向上します。TPU VM で使用されるデフォルトの malloc は tcmalloc です。TPU VM ソフトウェアで標準の malloc を使用するには、LD_PRELOAD 環境変数の設定を解除します。

   (vm)$ unset LD_PRELOAD

PyTorch バージョンの変更

TPU VM にプリインストールされた PyTorch バージョンを使用しない場合は、使用するバージョンをインストールします。たとえば、1.9 を使用する場合は、次のコマンドを実行します。

(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 がモデルのコードとデータにアクセスできる必要があります。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=tpu-vm-pt-1.11

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 にコマンドを発行します。