XProf を使用したマネージド プロファイリング

ML Diagnostics ワークロード モニタリング、ML Diagnostics CLI、または ML Diagnostics SDK を使用する場合は、マネージド XProf プロファイリングを使用できます。マネージド XProf は、ML 診断プラットフォームがGoogle Cloud プロジェクトでプロビジョニングして管理する、スケーラブルなバックエンドを備えた XProf インスタンスです。一方、セルフホスト型 XProf では、プロファイルを表示するために独自の XProf インスタンスを手動でデプロイして維持する必要があります。

XProf は、ML ワークロード用のオープンソースのプロファイリングとパフォーマンス分析ツールであり、OpenXLA エコシステムの一部です。セルフホスト型の XProf と比較して、マネージド XProf には次の機能があります。

  • セットアップ不要: XProf 可視化サーバーやその依存関係の手動インストールや構成は必要ありません。
  • 一元化されたマルチラン プロファイル セッション アクセス: プロジェクトごとに 1 つのマネージド XProf インスタンスを使用して、異なる ML 実行間でプロファイル セッションを読み込みます。
  • 自動スケーリングと高速読み込み: 複数のユーザーが同時にプロファイルを閲覧する場合に、大規模なプロファイルの読み込みを高速化し、バックエンド リソースを自動的にスケーリングします。
  • 組み込みのコラボレーションとセキュリティ: プロファイリング インスタンスをセキュリティの脆弱性から保護しながら、プロファイル セッションの共有可能なリンクを生成します。
  • 複数のプロファイル キャプチャ モード: マルチホスト キャプチャ、オンデマンド プロファイリング、スナップショット モードでの継続的プロファイリング、ワークロード モニタリングからの自動プロファイル キャプチャをサポートしています。
  • 追加のインフラストラクチャ費用なし: XProf インスタンスをセルフホストする場合とは異なり、マネージド XProf コンピューティング リソースを追加料金なしで提供します。

デプロイパス

マネージド XProf をデプロイして有効にするには、ML Diagnostics SDK、ワークロード モニタリング、または ML Diagnostics CLI を使用します。

ML Diagnostics プラットフォームは、プロジェクトごとにスケーラブルなバックエンドを備えたマネージド XProf インスタンスを 1 つデプロイします。このインスタンスを使用すると、プロジェクト内のすべての ML 実行でキャプチャされたすべてのプロファイル セッションを表示できます。プロファイルをキャプチャするには、キャプチャしたプロファイル データを保存するために、ML 実行に関連付けられた Cloud Storage バケットも必要です。

マネージド XProf を有効にして、Cloud Storage バケットを ML 実行に関連付けるには、次のいずれかの方法を使用します。

  • ML Diagnostics SDK: SDK を使用して最初の ML 実行をデプロイすると、ML Diagnostics プラットフォームはプロジェクトでマネージド XProf インスタンスを自動的に有効にしてデプロイします。プロファイルを保存する Cloud Storage バケットを指定するには、SDK で ML 実行を定義するときに gcs_path パラメータを設定します。
  • ワークロードのモニタリング: ワークロードのモニタリングでは、マネージド XProf はデフォルトで有効になりません。代わりに、自動プロファイル キャプチャがプロジェクトでプロファイル セッションを初めてトリガーすると、ML 診断プラットフォームはマネージド XProf インスタンスを自動的にプロビジョニングして有効にします。デフォルトでは、ML 診断はプロジェクトにリージョン Cloud Storage バケットもプロビジョニングします。
  • ML Diagnostics CLI: 任意の ML 実行を更新して XProf ツールを有効にできます。これにより、プロジェクトにマネージド XProf インスタンスがデプロイされ、Cloud Storage バケットが ML 実行に関連付けられます。

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

ML Diagnostics は、XProf プロファイルをキャプチャする 4 つの方法(プログラムによるキャプチャ、オンデマンド プロファイリング、スナップショット モードでの継続的なプロファイリング、ワークロード モニタリングからの自動プロファイル キャプチャ)をサポートしています。これらのメソッドは、キャプチャがトリガーされる方法と、記録される時間枠が異なります。

次の表は、プロファイル キャプチャ方法を比較したものです。

キャプチャ方法 説明 サポートされているフレームワーク
プログラムによるキャプチャ

プロファイリング コマンドをワークロード コードに直接埋め込み、指定されたトレーニング ステップまたは関数でプロファイル データの記録を開始および停止します。

既知のトレーニング ステップまたは特定のコード セクションの再現可能なプロファイルが必要な場合は、この方法を使用します。

JAX、vLLM、Tunix、TorchTPU
オンデマンド プロファイリング

ワークロードの実行中に、指定した期間にわたって Google Cloud コンソールまたは gcloud CLI からライブ プロファイルのキャプチャをトリガーします。

この方法は、ステップ範囲を変更したりワークロードを再起動したりせずに、実行中のワークロードをリアルタイムで検査する場合に使用します。

JAX、vLLM、Tunix、TorchTPU
スナップショット モードでの継続的プロファイリング(プレビュー)

TPU ホストの循環バッファに 2 GB のデバイス トレースのプロファイルを継続的に記録し、スナップショットをトリガーすると、最新の 2 GB のバッファを Cloud Storage バケットに保存します。

この方法は、断続的な速度低下や異常につながるワークロード アクティビティを検査する必要がある場合に使用します。

JAX、vLLM、Tunix
ワークロード モニタリングからのプロファイルの自動キャプチャ(プレビュー)

ワークロード モニタリングのパフォーマンス低下アナライザが "bad" 状態を報告すると、XProf プロファイルのキャプチャが自動的にトリガーされます。プロファイルの自動キャプチャを構成して、オンデマンド キャプチャをトリガーするか、継続的プロファイリングが有効になっている場合は継続的プロファイリング スナップショットをトリガーできます。

この方法は、パフォーマンスが低下したときに、手動でキャプチャをトリガーするのを待たずに、ML 診断でトレースデータをキャプチャする場合に使用します。

JAX、vLLM、Tunix、TorchTPU

プログラムによるキャプチャ

プログラムによるキャプチャでは、プロファイリング コマンドをワークロード コードに直接埋め込み、プロファイル データの記録を開始および停止するタイミングを明示的に指定します。ML 実行の gcs_path パラメータで指定された Cloud Storage バケットに書き込まれたプロファイルは、プロジェクトのマネージド XProf インスタンスにプロファイル セッションとして自動的に表示されます。これは、Google Cloud コンソールまたは Google Cloud CLI を使用して表示できます。

プログラムによるキャプチャは、次のいずれかの方法で構成できます。

オンデマンド プロファイリング

オンデマンド プロファイリングを使用すると、ワークロードの実行中にプロファイルのキャプチャをリアルタイムでトリガーできます。オンデマンド キャプチャをトリガーすると、指定した期間にわたってプロファイルが記録されます。オンデマンド プロファイリングは、ワークロード コードで直接有効にするか、ML Diagnostics SDK を使用して有効にできます。

前提条件

オンデマンド プロファイルをキャプチャする前に、次の設定手順を完了します。

  1. ワークロードのポート 9999 で XProf サーバーを起動して、オンデマンド キャプチャ トリガーに応答できるようにします。

    • JAX で直接: ポート 9999 で profiler.start_server メソッドを呼び出します。

      profiler.start_server(port=9999)
      
    • ML Diagnostics SDK: ML 実行を定義するときに、machinelearning_run で on_demand_xprof=True を設定します。

  2. クラスタ オーケストレーターまたはコンピューティング環境を構成します。

オンデマンド プロファイルをキャプチャする

オンデマンド プロファイル キャプチャは、Google Cloud コンソールまたは gcloud CLI からトリガーできます。キャプチャをトリガーするときに、次のパラメータを指定します。

  • Duration: キャプチャの長さ(ミリ秒または秒)。キャプチャできる最大時間は、ワークロードの特性と、2 GB のプロファイルの上限に収まるデータ量によって異なります。
  • ターゲット: プロファイルをキャプチャするターゲット ノード。
    • Google Cloud コンソール: ワークロードを実行しているすべてのノードが自動的に入力されます。すべてのノードまたはノードのサブセットを選択できます。
    • gcloud CLI: ワークロードのプロファイラ ターゲットを一覧表示し、それらのターゲットのすべてまたは一部のキャプチャをトリガーします。
  • デバイス、ホスト、Python のトレーサー レベル: デバイス、ホスト、Python のトレースでキャプチャする詳細レベル。

スナップショット モードでの継続的プロファイリング

継続的プロファイリングが有効になっている場合、ML Diagnostics プラットフォームは、TPU ホスト上の循環バッファに 2 GB のプロファイルを継続的に記録します。スナップショットをトリガーすると、ML 診断はバッファから最新の 2 GB のプロファイルをプロファイル セッションとして Cloud Storage バケットに保存します。このスナップショットを使用すると、問題につながるワークロード アクティビティを検査できます。継続的なプロファイリングはデバイス トレースのみをサポートします。ホスト トレースや Python トレースは収集されません。

前提条件

継続的プロファイリングを開始する前に、環境が次の要件を満たしていることを確認してください。

  • サポートされているライブラリのバージョン:
    • libtpu バージョン 0.0.43 以降。
    • jax バージョン 0.11.0 以降。
  • クラスタ オーケストレーターまたはコンピューティング環境の設定:

profiler.start_server(port=9999) を呼び出してワークロード コードで XProf サーバーを起動するか、ML Diagnostics SDK を使用する場合は machinelearning_run で on_demand_xprof=True を設定します。

継続的プロファイリングを開始

継続的プロファイリングは、ワークロード コードで直接開始できます。これには、gcloud CLI または Google Cloud コンソールを使用します。開始すると、継続的プロファイリングは、指定された TPU ホストの循環バッファに 2 GB のプロファイルを記録します。

gcloud CLI またはGoogle Cloud コンソールを使用して継続的なプロファイリングを開始する場合は、次のパラメータを指定します。

  • ターゲット: 継続的にプロファイリングするターゲット ノード。
    • Google Cloud コンソール: ワークロードを実行しているすべてのノードが自動的に入力されます。すべてのノードまたはノードのサブセットを選択できます。
    • gcloud CLI: ワークロードのプロファイラ ターゲットを一覧表示し、それらのターゲットのすべてまたは一部に対して継続的なプロファイリングを開始します。
  • デバイスのトレーサー レベル: 継続的にキャプチャするデバイス トレースのレベル。継続的なプロファイリング中は、ホストと Python のトレースが無効になります。

継続的プロファイリングは、Python ワークロード コードから直接、gcloud CLI、または Google Cloud コンソールで開始できます。

Python ワークロード コード

Python ワークロード コードから直接継続的プロファイリングを開始するには、start_continuous_profiling 関数を呼び出します。

from xprof.api import continuous_profiling_snapshot
continuous_profiling_snapshot.start_continuous_profiling('localhost:9999', {})

詳細については、XProf JAX プロファイリングをご覧ください。

gcloud CLI

gcloud CLI から継続的プロファイリングを開始するには、次のコマンドを使用して、アクティブなプロファイラ ターゲットを一覧表示し、それらのノード全体で継続的プロファイリングを開始します。

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RUN_ID="RUN_ID"

# 1. Discover active ProfilerTarget IDs on the ML run:
TARGETS=$(gcloud alpha mldiagnostics profiler-target list \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --format="csv[no-heading](name.basename())" | paste -sd, -)

# 2. Start continuous profiling across the target nodes:
gcloud alpha mldiagnostics machine-learning-run start-continuous-profiling \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --targets="${TARGETS}" \
  --device-tracer-level=enabled

次のように置き換えます。

  • PROJECT_ID: 実際の Google Cloud プロジェクト ID。
  • LOCATION: ML 実行が配置されているリージョン。
  • RUN_ID: ML 実行の ID。
Google Cloud コンソール

Google Cloud コンソールで継続的なプロファイリングを開始する手順は次のとおりです。

  1. Google Cloud コンソールで、[Cluster Director] または [GKE > AI/ML > 診断を実行] に移動し、ML 実行を選択します。
  2. [プロファイル] タブを開きます。
  3. 上部のアクションバーで、[継続的なプロファイリングを開始] をクリックします。
  4. プロファイリングするノードとデバイス トレーサー レベルを選択し、[Start] をクリックします。

継続的プロファイリングが開始されると、[継続的プロファイルのキャプチャ] バナーにセッションの開始時間が Since START_TIME として表示され、アクションバーに [スナップショットを取得] と [継続的プロファイリングを停止] が表示されます。

スナップショットをキャプチャする

継続的プロファイリングの実行中にスナップショットをキャプチャすると、循環バッファから最新の 2 GB のプロファイルをプロファイル セッションとして Cloud Storage バケットに保存できます。スナップショットをキャプチャするには、gcloud CLI または Google Cloud コンソールを使用します。

継続的プロファイリング スナップショットをキャプチャするときに、次のパラメータを指定します。

  • ターゲット: 継続的プロファイリングが実行されているターゲット ノードのすべてまたはサブセット。
    • Google Cloud コンソール: 継続的プロファイリング用に構成されたすべてのノードが自動的に入力されます。すべてのノードまたはノードのサブセットを選択できます。
    • gcloud CLI: 継続的プロファイリングが実行されているプロファイラ ターゲットのすべてまたはサブセットを指定します。
  • プロファイラ セッション名: キャプチャされたスナップショット セッションの名前。

スナップショットをキャプチャするには、gcloud CLI またはGoogle Cloud コンソールを使用します。

gcloud CLI

gcloud CLI を使用して継続的なプロファイリング スナップショットをキャプチャするには、次のコマンドを実行します。

PROJECT_ID="PROJECT_ID"
LOCATION="LOCATION"
RUN_ID="RUN_ID"

gcloud alpha mldiagnostics profiler-session capture "SNAPSHOT_NAME" \
  --machine-learning-run="${RUN_ID}" \
  --location="${LOCATION}" \
  --project="${PROJECT_ID}" \
  --targets="TARGET_NODES" \
  --kind=snapshot

次のように置き換えます。

  • PROJECT_ID: 実際の Google Cloud プロジェクト ID。
  • LOCATION: ML 実行が配置されているリージョン。
  • RUN_ID: ML 実行の ID。
  • SNAPSHOT_NAME: キャプチャされたプロファイル セッションの名前。
  • TARGET_NODES: キャプチャするターゲット ノードのカンマ区切りのリスト。
Google Cloud コンソール

Google Cloud コンソールで継続的なプロファイリング スナップショットをキャプチャする手順は次のとおりです。

  1. Google Cloud コンソールで、[Cluster Director] または [GKE] > [AI/ML] > [診断を実行] に移動し、ML 実行を選択して、[プロファイル] タブを開きます。
  2. 上部のアクションバーで、[スナップショットを作成] をクリックします。
  3. [継続的プロファイリング スナップショットをキャプチャ] ペインで、プロファイル名を入力し、ターゲット ノードを選択して、[送信] をクリックします。

スナップショットが [プロファイル] テーブルに表示されます。リンクがすぐに表示されない場合は、[更新] をクリックします。

ワークロード モニタリングからのプロファイルの自動キャプチャ

ワークロード モニタリングからの自動プロファイル キャプチャを使用すると、ML 診断は、パフォーマンス低下アナライザが "bad" 状態を報告したときに XProf プロファイルを自動的にキャプチャします。問題が発生した後にオンデマンド キャプチャを手動でトリガーする必要はありません。プロファイルの自動キャプチャを構成して、オンデマンド キャプチャをトリガーするか、継続的プロファイリングが有効になっている場合は継続的プロファイリング スナップショットをトリガーできます。このプロセスでは、イベントの発生時にキャプチャされたプロファイルを使用して、問題の根本原因を診断できます。

設定手順と前提条件については、ワークロード モニタリングのドキュメントのプロファイルの自動キャプチャをご覧ください。

エージェント型 AI のプロファイル分析

エージェント型 AI プロファイル分析は、TPU ワークロード用にキャプチャされた XProf パフォーマンス プロファイルを自動的に分析し、未加工の XProf トレースデータから HTML 分析レポートを生成します。Vertex AI で Gemini モデルを使用すると、分析エージェントが XProf プロファイリング ツールを実行し、パフォーマンスのボトルネックを要約して、コードまたは構成の変更を推奨します。必要な API が有効になると、ML 診断は、プラットフォームでキャプチャされたプロファイルごとにエージェント型 AI プロファイル分析レポートを生成します。

エージェント型 AI プロファイル分析では、オープンソースの XProf スキル リポジトリと同じ基盤となるツールと機能が使用されます。ML 診断では、スキルファイルを直接使用したり、オープンソース スキルに必要な XProf CLI をインストールしたりするのではなく、基盤となる XProf ツールの実装をインポートして、エージェントが直接呼び出せるようにします。その結果、ML Diagnostics の分析出力は、ローカルの xplane ファイルで XProf CLI を使用してオープンソースの XProf スキルを実行した場合と同等になります。

前提条件

エージェント型 AI プロファイル分析を使用する前に、 Google Cloud プロジェクトで次の設定手順を完了します。

必要な API の有効化

分析エージェントがプロジェクト内で Gemini モデルを実行できるように、Vertex AI API(aiplatform.googleapis.com)を有効にします。

  • Google Cloud コンソール: API ライブラリで Vertex AI API を有効にします。
  • gcloud CLI: 次のコマンドを実行します。

    gcloud services enable aiplatform.googleapis.com
    

IAM ロールを付与する

ユーザー アカウントに次のいずれかの Identity and Access Management(IAM)ロールを付与します。

  • Cluster Director 閲覧者(roles/hypercomputecluster.viewer): Cluster Director と ML 診断リソースへの読み取り専用アクセス権を付与します。これには、 Google Cloud コンソールの [診断を実行] > [プロファイル] で生成された分析レポートも含まれます。
  • カスタムロール(hypercomputecluster.locations.analysis): カスタムロールに追加すると、 Google Cloud コンソールの [診断を実行] > [プロファイル] で生成された分析レポートを表示する権限が付与されます。

仕組み

プログラムによるオンデマンド、継続的、または自動のプロファイル キャプチャが完了すると、分離されたバックグラウンド ジョブが XProf プロファイルで次の分析ワークフローを実行します。

  1. 広範なデータの収集: 15 個の XProf ツールを並行して実行し、ステップ時間、ハードウェア使用率、メモリ使用量、ルーフライン境界、最も遅いオペレーション、コンパイラの非効率性などの健全性スナップショットを収集します。このステップでは、トレーニング ステップごとにタイミングを正規化し、正常なオペレーションをボトルネックから分離します。
  2. Focused AI investigation: 収集されたヘルス スナップショットを確認します。ワークロードのコンピューティング使用率が高く、アイドル時間と通信オーバーヘッドが低く、メモリ不足のリスクが低い場合、エージェントは追加のツール呼び出しをスキップして OPTIMAL レポートを生成します。ボトルネックが検出されると、エージェントは仮説を立て、追加の XProf ツールを使用して最大 4 つのターゲットを絞ったフォローアップ チェックを実行し、計算グラフ内の周囲のオペレーション、ホスト間のタイムライン遅延、カスタム カーネルを検査します。
  3. レポートの生成: XProf の指標と AI エージェントの根本原因分析、推奨される修正を組み合わせます。生成された summary.html レポートは Cloud Storage バケットに保存され、Google Cloud コンソールの [Run Diagnostics] > [Profiles] に表示されます。

このワークフローでは、プロファイル データと分析がプロジェクト内で非公開に保たれます。未加工のプロファイラ トレースと生成された HTML レポートは Cloud Storage バケットに保存され、すべての AI モデル呼び出しは Google Cloud プロジェクトの Vertex AI を介して実行されます。

XProf ツール

次の表に、分析中にエージェントが実行する XProf ツールを示します。

カテゴリ XProf ツール 説明
ワークロードの概要とステップのタイミング get_xprof_overview、get_step_trace、get_kpi_metrics、get_device_information、get_hosts、get_utilization_viewer、get_smart_suggestions

これらのツールは、広範なデータの収集時に実行されます。get_hosts ツールは、AI に焦点を当てた調査時にも実行され、次の項目を確認します。

  • ハードウェア タイプ、ホストとチップの数、ハードウェアのピーク上限
  • 平均ステップ時間と、コンピューティング、ネットワーク通信、ホストまたはデータ読み込みのアイドル時間の内訳
  • 行列乗算ユニット(MXU)、TensorCore、メモリ帯域幅の使用率、組み込みの XProf の推奨事項
ルーフラインと最も遅いオペレーション get_roofline_model、get_hlo_op_profile

これらのツールは、広範なデータの収集と、AI に焦点を当てた調査中に実行され、次のことを確認します。

  • ワークロードまたは特定のオペレーションがコンピューティングまたはメモリ帯域幅によって制限されているかどうか
  • 行列乗算、アテンション、集団オペレーションなどのオペレーション タイプ別の所要時間、Python ソースコード行へのリンク
コンパイラの非効率性検出機能 detect_layout_mismatch_copies、detect_unfused_reshapes、detect_unnecessary_convert_reduce、list_hlo_modules

これらのツールは、広範なデータの収集時に実行され、次の項目をチェックします。

  • オペレーション間のテンソル レイアウトの不一致によって発生する余分なデータコピー
  • メモリ帯域幅を浪費する、融合されていない reshape、transpose、copy オペレーション
  • リダクション内の不要な精度変換(bf16 から f32 など)
メモリと OOM のリスク get_memory_profile、get_peak_allocations

これらのツールは、広範なデータの収集時に実行され、次の項目をチェックします。

  • デバイスのピーク時の高帯域幅メモリ(HBM)使用量と、合計容量およびメモリの断片化の比較
  • メモリ不足(OOM)のリスクが LOW、MEDIUM、HIGH のいずれであるか、およびメモリ内の最大のテンソル
計算グラフ エクスプローラ get_hlo_neighborhood

このツールは、AI の調査中に実行され、XLA グラフで低速なオペレーションの直前と直後のオペレーションをチェックし、コンパイラの融合をブロックしているものや、データ移動を強制しているものを特定します。

タイムラインとカーネルのタイミング aggregate_xplane_events、get_kernel_stats

これらのツールは、AI の詳細な調査中に実行され、次の点を確認します。

  • ホスト間のタイムライン統計情報。データローダの停止とネットワーク通信の遅延を区別します。
  • 個々のデバイス カーネルの実行時間と呼び出し回数
カスタム カーネル プロファイラ llo_load、llo_query

これらのツールは、AI の詳細な調査中に実行され、標準の高レベル プロファイリングでは不透明に見える Pallas や Mosaic などのカスタム カーネルの低レベルの実行トレースをチェックします。

AI プロフィール分析レポート

キャプチャされたプロファイルの分析が完了すると、ML 診断は生成された summary.html レポートを Cloud Storage バケットに保存し、 Google Cloud コンソールの [診断を実行] > [プロファイル] に表示します。このレポートを使用すると、ワークロードの全体的な健全性の確認、パフォーマンスのボトルネックの調査、推奨されるコードまたは構成の変更の特定を行うことができます。

生成された各レポートには、次のセクションが含まれています。

  • 分析の概要: OPTIMAL または NEEDS_ATTENTION の全体的なステータス、3 ~ 5 文の概要、ステップ時間、メモリ使用量、Roofline 効率などの主要なパフォーマンス指標が表示されます。
  • 優先順位付けされたボトルネックと対応可能な修正: 重大な問題の P0 から軽微な問題の P2 まで、優先順位付けされた検出結果を一覧表示します。各検出結果には次の情報が含まれます。
    • 何が起きているか: 根本原因の説明。関連する場合は、テンソルの形状とレイアウトも含まれます。
    • エビデンス: ボトルネックが測定された XProf 指標とオペレーションの直接引用。
    • 予想される影響と修正: 影響を受けるステップ時間またはメモリの量。推奨されるフレームワーク、シャーディング、コンパイラの変更と組み合わせて表示されます。
  • 詳細な XProf ツール分析: 次の分析から、サポートデータと詳細な内訳を提供します。
    • ルーフライン分析
    • メモリ解析
    • High-Level Optimizer(HLO)オペレーション プロファイルの分析
    • カーネル実行の統計情報
    • カスタム カーネルと低レベルのオプティマイザーの分析