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

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

このガイドの手順を開始する前に、Cloud TPU VM ユーザーガイドの説明に従って、TPU VM と ssh を設定します。

Pytorch/XLA で使用できる TPU ソフトウェアバージョンのリストについては、PyTorch でサポートされているバージョンをご覧ください。

基本設定

XRT TPU デバイス構成を設定する

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

   (vm)$ export TPU_NUM_DEVICES=4

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

   (vm)$ unset LD_PRELOAD

PyTorch バージョンの変更

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

(tpuvm):$ cd /usr/share/
(tpuvm):$ sudo git clone -b release/1.13 --recursive https://github.com/pytorch/pytorch
(tpuvm):$ cd pytorch/
(tpuvm):$ sudo git clone -b r1.13 --recursive https://github.com/pytorch/xla.git
(tpuvm):$ cd xla/
(tpuvm):$ yes | sudo pip3 uninstall torch_xla
(tpuvm):$ yes | sudo pip3 uninstall torch
(tpuvm):$ yes | sudo pip3 uninstall torch_vision
(tpuvm):$ sudo pip3 install torch==1.13.0
(tpuvm):$ sudo pip3 install torchvision==0.14.0
(tpuvm):$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.13-cp38-cp38-linux_x86_64.whl
(tpuvm):$ sudo rm -rf /usr/local/lib/python3.8/dist-packages/libtpu*
(tpuvm):$ sudo pip3 install torch_xla[tpuvm]

簡単な計算を実行する

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 compute tpus tpu-vm delete tpu-name \
--zone=zone

削除には数分かかることがあります。gcloud 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 Pod の例

TPU VM Pod 上で PyTorch/XLA を実行するための設定情報と例については、PyTorch TPU VM Pod をご覧ください。

TPU VM 上の Docker

このセクションでは、PyTorch/XLA がプリインストールされた TPU VM 上で Docker を実行する方法について説明します。

利用可能な Docker イメージ

使用可能なすべての TPU VM Docker イメージを確認するには、GitHub の README をご覧ください。

TPU VM で Docker イメージを実行する

(tpuvm): sudo docker pull gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker run --privileged  --shm-size 16G --name tpuvm_docker -it -d  gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker exec --privileged -it tpuvm_docker /bin/bash
(pytorch) root:/#

libtpu を確認する

libtpu がインストールされていることを確認するには、次のコマンドを実行します。

(pytorch) root:/# ls /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/ | grep libtpu

その結果、次のような出力が生成されます。

libtpu
libtpu_nightly-0.1.dev20220518.dist-info

結果が表示されない場合は、次のコマンドを使用して、対応する libtpu を手動でインストールできます。

(pytorch) root:/# pip install torch_xla[tpuvm]

`tcmalloc` の所有権を証明してください

tcmalloc は、TPU VM で使用するデフォルトの Malloc です。詳細については、このセクションをご覧ください。このライブラリは、新しい TPU VM Docker イメージにプリインストールする必要がありますが、手動で確認することをおすすめします。次のコマンドを実行して、ライブラリがインストールされていることを確認できます。

(pytorch) root:/# echo $LD_PRELOAD