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 を生成します。その後、トレーニングを開始します。(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 にコマンドを発行します。
次の環境変数をエクスポートします。
$ export TPU_NAME=tpu-name $ export ZONE=zone $ export PROJECT_ID=project-id
TPU VM 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 tpu-vm-pt-1.11 --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.11
環境を有効にして、その環境からトレーニングを実行します。(vm)$ gcloud compute config-ssh
(vm)$ conda activate torch-xla-1.11 (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