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

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

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

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

  1. 最新のイメージ: tpu-vm-pt-1.11
  2. tpu-vm-pt-1.10
  3. 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 で使用されるデフォルトの malloctcmalloc です。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

簡単な計算を実行する

  1. TPU VM で Python インタープリタを起動します。

    (vm)$ python3
    
  2. 次の PyTorch パッケージをインポートします。

    import torch
    import torch_xla.core.xla_model as xm
    
  3. 次のスクリプトを入力します。

    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 にダウンロードされます。

  1. Cloud Shell で次のコマンドを実行して、gcloud の最新バージョンを実行していることを確認します。

    $ sudo /opt/google-cloud-sdk/bin/gcloud components update
    
  2. 次の環境変数をエクスポートします。

    $ export PROJECT_ID=project-id
    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export RUNTIME_VERSION=tpu-vm-pt-1.11
    
  3. 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)$ で始まる各コマンドを実行します。

  4. ssh を任意の TPU ワーカー(ワーカー 0 など)に接続します。データ/モデルのダウンロードが完了したかどうかを確認します。ダウンロードは高速ですが、ファイル処理に数分かかることがあります。ファイル処理の完了を確認するには、ダウンロード ディレクトリで ls -al を実行します。ファイルは、起動スクリプトで指定した順序でダウンロードされます。次に例を示します。

    download file1
    download file2
    download file3
    

    file3ls コマンドの出力に含まれている場合、ファイル処理は完了しています。

    ファイル処理が完了したら、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 サーバーの設定は、サーバーを再起動するまで有効になりません。

  5. ResNet トレーニングが終了したら、TPU VM を削除します。

    (vm)$ exit
    
    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
     --zone=zone
    
  6. gcloud alpha compute tpus list --zone=${ZONE} を実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。

リモート コーディネーターを使用した Pod

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

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

    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export PROJECT_ID=project-id
    
  2. 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'
    
  3. リモート コーディネーター VM の作成に必要な次の環境変数をエクスポートします。

    $ export VM_NAME=vm-name
    $ export ZONE=zone
    
  4. 次のコマンドを実行して、リモート コーディネーター 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 にログインしていることを示しています。

  5. ssh をリモート コーディネーター インスタンスに接続します。

    (vm)$ gcloud compute ssh vm-name --zone=zone
    
  6. 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%
    

  7. 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