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

重要: TPU は、2 つの異なるアーキテクチャ、TPU ノードと TPU VM で使用できます。このチュートリアルでは、TPU VM を使用していることを前提としています。この 2 つの VM アーキテクチャについては、システム アーキテクチャをご覧ください。TPU ノードを使用する必要がある場合を除き、TPU VM を使用することをおすすめします。TPU ノードを使用した TPU Pod での PyTorch の使用の詳細については、Cloud TPU Pod での PyTorch モデルのトレーニングをご覧ください。

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

このガイドでは、PyTorch を使用して v4 TPU で簡単な計算を行う方法について説明します。

基本設定

1. Pytorch 2.0 の TPU VM ランタイムを実行している v4 TPU を使用して TPU VM を作成します。

     gcloud compute tpus tpu-vm create your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8 \
     --version=tpu-vm-v4-pt-2.0
   

   注: 別のバージョンの PyTorch をインストールする場合は、使用する TPU ランタイム バージョンに tpu-vm-v4-pt-2.0 を置き換えます。PyTorch を使用した TPU ランタイムの一覧については、PyTorch でサポートされるバージョンをご覧ください。

2. SSH を使用して TPU VM に接続します。

     gcloud compute tpus tpu-vm ssh your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8
   

   重要なポイント: このガイド全体を通して、接頭辞 (vm) $ は TPU VM インスタンスでコマンドを実行する必要があることを意味します。

3. PJRT または XRT TPU デバイス構成を設定します。

   注: PyTorch/XLA ランタイム オプションには、PJRT と XRT の 2 つがあります。XRT を使用する理由がない限り、PJRT を使用することをおすすめします。さまざまなランタイム構成の詳細と、XRT を使用する理由については、PJRT ランタイムのドキュメントをご覧ください。

   PJRT

       (vm)$ export PJRT_DEVICE=TPU
    

   XRT

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

4. Cloud TPU v4 でトレーニングする場合は、次の環境変数も設定します。

     (vm)$ export TPU_NUM_DEVICES=4
   

簡単な計算を実行する

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)
   

   注: OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op エラー メッセージは今後のリリースで修正される予定です。

   次の出力が表示されます。

   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} を実行して、リソースが削除されたことを確認します。

高度なトピック

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

   (vm)$ unset LD_PRELOAD

上記の例（簡単な計算と ResNet50）では、PyTorch/XLA プログラムが Python インタープリタと同じプロセスでローカル XRT サーバーを起動します。XRT ローカル サービスを別のプロセスで開始することもできます。

(vm)$ python3 -m torch_xla.core.xrt_run_server --port 51011 --restart

注: 上記で指定したポートは、先ほど export XRT_TPU_CONFIG="localservice;0;localhost:51011" コマンドで指定したポートと同じです。

このアプローチの利点は、トレーニングの実行中、コンパイル キャッシュが保持されることです。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 を実行する方法について説明します。

注: 現在、TPU VM で Docker を実行することは、単一デバイスの TPU（v2-8、v3-8）でのみサポートされています。

利用可能な Docker イメージ

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

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

注: 最後のコマンド出力（root:/#）は、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:/#

注: TPU デバイスがコンテナ内に表示されるようにするには、--privileged フラグを使用してコンテナを実行する必要があります。このコンテナには、pytorch conda 環境がプリインストールされています。(pytorch) で始まるすべてのコマンドをコンテナ内で実行する必要があります。

libtpu を確認する

注:libtpu は新しい TPUVM Docker イメージにプリインストールする必要がありますが、常に手動で確認することをおすすめします。

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