PyTorch XLA ワークロードをプロファイリングする

プロファイリングにより、モデルのパフォーマンスを分析して改善できます。プロファイリングには他にも多くの要素がありますが、プロファイリングはタイミング オペレーションとして、またデバイス（TPU）とホスト（CPU）の両方で実行されるコードの一部と考えると役に立つ場合があります。このガイドでは、トレーニングや推論用にコードをプロファイリングする方法の概要を簡単に説明します。生成されたプロファイルを分析する方法の詳細については、次のガイドをご覧ください。

使ってみる

TPU を作成する

  1. 環境変数をエクスポートする

    $ export TPU_NAME=your_tpu_name
$ export ZONE=us-central2-b
$ export PROJECT_ID=project-id
$ export ACCELERATOR_TYPE=v4-8
$ export RUNTIME_VERSION=tpu-vm-v4-pt-2.0

    変数の説明をエクスポートする

    TPU name
    Cloud TPU に使用する名前。
    zone
    Cloud TPU を作成するゾーン。
    project ID
    モデルのトレーニングとプロファイリングに使用するプロジェクト ID。
    accelerator-type
    アクセラレータ タイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータ タイプの詳細については、TPU のバージョンをご覧ください。
    version
    Cloud TPU ランタイムのバージョン。デフォルトはエクスポートされた変数に表示されますが、サポートされている構成のリストから使用することもできます。

  2. TPU リソースを起動する

    $ gcloud compute tpus tpu-vm create ${TPU_NAME} \
--zone us-central2-b \
--accelerator-type ${ACCELERATOR_TYPE} \
--version ${RUNTIME_VERSION} \
--project $PROJECT_ID \
--subnetwork=tpusubnet

  3. gcloud scp コマンドを使用して、TPU VM のホーム ディレクトリにコードを移動します。次に例を示します。

    $ gcloud compute tpus tpu-vm scp my-code-file ${TPU_NAME}: --zone ${ZONE}

プロファイリング

プロファイルは capture_profile.py を使用して手動でキャプチャすることも、torch_xla.debug.profiler API を使用してトレーニング スクリプト内からプログラマティックにキャプチャすることもできます。

プロファイル サーバーを起動する

プロファイルをキャプチャするには、トレーニング スクリプト内でプロファイル サーバーが実行されている必要があります。任意のポート番号（次のコマンドにある 9012 など）でサーバーを起動します。

import torch_xla.debug.profiler as xp
server = xp.start_server(9012)

サーバーは main 関数の開始と同時に起動できます。

これで、次のセクションで説明するようにプロファイルをキャプチャできるようになりました。このスクリプトは、1 つの TPU デバイスで発生するすべての状況をプロファイリングします。

トレースの追加

ホストマシンでオペレーションをプロファイリングする場合は、コードに xp.StepTrace または xp.Trace を追加します。これらの関数は、ホストマシン上の Python コードをトレースします。（これは「グラフ」を TPU デバイスに渡すまで、ホスト（CPU）で Python コードを実行するのにかかる時間の測定値と考えることができます。そのため、これはトレースのオーバーヘッドを分析する際に非常に役立ちます）。これは、コードがデータのバッチを処理するトレーニング ループ内に追加できます。たとえば、

for step, batch in enumerate(train_dataloader):
    with xp.StepTrace('Training_step', step_num=step):
       ...

または、コードの個々の部分を以下のものでラップします

 with xp.Trace('loss'):
    loss = ...

Lighting を使用している場合は、コードの一部で自動的に行われるため、トレースの追加をスキップできます。ただし、さらにトレースを追加する場合は、トレーニング ループ内に挿入できます。

最初のコンパイル後にデバイスのアクティビティをキャプチャできます。モデルがトレーニングまたは推論のステップを開始するまで待機してください。

手動キャプチャ

Pytorch XLA リポジトリの capture_profile.py スクリプトを使用すると、プロファイルを迅速にキャプチャできます。これを行うには、キャプチャ プロファイル ファイルを TPU VM に直接コピーします。次のコマンドは、このファイルをホーム ディレクトリにコピーします。

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="wget https://raw.githubusercontent.com/pytorch/xla/master/scripts/capture_profile.py"

トレーニングの実行中に、次のコマンドを実行してプロファイルをキャプチャします。

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="python3 capture_profile.py --service_addr "localhost:9012" --logdir ~/profiles/ --duration_ms 2000"

このコマンドは .xplane.pb ファイルを logdir に保存します。ロギング ディレクトリ ~/profiles/ を任意の場所と名前に変更できます。Cloud Storage バケットに直接保存することもできます。そのためには、logdir を gs://your_bucket_name/ に設定します。

プログラマティック キャプチャ

スクリプトをトリガーしてプロファイルを手動でキャプチャするのではなく、トレーニング スクリプト内で torch_xla.debug.profiler.trace_attached API を使用してプロファイルを自動的にトリガーするようにトレーニング スクリプトを構成できます。

たとえば、特定のエポックとステップでプロファイルを自動的にキャプチャするには、PROFILE_STEP、PROFILE_EPOCH、PROFILE_LOGDIR 環境変数を使用するようにトレーニング スクリプトを構成できます。

import os
import torch_xla.debug.profiler as xp

# Within the training script, read the step and epoch to profile from the
# environment.
profile_step = int(os.environ.get('PROFILE_STEP', -1))
profile_epoch = int(os.environ.get('PROFILE_EPOCH', -1))
...

for epoch in range(num_epoch):
   ...
   for step, data in enumerate(epoch_dataloader):
      if epoch == profile_epoch and step == profile_step:
         profile_logdir = os.environ['PROFILE_LOGDIR']
         # Use trace_detached to capture the profile from a background thread
         xp.trace_detached('localhost:9012', profile_logdir)
      ...

これにより、PROFILE_LOGDIR 環境変数で指定されたディレクトリに .xplane.pb ファイルが保存されます。

TensorBoard で分析する

プロファイルをさらに分析するには、同じまたは別のマシンの TPU TensorBoard プラグイン
で TensorBoard を使用します（推奨）。

リモートマシンで TensorBoard を実行するには、SSH を使用して TensorBoard に接続し、ポート転送を有効にします。たとえば

$ ssh -L 6006:localhost:6006 remote server address

or

$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --ssh-flag="-4 -L 6006:localhost:6006"

リモートマシンで、必要なパッケージをインストールし、TensorBoard を起動します（マシンの ~/profiles/ にプロファイルがある場合）。プロファイルを別のディレクトリまたは Cloud Storage バケットに保存した場合は、必ずパスを正しく指定してください（例: gs://your_bucket_name/profiles）。

(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile
 
(vm)$ tensorboard --logdir ~/profiles/ --port 6006
 
(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly

TensorBoard の実行

ローカル ブラウザで、http://localhost:6006/ に移動し、プルダウン メニューから [PROFILE] を選択してプロファイルを読み込みます。

TensorBoard ツールとその出力の解釈方法については、TPU ツールをご覧ください。

画像

画像