PyTorch XLA のパフォーマンス プロファイリング

概要

このガイドでは、PyTorch を使用して Cloud TPU のパフォーマンス ツールと指標の自動分析機能を使用する方法について説明します。これらのツールにより、トレーニング ワークロードのパフォーマンスのデバッグと最適化が容易になります。

コンセプト

PyTorch / XLA を初めて使用する場合は、PyTorch の API_GUIDEトラブルシューティングのドキュメントをご覧ください。Cloud TPU については、コンセプト ドキュメントをご覧ください。

TPU ノード + PyTorch/XLA のプロファイリング

必要な Cloud TPU 関連のリソースを作成して初期化します。

  1. プロジェクト ID、Cloud Storage バケットと、TPU リソースに使用するゾーンの変数を作成します。

    export PROJECT_ID=PROJECT_ID
    export BUCKET_NAME=BUCKET_NAME
    export ZONE=ZONE
    
    gcloud --project=$PROJECT_ID compute project-info add-metadata \
    --metadata BUCKET_NAME=$BUCKET_NAME
    
  2. Compute Engine VM インスタンスを作成します。 このインスタンスには、すべての Python スクリプトとモデルが格納されます。

    gcloud compute instances create profiler-tutorial-vm \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --machine-type=n1-standard-16 \
      --image-project=ml-images \
      --image-family=torch-xla \
      --boot-disk-size=300GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    
  3. TPU リソースを作成します。

    gcloud compute tpus create profiler-tutorial-tpu \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --network=default \
     --version=pytorch-1.8 \
     --accelerator-type=v3-8
    
  4. Cloud Storage バケットを作成する。

    まだインストールしていない場合は、まず gsutil CLI をインストールします(インストール手順)。

  5. gsutil mb コマンドを使用して、プロファイリング アーティファクトがすべて保存される Cloud Storage バケットを作成します。リージョンとバケット名の変数は、トレーニングで使用する値に置き換えます。

    gsutil mb -p ${PROJECT_ID} -c standard -l REGION gs://${BUCKET_NAME}
    

    ここで

    • REGION は、Cloud TPU を作成したリージョンです(例: europe-west4)。
  6. Cloud TPU プロジェクトのサービス アカウントを作成します。

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    このコマンドでは、Cloud TPU サービス アカウントを次の形式で返します。

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

    例: service-164006649440@cloud-tpu.iam.gserviceaccount.com

  7. サービス アカウントをエクスポートして、ストレージ バケットに関するサービス アカウント権限を付与します。account-number は、サービス アカウント作成の出力で返された PROJECT_NUMBER に置き換えます。

    export SERVICE_ACCOUNT=service-ACCOUNT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    gsutil acl ch -u $SERVICE_ACCOUNT:READER gs://${BUCKET_NAME}
    gsutil acl ch -u $SERVICE_ACCOUNT:WRITER gs://${BUCKET_NAME}
    

TensorBoard を設定する

  1. ローカルマシンのポート 9001 に対する VM の VM 転送ポート 9001 に ssh 接続します。このポートは、ローカル ブラウザで TensorBoard UI を開くために使用されます。

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
  2. TensorBoard のインストール専用の Conda 環境を作成します。

    conda create -y -n tensorboard python=3.6
    conda activate tensorboard
    pip install tf-nightly==2.5.0.dev20210211 tb-nightly==2.5.0a20210211 tbp-nightly==2.5.0a20210211
    
  3. Compute Engine VM で TensorBoard サーバーを稼働させてインストールをテストし、次にローカルマシン上で http://localhost:9001/#profile にアクセスして、サーバーへの接続を試みます。

     # Get bucket name
     BUCKET_NAME=$(curl "http://metadata.google.internal/computeMetadata/v1/project/attributes/BUCKET_NAME" -H "Metadata-Flavor: Google")
    
     tensorboard --logdir gs://${BUCKET_NAME} --port 9001
    

ローカルマシンで http://localhost:9001/#profile にアクセスすると、次のように表示されます。

画像

http://localhost:9001 にアクセスし、上部の右隅の [アップロード] ボタンの横にあるプルダウン メニューから、[プロファイル] オプションを選択して、先ほどのプロフィール ページを表示することもできます。

画像

モデルのプロファイリング

TensorBoard サーバーを引き続き動作させるには、ローカルマシンから新しいターミナル ウィンドウを開始し、GCE VM にもう一度 ssh 接続を行います(今回のポート転送には -L を使用しません)。

  1. 新しいシェル内のため、新しいターミナル ウィンドウで、プロジェクト ID、ストレージ バケット環境、ゾーン変数を再度エクスポートします。

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. VM に ssh 接続します。

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
    conda activate torch-xla-1.8
    PROJECT_ID=$(curl "http://metadata.google.internal/computeMetadata/v1/project/project-id" -H "Metadata-Flavor: Google")
    export TPU_IP_ADDRESS=$(gcloud compute tpus describe profiler-tutorial-tpu --zone=${ZONE} --project=${PROJECT_ID} \
    --format="value(ipAddress)")
    echo TPU_IP_ADDRESS=${TPU_IP_ADDRESS}
    export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    
  3. 統合テストが環境内のエンドツーエンドで動作することを確認します。

    python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profiler.py # takes <1 min
    >
  4. トレーニングを開始する前に、/usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py 内の次の行を編集します。

    以下の行を変更します。

    accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    To:
    accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    前述の train_mnist の 2 つの引数により、動的グラフとテンソルの取得が人為的に行われます。これらについては、後の自動指標分析のセクションで説明します。ここでは、TPU をプロファイリングするだけのため、次の例は名目上のパフォーマンスで動作します。

  5. サーバー プロファイリングに使用されるトレーニング実行を開始します。

    PT_XLA_DEBUG=1 XLA_HLO_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU(サーバー)のプロファイリング

トレーニングが開始されたら、http://localhost:9001/#profile にアクセスして、次の手順でプロファイルを取得します。

画像

次のページは自動的に再読み込みされます。

画像

サポートされるツールが、左側のペインの [ツール] プルダウンに表示されます。

  • 概要ページ(このページには、TensorFlow / TPU と PyTorch / XLA TPU とでは根本的に異なっている、入力パイプライン分析ツールの表示は含まれていません)
  • メモリビューア
  • Op Profile
  • Pod ビューア
  • Tensorflow 統計情報(PyTorch の統計情報など、フレームワーク レベルの統計情報)
  • Trace ビューア(Chrome ブラウザが必要)

概要ページ

このページには、キャプチャされたプロフィールの概要が表示されます。この例では、MNIST データセットで小さなモデルをトレーニングしているため、非常に高いアイドル時間を示しています。

メモリビューア

テンソルと HLO PyTorch 演算ごとのデバイスメモリ(HBM)使用量が表示されます。メモリビューアは HLO モジュールごとにメモリのビューを取得するため、デバイスデータ割り当て(入力とラベルの一括送信)グラフなどのモジュールが存在する場合があります。特定の HLO モジュールのメモリ使用量を表示するには、左側の [Hosts] プルダウンからモジュールを選択します。

画像

選択した HLO モジュールを表示すると、モジュール実行 HBM フットプリントのタイムラインを包括的に確認できます。これは、割り当てサイズ、プログラムの実行サイズ、パディング サイズの順に表示されます。

画像

それぞれのバッファ割り当てにカーソルを合わせると、各バッファの詳細を確認できます。たとえば、デバイス HBM の大半を占める割り当ては、次のように見つけます。

画像

上の例では、(1)はユーザーコードによって追加された torch_xla.debug.profiler.Trace アノテーションに対応します。この行に対応する test_profile_mp_mnist.py コードを調査します。

   class MNIST(nn.Module):
   ...
     def forward(self, x):
       with xp.Trace('conv1'):
         x = F.relu(F.max_pool2d(self.conv1(x), 2))
         x = self.bn1(x)
   ...
   

また、test_mnist 演算名前空間から、この HLO モジュールが、xp.trace('test_mnist') コンテキスト マネージャーを持つ eval ループに対応していることがわかります。

XLA 演算プロファイル

演算プロファイルは、プロファイリング期間中に実行された XLA 演算のパフォーマンス統計を表示する Cloud TPU ツールです。演算プロファイルには次の情報が表示されます。

  • アプリケーションによる Cloud TPU の使用状況。演算に費やされた時間の割合がカテゴリ別に表示されます。また、TPU FLOPS の使用率も表示されます。
  • 最も時間のかかる演算。これらの演算は、最適化の潜在的なターゲットとなります。
  • 個々の演算の詳細。形状、パディング、演算を使用する式などが表示されます。

演算プロファイルを使用すると、最適化に適したターゲットを見つけることができます。たとえば、モデルで TPU のピーク FLOPS が 5% にすぎない場合、このツールを使用して、実行に最も時間を費やした XLA 演算とその演算による TPU FLOPS の使用量を特定できます。

画像

それぞれの説明は以下のとおりです。

  1. 概要セクション。Cloud TPU の使用率を表示し、最適化のための推奨事項が表示されます。
  2. コントロール パネル。表に表示される演算の数、表示される演算、並べ替え方法を設定するためのコントロールが含まれます。
  3. 演算テーブル。XLA 演算に関連付けられている TensorFlow 演算の上位カテゴリを一覧表示する表。これらの演算は、Cloud TPU 使用率の順に並んでいます。
  4. 演算詳細カードテーブル内の演算にカーソルを合わせると、演算に関する詳細情報が表示されます。FLOPS 使用率、演算が使用されている式、演算レイアウト(fit)などが表示されます。

Pod ビューア

Pod ビューアツールの詳細については、TPU ツールをご覧ください。

フレームワークの統計情報(TensorFlow / PyTorch の統計情報)

フレームワーク統計情報には、TPU デバイスと TPU ホストで稼働している PyTorch と XRT 演算の詳細な内訳が表示されます。

トレース ビューア

Trace Viewer は、Cloud TPU のパフォーマンス分析ツールです。トレース ビューアは Chrome トレース イベント プロファイリング ビューアを使用しているため、Chrome ブラウザを使用する必要があります。

トレース ビューアには次のタイムラインが表示されます。

  • TensorFlow モデルによって実行された演算の実行期間。
  • 演算を実行したシステムの部分(TPU またはホストマシン)。通常、PyTorch / XLA の場合、ホストマシンは主にコンパイルとバッファの割り当て / 割り当て解除で動作するのに対して、TPU は実際のモデル トレーニングを行います。
  • トレース ビューアを使用して、モデル内のパフォーマンスの問題を特定し、この問題を解決する対策を講じることができます。ドリルダウンして、実行に最も時間がかかっている PyTorch / XLA の演算を特定できます。

xp.Trace(NAME) アノテーションを追加することで、モデルの特定の部分の実行にかかる時間を測定するために、トレースを直接追加できます。たとえば、次のトレースは、以下のように表示されます。

画像

  1. test_profile_mp_mnist.py のモデルコード内で、明示的なユーザー アノテーションによって生成された。
  2. PyTorch 演算が実行された(事前低下)。
  3. PyTorch / XLA 自動生成 HLO モジュール名。
  4. XLA Ops がデバイスで実行された(融合)。

詳細については、トレース ビューアの一般的な TPU ドキュメントをご覧ください。ただし、入力パイプラインや他の TensorFlow 固有の部分に関するセクションは、このドキュメントの文脈と関連がないため、無視してください。

PyTorch / XLA クライアントのプロファイリング

モデルの実行中に TPU 側をプロファイリングする場合と同様に、次は PyTorch / XLA クライアント側をトレーニング中にプロファイリングします。クライアント側で使用される主なモニタリング ツールはトレース ビューアです。

プロファイリング サーバーは、トレーニング スクリプトで起動する必要があります。例については、こちらをご覧ください。これにより、TensorBoard からクエリしてトレースを取得できます。

画像

複数のプロセスからトレースを取得するため、各プロセスでは異なるポートでプロファイリング サーバーを起動(たとえば、ベースポート番号に xm.get_ordinal() を追加する)できます。つづいて、「,」で連結された localhost:port のリストを指定します。TensorBoard では、複数のプロセスからのトレースを一度に表示することができないため、プロセスごとに別の [Host] プルダウンを確認します。

次の図は、サンプルのトレースを示したものです。

画像

TPU トレースにさまざまな名前空間トレースを追加する方法と同様に、同じ API を使用して、名前空間トレースをクライアント側トレース(xp.Trace(NAME))に追加できます。このモデルは小さく、小さい NIST イメージで使用されているため、ステップ時間は短く、必ずしも均一ではありません。演習として、ResNet50 のトレースを取得するための test_train_mp_imagenet.py --fake_data の例と同様に、トレースを追加し、プロファイラ サーバーの起動を試すことができます。

トレースには、調査可能な別のメタデータがあります。たとえば、TransferToServer と TransferFromServer のトレースでは、送受信されたテンソルの正確な数と合計サイズが表示されます。

画像

XLA グラフのコンパイルでは、問題の診断に役立つグラフハッシュが表示されます。

画像

また、TensorBoard UI によるプロファイリングの代わりに、PyTorch / XLA から TPU とクライアントの両方をプログラムでプロファイリングする API(xp.trace())も提供されています。

自動指標分析

このセクションでは、デバッグモードを使用して、以下のようなパフォーマンスの問題を検出する方法を説明します。

  • 動的グラフ / 連続コンパイル
  • 非常に遅いグラフ コンパイル
  • 非常に遅いグラフの実行
  • 頻繁な XLA→CPU 転送
  • デバイス HBM からホスト RAM への度重なるスワップ
  • 度重なる HBM 断片化
  • 減少しない aten:: 演算

トレーニングを開始する前に、/usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py の次の行を元に戻します。

以下の行を変更します。

   accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
   
To:
   accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
   

これらの変更により、コンパイルとテンソルの取得が人為的に行われます。dynamic_graph=True は各ステップのバッチサイズを人為的に変更し、XLA の低下したグラフがステップごとに変わり、再コンパイルされます。fetch_often=True は、各ステップで loss.item() の呼び出しを挿入し、各ステップでデバイスからテンソルの値を取得するため、パフォーマンスが低下します。

サンプル トレーニング スクリプトを実行します。

   PT_XLA_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --fake_data --num_cores=1
   

デバッグ時には、デバッグ プロセスが簡素化されるため、--num_cores=1 を指定して実行することをおすすめします。出力例の一部は、次のようになります。

Epoch 1 train begin 01:18:05
| Training Device=xla:1/0 Step=0 Loss=0.00000 Rate=1905.00 GlobalRate=1904.72 Time=01:18:05
pt-xla-profiler: TransferFromServerTime too frequent: 3 counts during 3 steps
pt-xla-profiler: TransferFromServerTime too frequent: 4 counts during 4 steps
pt-xla-profiler: TransferFromServerTime too frequent: 5 counts during 5 steps
pt-xla-profiler: TransferFromServerTime too frequent: 6 counts during 6 steps
pt-xla-profiler: TransferFromServerTime too frequent: 7 counts during 7 steps
pt-xla-profiler: TransferFromServerTime too frequent: 8 counts during 8 steps
pt-xla-profiler: TransferFromServerTime too frequent: 9 counts during 9 steps
pt-xla-profiler: TransferFromServerTime too frequent: 10 counts during 10 steps
pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromServerTime too frequent: 11 counts during 11 steps
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps

接頭辞 pt-xla-profiler がある行は、自動指標分析の出力に対応しています。この例では、TransferFromServerTime の出現頻度が高すぎる(1 ステップに 1 回)ことがわかります。これは、トレーニング ループがステップごとに loss.item() の値を取得するためです。また、グラフ内にダイナミック シェイプがあることによってモデルを繰り返し再コンパイルする必要がある場合、CompileTime too frequent という警告が表示されることがあります。この種の問題を引き起こすコード スニペットを次に示します。test_profile_mp_mnist.py

    for step, (data, target) in enumerate(loader):
      if dynamic_graph:
        # The batch dimension is different every step.
        index = max(-step, -flags.batch_size + 1)  # non-empty
        data, target = data[:-index, :, :, :], target[:-index]
      ...
      if fetch_often:
        # Fetch tensor value from XLA:TPU to CPU every step.
        loss_i = loss.item()

次に、Ctrl+C でトレーニング スクリプトを終えると、減少していない演算の概要が最後に表示されます。aten::_local_scalar_dense は、XLA テンソルを CPU コンテキストに戻すことに対応している特殊な演算です。

このレポートでは、aten::_local_scalar_dense 演算が呼び出される主な場所が 2 か所、どちらも loss.item() のソースコードにあります。

  • test/test_profile_mp_mnist.py:158
  • test/test_profile_mp_mnist.py:61
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=27):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:158)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: FRAME (count=2):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   _train_update (test/test_profile_mp_mnist.py:61)
pt-xla-profiler:    (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:700)
pt-xla-profiler:   _run_step_closures (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:709)
pt-xla-profiler:   mark_step (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:723)
pt-xla-profiler:   __exit__ (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/debug/profiler.py:153)
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:162)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

次のスクリプトで自動指標分析を実行します。このスクリプトには、低下しない演算(_ctc_loss)が含まれています(1.8 リリース時点)。

PT_XLA_DEBUG=1 python <<EOF
import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t = torch.randn(50, 16, 20).log_softmax(2).to(dev)
target = torch.randint(low=1, high=20, size=(16, 30), dtype=torch.long).to(dev)
input_lengths = torch.full(size=(16,), fill_value=50, dtype=torch.long).to(dev)
target_lengths = torch.randint(low=10, high=30, size=(16,), dtype=torch.long).to(dev)

for _ in range(10):
  loss = torch.nn.CTCLoss()(t, target, input_lengths, target_lengths)
  xm.mark_step()
EOF

上のスクリプトを PT_XLA_DEBUG=1 で実行すると、出力は次のようになります。

…
pt-xla-profiler: TransferFromServerTime too frequent: 30 counts during 10 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=10):
pt-xla-profiler: Unlowered Op: "_ctc_loss"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   ctc_loss (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/functional.py:2305)
pt-xla-profiler:   forward (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/loss.py:1593)
pt-xla-profiler:   _call_impl (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/module.py:889)
pt-xla-profiler:    (:11)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

自動指標アナライザは、STDIN の 11 行目によって、この低下しない演算が発生していることを示しています(つまり、torch.nn.CTCLoss() を含む行)。現在、ctc_loss 演算は低下しておらず、その理由を上のレポートで確認します。実行前は、テンソルが最初は XLA:TPU にあるため、TransferFromServerTime の警告がいくつか表示されることもありますが、この演算は減少していないため、最初に XLA テンソルを CPU に戻して、CPU 上で aten:: 演算を実行し、それを元に戻す必要があります。

代わりに pt-xla-profiler の出力をファイルに書き込む場合は、PT_XLA_DEBUG=1PT_XLA_DEBUG_FILE=$PATH_TO_FILE を設定します。

クリーンアップ

VM を終了してから、次のコマンドを実行して TPU、VM、Cloud Storage バケットを削除します。

(vm)$ exit
gcloud compute instances delete profiler-tutorial-vm \
  --zone=${ZONE} \
  --project=${PROJECT_ID}
gcloud compute tpus delete profiler-tutorial-tpu \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --async
gsutil rm -fr gs://${BUCKET_NAME}

TPU VM + PyTorch / XLA のプロファイリング

このセクションでは、TPU VM アーキテクチャを使用して PyTorch / XLA をプロファイリングします。

環境変数のエクスポート

  1. TPU リソースに使用するプロジェクト ID とゾーンの変数を作成します。

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    

Cloud TPU の作成

TPU VM ユーザーガイドを参照し、設定後に torchtorch_xlatorchvisiontensorboard がプリインストールされた v3-8 TPU VM を作成してください。

  1. TPU リソースを作成します。

    gcloud alpha compute tpus tpu-vm create profiler-tutorial-tpu-vm \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --version=v2-alpha \
     --accelerator-type=v3-8
    

TensorBoard サーバーの起動

  1. VM に SSH 接続し、tensorboard-plugin-profile をインストールして、TensorBoard サーバーを起動します。

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
      pip3 install tf-nightly==2.5.0.dev20210211 tb-nightly==2.5.0a20210211 tbp-nightly==2.5.0a20210211
    
      tensorboard --logdir ./tensorboard --port 9001
    

ローカルマシンの http://localhost:9001 で TensorBoard の出力を確認すると、次のように表示されます。

画像

http://localhost:9001 に TensorBoard の出力を表示する場合は、右上の [UPLOAD] ボタンの横にあるプルダウンで [PROFILE] オプションを選択して、上記のプロファイル ページにアクセスすることも可能です。

画像

モデルのプロファイリング

開発環境の新しいターミナル ウィンドウで、上記と同じ環境変数をエクスポートして、TPU VM に ssh 接続します。

  1. 新しいターミナル ウィンドウで新しいシェルに入っているため、プロジェクト ID とゾーン変数をもう一度エクスポートします。

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. VM に ssh 接続します。

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
  3. PyTorch / XLA リポジトリのクローンを作成し、e2e テストを実行します。

      git clone -b r1.8 https://github.com/pytorch/xla
      export XRT_TPU_CONFIG="localservice;0;localhost:51011"
      python3 xla/test/test_profiler.py  # takes <1 min
    >
  4. トレーニングを開始する前に、xla/test/test_profile_mp_mnist.py の次の行を編集します。

    以下の行を変更します。

        accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    To:
        accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    前述の train_mnist の 2 つの引数により、動的グラフとテンソルの取得が人為的に行われます。これらについては、後の自動指標分析のセクションで説明します。ここでは、TPU をプロファイリングするだけのため、次の例は名目上のパフォーマンスで動作します。

  5. トレーニングを開始します。

     XLA_HLO_DEBUG=1 python3 xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU + クライアントのプロファイリング

トレーニングが開始されたら、http://localhost:9001 で TensorBoard の出力を表示し、次の手順でプロファイルを取得します。

画像

次のページが再読み込みされます。

画像

現在 TPU VM 設定では、トレース ビューアツールのみが選択されているため、[Tools] プルダウンで [trace_viewer] を選択し、トレースを調べます。TPU VM 設定では、「クライアント側」と TPU デバイス側の両方のトレースが 1 つのフルビューで表示されることを確認できます。

画像

クリーンアップ

  1. VM を終了してから、次のコマンドを実行して TPU、VM、Cloud Storage バケットを削除します。

    (vm)$ exit
    

作成した TPU VM を削除します。

  1. Cloud TPU と Compute Engine リソースを削除します。

    $ gcloud alpha compute tpus tpu-vm delete profiler-tutorial-tpu-vm \
      --project ${PROJECT_ID} --zone=${ZONE}
    
  2. 次のコマンドを実行して、リソースが削除されたことを確認します。削除には数分かかることがあります。以下のようなレスポンスは、インスタンスが正常に削除されたことを示します。

    $ gcloud alpha compute tpus tpu-vm list --project ${PROJECT_ID} --zone=${ZONE}
    
    Listed 0 items.