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 仮想環境を設定する必要があります。設定後、次のコマンドを実行します。