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
その結果、次のような出力が生成されます。/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4
LD_PRELOAD
が設定されていない場合は、手動で実行できます。
(pytorch) root:/# sudo apt-get install -y google-perftools (pytorch) root:/# export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4"
デバイスを確認する
次のコマンドを実行して、TPU VM デバイスが使用可能であることを確認します。
(pytorch) root:/# ls /dev | grep accel
次の結果が生成されます。accel0 accel1 accel2 accel3
結果が表示されない場合は、コンテナを --privileged
フラグで開始していない可能性があります。
モデルを実行する
次のコマンドを実行して、TPU VM デバイスが使用可能かどうかを確認できます。
(pytorch) root:/# export XRT_TPU_CONFIG="localservice;0;localhost:51011" (pytorch) root:/# python3 pytorch/xla/test/test_train_mp_imagenet.py --fake_data --num_epochs 1