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

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

このガイドの手順を開始する前に、GCP プロジェクトの準備の説明に沿って、TPU VMssh を設定します。これにより、このガイドのコマンドの実行に必要なリソースが設定されます。

PyTorch 1.8.1 と PyTorch/XLA 1.8.1 は、TPU VM にプリインストールされています。

基本設定

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

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

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

   (vm)$ unset LD_PRELOAD

PyTorch バージョンの変更

TPU VM にプリインストールされた PyTorch バージョンを使用しない場合は、使用するバージョンをインストールします。次に例を示します。

(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 がモデルのコードとデータにアクセスできる必要があります。これを行う 1 つの方法は、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=v2-alpha
    
  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
    

    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 サーバーの設定は、サーバーを再起動するまで有効になりません。

  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 v2-alpha --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.8.1 環境を有効にして、その環境からトレーニングを実行します。

    (vm)$ gcloud compute config-ssh
    
    (vm)$ conda activate torch-xla-1.8.1
    (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