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 |
| ワークロード モニタリングからのプロファイルの自動キャプチャ(プレビュー) |
ワークロード モニタリングのパフォーマンス低下アナライザが この方法は、パフォーマンスが低下したときに、手動でキャプチャをトリガーするのを待たずに、ML 診断でトレースデータをキャプチャする場合に使用します。 |
JAX、vLLM、Tunix、TorchTPU |
プログラムによるキャプチャ
プログラムによるキャプチャでは、プロファイリング コマンドをワークロード コードに直接埋め込み、プロファイル データの記録を開始および停止するタイミングを明示的に指定します。ML 実行の gcs_path パラメータで指定された Cloud Storage バケットに書き込まれたプロファイルは、プロジェクトのマネージド XProf インスタンスにプロファイル セッションとして自動的に表示されます。これは、Google Cloud コンソールまたは Google Cloud CLI を使用して表示できます。
プログラムによるキャプチャは、次のいずれかの方法で構成できます。
- ML Diagnostics SDK: SDK の
profメソッドを使用します。これらのメソッドは、基盤となるフレームワーク レベルのプロファイリング API をラップするため、プロファイリング コードは JAX、vLLM、Tunix、TorchTPU で一貫性を保ちます。 - フレームワーク ネイティブのプロファイリング API: ML フレームワークの組み込みプロファイリング メソッドを直接呼び出し、出力を ML 実行の Cloud Storage バケットに保存します。
- JAX: プロファイリング計算
- vLLM: Cloud TPU での vLLM のプロファイリング
- TorchTPU: XProf を使用したパフォーマンス プロファイリング
オンデマンド プロファイリング
オンデマンド プロファイリングを使用すると、ワークロードの実行中にプロファイルのキャプチャをリアルタイムでトリガーできます。オンデマンド キャプチャをトリガーすると、指定した期間にわたってプロファイルが記録されます。オンデマンド プロファイリングは、ワークロード コードで直接有効にするか、ML Diagnostics SDK を使用して有効にできます。
前提条件
オンデマンド プロファイルをキャプチャする前に、次の設定手順を完了します。
ワークロードのポート
9999で XProf サーバーを起動して、オンデマンド キャプチャ トリガーに応答できるようにします。JAX で直接: ポート
9999でprofiler.start_serverメソッドを呼び出します。profiler.start_server(port=9999)ML Diagnostics SDK: ML 実行を定義するときに、
machinelearning_runでon_demand_xprof=Trueを設定します。
クラスタ オーケストレーターまたはコンピューティング環境を構成します。
- Google Kubernetes Engine(GKE)またはセルフマネージド Kubernetes: オペレーターと Webhook、または ML 診断 DaemonSet をデプロイします。
- Cluster Director で管理される Slurm またはセルフマネージド Slurm: Slurm プロローグ スクリプトを使用して ML 診断デーモンをデプロイします。
- Compute Engine: インスタンス起動スクリプトを使用して ML Diagnostics デーモンをデプロイします。
オンデマンド プロファイルをキャプチャする
オンデマンド プロファイル キャプチャは、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以降。
- クラスタ オーケストレーターまたはコンピューティング環境の設定:
- GKE またはセルフマネージド Kubernetes: ML Diagnostics DaemonSet をデプロイします。
- Cluster Director で管理される Slurm またはセルフマネージド Slurm: Slurm prolog スクリプトを使用して ML Diagnostics デーモンをデプロイします。
- Compute Engine: インスタンス起動スクリプトを使用して ML Diagnostics デーモンをデプロイします。
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 コンソールで継続的なプロファイリングを開始する手順は次のとおりです。
- Google Cloud コンソールで、[Cluster Director] または [GKE > AI/ML > 診断を実行] に移動し、ML 実行を選択します。
- [プロファイル] タブを開きます。
- 上部のアクションバーで、[継続的なプロファイリングを開始] をクリックします。
- プロファイリングするノードとデバイス トレーサー レベルを選択し、[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 コンソールで継続的なプロファイリング スナップショットをキャプチャする手順は次のとおりです。
- Google Cloud コンソールで、[Cluster Director] または [GKE] > [AI/ML] > [診断を実行] に移動し、ML 実行を選択して、[プロファイル] タブを開きます。
- 上部のアクションバーで、[スナップショットを作成] をクリックします。
- [継続的プロファイリング スナップショットをキャプチャ] ペインで、プロファイル名を入力し、ターゲット ノードを選択して、[送信] をクリックします。
スナップショットが [プロファイル] テーブルに表示されます。リンクがすぐに表示されない場合は、[更新] をクリックします。
ワークロード モニタリングからのプロファイルの自動キャプチャ
ワークロード モニタリングからの自動プロファイル キャプチャを使用すると、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 プロファイルで次の分析ワークフローを実行します。
- 広範なデータの収集: 15 個の XProf ツールを並行して実行し、ステップ時間、ハードウェア使用率、メモリ使用量、ルーフライン境界、最も遅いオペレーション、コンパイラの非効率性などの健全性スナップショットを収集します。このステップでは、トレーニング ステップごとにタイミングを正規化し、正常なオペレーションをボトルネックから分離します。
- Focused AI investigation: 収集されたヘルス スナップショットを確認します。ワークロードのコンピューティング使用率が高く、アイドル時間と通信オーバーヘッドが低く、メモリ不足のリスクが低い場合、エージェントは追加のツール呼び出しをスキップして
OPTIMALレポートを生成します。ボトルネックが検出されると、エージェントは仮説を立て、追加の XProf ツールを使用して最大 4 つのターゲットを絞ったフォローアップ チェックを実行し、計算グラフ内の周囲のオペレーション、ホスト間のタイムライン遅延、カスタム カーネルを検査します。 - レポートの生成: 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_roofline_model、get_hlo_op_profile |
これらのツールは、広範なデータの収集と、AI に焦点を当てた調査中に実行され、次のことを確認します。
|
| コンパイラの非効率性検出機能 |
detect_layout_mismatch_copies、detect_unfused_reshapes、detect_unnecessary_convert_reduce、list_hlo_modules |
これらのツールは、広範なデータの収集時に実行され、次の項目をチェックします。
|
| メモリと OOM のリスク |
get_memory_profile、get_peak_allocations |
これらのツールは、広範なデータの収集時に実行され、次の項目をチェックします。
|
| 計算グラフ エクスプローラ | 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)オペレーション プロファイルの分析
- カーネル実行の統計情報
- カスタム カーネルと低レベルのオプティマイザーの分析