デベロッパー

新しい XProf プロファイラと Cloud Diagnostics XProf ライブラリで xPU の ML パフォーマンスを強化

2025年10月1日

https://storage.googleapis.com/gweb-cloudblog-publish/images/xprof-hero-image-2436x1200_k5YQpVQ.max-2200x2200.png

Rajesh Anantharaman

Product Management Lead, Google Cloud

※この投稿は米国時間 2025 年 9 月 15 日に、Google Cloud blog に投稿されたものの抄訳です。

ML モデルの構築よりも、パフォーマンスのデバッグに時間を費やしている感覚をお持ちの方は多くいらっしゃいます。昨今の急速なペースで進化する AI の状況においては、新しいモデルアーキテクチャの探求から、常に変化するハードウェアとソフトウェアのスタックへの対応までが必要になるため、モデルの最適化は複雑な課題です。Gemini、検索、YouTube などのプロダクト向けにモデルを最適化している Google でさえ、まさにこうした課題に直面してきました。

そこで、Google は社内ツールボックスからコミュニティに貴重なツールを提供することになりました。それが、Google のコア ML プロファイラである XProf です。この投稿では、この強力なツールと新しい Cloud Diagnostics XProf ライブラリを使用して、モデルのボトルネックを簡単に特定し、最適化の隠れた機会を見つけ、本来の業務であるべきイノベーションに再び取り組める方法をご紹介します。

XProf のアップデート

数年前、Google は XProf ツールのフォークを「TensorFlow Profiler」というオープンソースとしてリリースしました。これは TensorBoard を通じて「tensorboard_plugin_profile」としてアクセスでき、コミュニティの多くの ML エンジニア（MLE）が xPU 上の ML モデルを最適化するために使用していました。しかし、市場の急速な変化と Google の社内ツールである XProf の大きな進歩を考慮して、これらの高度な機能をコミュニティと共有する時期が来たと判断しました。最近、社内用 XProf の機能を社外用ツールに統合し、社外用コードベースを OpenXLA に移行しました。また、TensorFlow と TensorBoard 以外にも進化していることを反映して、社外用プロファイラを「XProf」というシンプルな名前に変更しました。これにより、コミュニティは Google と Google DeepMind で使用されているのと同じプロファイリングツールの利便性を享受できます。XProf を OpenXLA の下に移行するもう一つの大きなメリットは、JAX、Pytorch / XLA、TensorFlow / Keras など、すべての XLA ベースのフレームワークにおいて一貫したエクスペリエンスで平等にサポートできるようになることです。

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

XProf を使用するには、まずモデルのワークロードコード内でプロファイルキャプチャを有効にする必要があります。プロファイルをキャプチャするには、次の 2 つの方法があります。

プログラマティックキャプチャ – モデルのコードにアノテーションを付け、コード内でプロファイルをキャプチャする場所を指定します。
オンデマンドキャプチャ（手動キャプチャ）– プログラマティックプロファイルキャプチャが無効であった期間でも、実行中にアドホックにプロファイルをキャプチャできます。また、実行中にモデル指標で問題が発生した場合、問題を診断するために、その瞬間のプロファイルを一定期間キャプチャする場合にも使用できます。

プロファイルの可視化

キャプチャしたプロファイルを可視化するには、ローカルまたは Google Cloud で TensorBoard を使用することをおすすめします。このセクションでは、GitHub でオープンソースとして提供されている Cloud Diagnostics XProf（別称: XProfiler ライブラリ）というライブラリを使用して、XProf と TensorBoard で Google Cloud 上のシームレスなセルフホスト型プロファイリングエクスペリエンスを実現する方法に焦点を当てます。

このライブラリを使用すると、XProf と TensorBoard をローカルで実行する場合と比べて、次のような改善が図られます。

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

Cloud Diagnostics XProf ライブラリを設定する

まず、gs://<some_bucket>/<some_run>/ のようなストレージパスでプロファイルをキャプチャします。

次に、gCloud CLI と Python 仮想環境を設定します。

最後に、次のコマンドでライブラリをインストールします。

pip install cloud-diagnostics-XProf

XProfiler を使用して TensorBoard インスタンスを作成する

TensorBoard をホストする VM または GKE Pod を作成します。

XProfiler create -z $ZONE -l $GCS_PATH

または

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

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

TensorBoard プロファイルプラグインに最近加えた大きな変更の一つは、プロファイルパスの要件を緩和したことです。これにより、すべての TensorBoard インスタンスが、ディレクトリ内の複数の実行からプロファイルを読み込めるようになりました。つまり、TensorBoard インスタンスを別途起動しなくても、実行間や長時間実行されるジョブのセッション間でプロファイルを簡単に比較できるようになるため、分析ワークフローが高速になります。XProfiler の作成時に、TensorBoard インスタンスをルートストレージパス gs://<some_bucket> に指定すると、TensorBoard はこのルートディレクトリの下にあるすべてのプロファイルを取得できます。TensorBoard UI では、すべての実行とセッションのプロファイルが、プルダウンで run1/session1、run1/session2、run2/session1 として表示されます。

https://storage.googleapis.com/gweb-cloudblog-publish/images/XProf_Blog_images.max-1000x1000.png

プロファイルの読み込みとツール切り替えの高速化

TensorBoard XProf プラグインに加えたもう一つの大きな変更は、読み込み時間の改善です。プラグインにマルチスレッドのサポートを追加したことで、より大きなプロファイルをはるかに高速に読み込めるようになりました。また、キャッシュ機能を追加したことで、2 回目の読み込みがさらに高速になり、ユーザーがパフォーマンスの最適化を行いながら、XProf のツール間をスムーズに移動できるようになりました。

Xprofiler を使用して TensorBoard と XProf をホストするためのマシンオプション

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

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

プロフィールを共有するリンク

XProfiler create を実行すると、次のようなメッセージが表示されます。

Instance for gs://<some-bucket> has been created.

You can access it via following,

1.https://<id>-dot-us-<region>.notebooks.googleusercontent.com.

2. XProfiler connect -z <some zone> -l gs://<some-bucket> -m ssh

Instance is hosted at XProf-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

注: 最初のオプション（1）が、作成されたリンクです。クリックすると、TensorBoard で XProf プロファイルを表示できます。パフォーマンスの最適化は、非常に反復的で共同作業を必要とするプロセスです。この共同作業を可能にするために、Cloud Diagnostics XProf ライブラリは TensorBoard インスタンスへのリンクを作成し、チーム全体や Google Cloud でのパフォーマンスの最適化を支援する Google エンジニアとプロファイルを簡単に共有できます。リンクにアクセスできるユーザーは、TensorBoard インスタンスが指す Cloud Storage バケットに設定された権限に基づいて制御されます。

リンクが機能しない場合は、SSH を介して TensorBoard インスタンスに接続し、XProfiler connect コマンドを実行してプロファイルを表示する方法もあります。

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

ワークロードコードでプロファイラサーバーを有効にして、オンデマンドプロファイリングを実行する場合は、次の 2 つの方法があります。

TensorBoard UI の [プロファイルのキャプチャ] ボタンをクリックします。GKE と Compute Engine で実行されているワークロードのオンデマンドキャプチャがサポートされています。
CLI で XProfiler キャプチャを使用し、TensorBoard UI の [プロファイルのキャプチャ] ボタンで指定するのと同様の情報を指定します。

XProf の新機能

アップデートされた XProf では、最も人気のあるツールのさまざまな機能が更新されています。

トレースビューア
メモリビューア / メモリプロファイル
グラフビューア
HLO Op プロファイル / HLO Op 統計
概要ページ

特に、メモリビューアでは、HBM（高帯域幅メモリ）、ホスト、TPU の SparseCore、VMEM、SMEM、CMEM、Sync Flags（SFlag）など、7 種類のメモリを確認できるようになりました。

また、トレースビューアと HLO Op プロファイルからのリンクを使用すると、すべてのオペレーションでグラフビューアにシームレスに移動できます。また、より多くのオペレーションをカバーできるよう、ソースコードの行の可視性も向上させました。

XProf を使用してパフォーマンスのボトルネックを見つけるために使用される最も一般的なフローは、次のようになります。

MLE が TensorBoard で XProf を開き、トレースビューアまたは HLO Op プロファイルでさまざまなオペレーションを確認する
関心のあるオペレーションをクリックし、詳細を調べる
このオペレーションに対するグラフビューアへのリンクをクリックして、モデル内でオペレーションがどのように配置されているかを確認する
メモリビューアで、モデルと特定のオペレーションに対して HBM / ホスト / SparseCore メモリが効率的に使用されているかどうかを確認する
最適化したいオペレーションを決定し、その実装のために、オペレーションのソースコード行を確認する

アップデートされた XProf ツールは、このフロー全体をスムーズかつ簡単にします。

新しい XProf ツール

さらに、XProf には以下のような新しいツールもいくつかリリースされています。

フレームワークオペレーション統計 – フレームワークレベルのオペレーション（JAX または TensorFlow）におけるパフォーマンス統計情報。
ルーフライン – プログラムがメモリまたはコンピューティングの制約を受けているかを視覚的に確認し、プログラムのパフォーマンスがハードウェアの理論上のピーク時のパフォーマンスにどれだけ近いかを「ルーフライン」として表します。
メガスケール統計 – データセンターネットワーク（DCN）を介して通信する複数の TPU スライスにまたがるワークロードのマルチスライス通信パフォーマンスを分析します。
GPU カーネル統計 – プロファイリングセッション中に起動されたすべての GPU アクセラレーテッドカーネルのパフォーマンス統計と、そのカーネルの元になったフレームワークオペレーション。

XProf での Pallas カーネルの可視性

XProf で要望の多かったパフォーマンスの可視化の大きな領域の一つは、Pallas カーネルのパフォーマンスの可視化です。これらのカーネルは XProf では「カスタム呼び出し」として表示されますが、カスタム呼び出しのパフォーマンスとその実装の詳細を確認することは困難でした。このたび、XProf 内で Pallas カーネルのサポートと可視性が向上したことをお知らせいたします。これで、HLO Op プロファイルとグラフビューアの両方で、Pallas カーネルの詳細を確認できます。Pallas カーネルのカスタム呼び出しごとに、共通の Pallas カーネルの場合はカーネルの名前が表示されます。名前をクリックすると、サイドパネルにカーネルのパフォーマンスなどの情報が表示されます。サイドパネルでカーネルをクリックしたときに正確なパフォーマンス指標を取得するには、カーネルの作成者が pl.CostEstimate オブジェクトを pallas_call 関数に渡して、費用モデルを提供する必要があります。また、[カスタム呼び出しテキスト] ボタンをクリックすると、Pallas カーネルの実装に関する詳細を確認できます。

https://storage.googleapis.com/gweb-cloudblog-publish/images/XProf_Blog_images_1.max-1000x1000.png

XProf を使用した CUDA グラフのトレース

XProf は、XLA を使用するワークロードの NVIDIA GPU でのプロファイリングもサポートしています。XProf ツールのほとんどで NVIDIA GPU がサポートされていますが、GPU トレースは CUDA ストリームごとに整理されるため、ビューが異なります。前述したように GPU での JAX のプロファイルをキャプチャします。ほとんどの XProf ツールは、NVIDIA GPU のパフォーマンスビューをサポートしています。XProf でサポートが追加された新機能の一つが CUDA グラフです。CUDA グラフを使用すると、個々のオペレーションを組み合わせたグラフとしてより大きな作業単位を定義できるため、カーネルの起動オーバーヘッドを最適化できます。XProf では、CUDA グラフを使用するユーザーが、XProf でそれらをトレースして、パフォーマンスをさらに最適化できるようになりました。

CUDA グラフのトレースを有効にするには、関連する XLA フラグ <XLA_FLAGS=--xla_enable_command_buffers_during_profiling> を有効にする必要があります。

これにより、XProf トレースで CUDA グラフを確認できます。グラフビューアで CUDA グラフを確認することもできます。

CUDA グラフ内で分割されたさまざまなノードのトレースを確認するには、ProfilerOptions の詳細構成値 <gpu_enable_cupti_activity_graph_trace , True> を追加します。

グラフ内でトレースされた詳細なノードと、グラフノードの実行が開始された特定のストリームにホストをリンクする矢印が表示されます。

https://storage.googleapis.com/gweb-cloudblog-publish/images/Screenshot_2025-09-15_at_11.47.25AM.max-2200x2200.png

お客様の声

さまざまなフレームワークで TPU と GPU を使用している多くのお客様が、更新された XProf ツールと Cloud Diagnostics XProf ライブラリの恩恵を受けています。以下に、その一部をご紹介します。

「当社はテクノロジーハブとして、世界最高峰の専門家たちが最先端テクノロジーを駆使して次世代のアプリをビルドしています。XProf ツールが改善されたことで、モデルを最適化し、TPU 上の JAX と PyTorch XLA の両方で生産性を高めることができました。XProf を使用すると、トレースビューアで特定されたボトルネックから、モデルのグラフビュー、ソースコードへとドリルダウンできるため、何を変更すべきか把握できます。これで、TPU で最高のパフォーマンスを発揮できるようにコードをチューニングすることに集中できます。Cloud Diagnostics XProf ライブラリを使用すると、プロファイルを簡単にキャプチャして読み込み、長期的な分析を行うことができます。また、リンクを共有することで、他のユーザーとのコラボレーションも簡単に行えます。」

- HubX、シニア ML 研究者 Mustafa Ozuysal 氏

「Mathpix は、コンピュータビジョンの最先端で、個人や企業がドキュメントを扱う方法を変革しています。画像推論のパフォーマンスに関する課題では、XProf を使用してプロファイルトレースを収集、分析し、キャッシュの使用が問題であることを発見しました。このコードをチューニングすることで、TPU での JAX のパフォーマンスが向上しました。Cloud Diagnostics XProf ライブラリのおかげで、Google Cloud で XProf を簡単に設定し、プロファイルを非常に高速に読み込むことができました。」

- Mathpix、シニア ML エンジニア Remy Ochei 氏

まとめ

アップデートされた XProf と新しい Cloud Diagnostics XProf ライブラリは包括的でエンドツーエンドの共同プロファイリングソリューションを提供します。これらを使用すると、コンピューティングと通信のオーバーラップのボトルネックの特定、デバイス上のメモリ使用量の最適化、カスタムカーネルのパフォーマンスに関する詳細な分析情報の取得など、さまざまなニーズに対応できます。Google Cloud でモデルのパフォーマンスを最大限に引き出すために、Google の社内ツールをコミュニティ全体に提供することとなりました。

以下のリソースを確認して、今すぐお試しください。

XProf ツールのドキュメント: https://openxla.org/XProf
XProf github: https://github.com/openxla/XProf
Cloud-diagnostics-XProf ライブラリの GitHub: https://github.com/AI-Hypercomputer/cloud-diagnostics-XProf
Google Cloud でのプロファイリング: https://cloud.google.com/tpu/docs/profile-tpu-vm

また、「TPU と GPU で ML パフォーマンスを解き放つ」という講演もご覧ください。この講演では、前述のツールの使用方法、XProf ツールの簡単な紹介、TPU と GPU で XProf を使用した実際のパフォーマンス最適化シナリオについて説明しています。

謝辞

この作業は、Google Cloud だけでなく、CoreML や Google 社内の複数のチームが連携する大規模な取り組みでした。これらのプロダクトに多大な貢献とサポートをしてくれた Kan Cai、Vaibhav Tyagi、Victor Geislinger、Pavel Dournov、Navid Khajouei、Clive Verghese、Yin Zhang、Kelvin Le、Matt Hurd、Mudit Gokhale、Sai Ganesh、Vikas Aggarwal、Sannidhya Chauhan、Stephanie Morton、Subham Soni、Ani Udipi、Newfel Harrat、Aspi Siganporia、Bryan Massoth、Chetna Jain、David Duan、George Vanica、Jiten Thakkar、Lei Zhang、Jiya Zhang、Jonah Weaver、Aditya Sharma、Fenghui Zhang、Bill Jia、Alex Spiridonov、Niranjan Hira に心より感謝いたします。

ー Google Cloud、プロダクト管理リード Rajesh Anantharaman

投稿先