Cloud TPU VM でモデルをプロファイリングする

プロファイリングは、Cloud TPU でモデルのパフォーマンスを最適化するための主なツールの一つです。コアとなるプロファイリングツールは XProf で、このツールは OpenXLA/XProf GitHub リポジトリから入手できます。XProf は、JAX、Pytorch XLA、Tensorflow / Keras など、すべての XLA ベースのフレームワークのプロファイリングをサポートしています。

プロファイルのキャプチャ

XProf を使用するには、まずモデルのプロファイルをキャプチャする必要があります。プロファイルをキャプチャするには、次の 2 つの方法があります。

プログラムによるキャプチャ
オンデマンドキャプチャ（手動キャプチャ）

プログラムによるキャプチャでは、コード内でプロファイルをキャプチャする場所を指定するため、モデルのコードにアノテーションを付ける必要があります。通常、いくつかのトレーニングステップでプロファイルをキャプチャするか、モデル内の特定のコードブロックをプロファイリングします。JAX、Pytorch XLA、TensorFlow などのフレームワークでは、トレースのキャプチャ方法がいくつかあります（トレースの開始または停止を API ベースで行うか、コンテキストマネージャーベースで行うかなど）。MaxText などの上位レベルのフレームワークでは、MaxText の実行時に profiler=xplane フラグを有効にするだけで、プロファイルのキャプチャを有効にできます。

プロファイルのオンデマンドキャプチャは、アドホックな方法でプロファイルをキャプチャする場合や、プログラムによるキャプチャを有効にしていない場合に利用できます。これは、実行中にモデル指標で問題が発生し、問題を診断するために、その瞬間のプロファイルを一定期間キャプチャする場合に役立ちます。

オンデマンドキャプチャを有効にするには、コード内で XProf サーバーを起動する必要があります。profiler.start_server を有効にすると、ML ワークロードで XProf サーバーが起動します。このサーバーがオンデマンドキャプチャトリガーをリッスンして、プロファイルのキャプチャを開始します。オンデマンドキャプチャは、TensorBoard UI または XProfiler ツールを使用した CLI でトリガーできます。

異なるフレームワークでプログラムによるキャプチャとオンデマンドキャプチャを有効にする方法については、以下をご覧ください。

プログラムによるキャプチャとオンデマンドキャプチャのいずれの方法でも、キャプチャしたプロファイルの保存先を指定する必要があります。プロファイルは、/profiles/run-name のようなパスを使用して TPU VM のディレクトリに保存することも、gs://bucket-name/run-name/ のようなパスを使用して Cloud Storage に保存することもできます。

このようなパスを使用すると、プロファイルは次の場所に保存されます。

/profiles/run-name/plugins/profile/session1/

または

gs://bucket-name/run-name/plugins/profile/session1/

セッションは、モデルの実行中に行った 1 回のキャプチャに対応しています。たとえば、トレーニング実行でステップ 1～3 のプロファイルをキャプチャし、その後、ステップ 8～10 のプロファイルをキャプチャした場合、これらは同じ実行のプロファイルになりますが、最初のステップ 1～3 のキャプチャは session1 になり、2 回目のステップ 8～10 のキャプチャは session2 になります。

実行ごとに、異なる日付スタンプの付いたセッションが示されます。異なるセッションにプログラム、オンデマンド、またはその両方の組み合わせでプロファイルをキャプチャできます。これは、モデルトレーニングのさまざまな部分のプロファイルを比較する場合に便利です。たとえば、トレーニング開始時のプロファイルとトレーニング終了時のプロファイルを比較する場合などです。

TensorBoard でプロファイルを可視化する

TensorBoard で XProf プロファイルを表示するには、tensorboard-plugin-profile プラグインをインストールする必要があります。

キャプチャしたプロファイルを Google Cloudで可視化するには、TPU VM で TensorBoard を使用するか、XProfiler ツールを使用して VM または GKE Pod に TensorBoard をホストすることをおすすめします。

TPU VM でプロファイルをキャプチャした場合は、その TPU VM の TensorBoard でキャプチャしたプロファイルを表示できます。

TPU VM に XProf と TensorBoard をインストールするには、次のコマンドを実行します。

pip install tensorboard_plugin_profile tensorboard
pip install xprof

TensorBoard でキャプチャしたプロファイルを表示するには、次のコマンドを実行します。

$ tensorboard --logdir=profiles/run-name

または

$ tensorboard --logdir=profiles

複数の実行から複数のプロファイルを読み込むには、特定の実行のプロファイルを含むサブディレクトリではなく、すべての実行とプロファイルを含むルートディレクトリを TensorBoard に指定します。

Google Cloudで Tensorboard プロファイルを可視化するための Cloud-Diagnostics-XProf ライブラリ

cloud-diagnostics-xprof（XProfiler ライブラリ）を使用すると、TensorBoard をホストし、 Google Cloudでプロファイルを可視化することが容易になります。詳細については、cloud-diagnostics-xprof GitHub リポジトリをご覧ください。

cloud-diagnostics-xprof リポジトリでは、XProf と TensorBoard をローカルで実行する場合と比べて、次のような改善が図られています。

XProf と TensorBoard の依存関係の設定とパッケージ化
プロファイルを Cloud Storage に保存します。これは、長期的な保持や、実行後の分析に役立ちます（実行を終了すると、ローカルでキャプチャされたプロファイルは削除されます）。
Compute Engine VM または GKE Pod に TensorBoard をプロビジョニングすることで、大規模なプロファイルや複数のプロファイルを高速に読み込みます。読み込み速度や費用に関するニーズに応じてマシンタイプを変更することもできます。
プロファイルの共有、チームメンバーや Google エンジニアとのコラボレーションのためにリンクを作成します。
GKE と Compute Engine のワークロードのオンデマンドプロファイリングが容易になり、ワークロードを実行しているホストを選択してプロファイルをキャプチャできます。

XProfiler ライブラリを使用する前に、プログラムでワークロードコードのプロファイルをキャプチャするか、プロファイルサーバーを起動して、後でオンデマンドでキャプチャできるようにする必要があります。

XProfiler ライブラリを設定するには、Google Cloud CLI と Python 仮想環境を設定する必要があります。設定後、次のコマンドを実行します。

pip install cloud-diagnostics-xprof

これにより、必要な XProf と TensorBoard の依存関係がすべてインストールされます。

次に、TensorBoard をホストする VM または GKE Pod を作成するため、次のコマンドを実行します。

xprofiler create -z $ZONE -l $GCS_PATH

または

xprofiler create --GKE -z $ZONE -l $GCS_PATH

$ZONE は任意のゾーンに置き換え、$GCS_PATH はプロファイルトレースのパスに置き換えます。複数の実行のすべてのプロファイルトレースを含むルートディレクトリを指定するか、1 回の実行の特定のプロファイルトレースのセットを指定します。

たとえば、次のディレクトリでプロファイルをキャプチャするとします。

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

ルートディレクトリ（gs://bucket-name）に GCS パスを設定できます。

xprofiler create -z $ZONE -l gs://bucket-name/

TensorBoard UI には、すべての実行とセッションのすべてのプロファイルが run1/session1、run1/session2、run2/session1 として表示されます。

デフォルトでは、xprofiler create は Compute Engine VM（具体的には c4-highmem-8 VM）を作成します。マシンタイプは -m フラグで変更できます。Compute Engine VM ではなく GKE Pod に TensorBoard インスタンスを作成する場合は、–GKE フラグを xprofiler create に渡すことができます。GKE Pod で TensorBoard インスタンスをホストすると、GKE にデプロイされた残りのワークロードと一緒に TensorBoard を管理しやすくなります。

TensorBoard をホストする VM または GKE Pod を使用すると、 Google Cloud での大きなプロファイルや複数のプロファイルの読み込みが、ローカルでホストされている TensorBoard よりもはるかに高速で行われます。Google のベンチマークによると、デフォルトの c4-highmem-8 VM を使用した最初の読み込みでは、数分以内に 1 GB 程度のプロファイルが読み込まれます。また、パフォーマンスと費用のニーズに応じて適切なマシンタイプを選択することもできます。

xprofiler create を実行すると、次のような出力が表示されます。

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

最初のオプションは、TensorBoard で XProf プロファイルを表示できるリンクです。このリンクを使用すると、所属するチームや、 Google Cloudでパフォーマンスの最適化を支援する Google エンジニアとプロファイルを共有できます。

リンクにアクセスできるユーザーは、プロファイルデータを保存する Cloud Storage バケットに設定された権限に基づいて制御されます。リンクが機能しない場合は、SSH を介して TensorBoard インスタンスに接続し、xprofiler connect コマンドを実行してプロファイルを表示できます。

ワークロードコードで Cloud Profiler サーバーを有効にして（これは、プログラムによるキャプチャとオンデマンドキャプチャの両方で必要です）、オンデマンドプロファイリングを実行する場合は、次の 2 つの方法があります。

a. TensorBoard の [Capture profile] ボタンをクリックします。これで、ワークロードが実行されているデバイスホストをプロファイリング対象として選択できるようになります。GKE または Compute Engine で実行されているワークロードのオンデマンドキャプチャがサポートされています。

b. XProfiler キャプチャコマンドを使用し、ゾーン、Cloud Storage バケット、フレームワーク、ホスト VM / Pod 名、キャプチャ期間（ミリ秒）などの情報を指定します。これは、TensorBoard UI に入力する必要のある情報と同じです。

cloud-diagnostics-xprof ライブラリの使用方法について詳しくは、こちらの GitHub ページをご覧ください。