トレース ビューア ツール

トレース ビューアを使用すると、プロファイリング セッション中に発生したイベントのタイムラインを可視化できます。ホスト（CPU）やアクセラレータ（GPU または TPU）など、システムのさまざまな部分でモデルによって実行されたオペレーションの実行時間が表示されます。これにより、モデルがハードウェア リソースをどのように使用しているかを把握し、パフォーマンスのボトルネックを特定し、モデルを最適化して実行速度を上げることができます。Trace ビューアの UI は chrome://tracing で使用されている UI に基づいているため、Chrome ブラウザを使用する必要があります。

サポートされているプラットフォーム

TPU: サポート対象

GPU: サポートされている

GPU のトレース ビューアでは、GPU チップごとにタイムラインにイベントがグループ化され、その中のストリームがグループ化されます。一方、TPU では、TPU チップとコアごとにタイムラインにイベントが整理されます。どちらの場合も、Trace Viewer には、アクセラレータが接続されているホスト CPU のスレッドごとのタイムラインにイベントが表示されます。

タイムラインの操作

Trace Viewer には、タイムラインを移動して調べるためのツールと方法がいくつか用意されています。

• ナビゲーション: 次のキーボード ショートカットを使用できます。
  • W: 拡大します。
  • S: ズームアウトします。
  • A: 左にパンします。
  • D: 右にパンします。
• ツール選択ツール: ツール選択ツールでツールをクリックするか、対応するキーボード ショートカットを使用します。
  • 選択ツール（1 または !）: イベントをクリックして選択し、[詳細] ペインで詳細を表示します。Ctrl キーを押しながらクリックして複数のイベントを選択し、概要を表示します。
  • パンツール（2 または @）: ドラッグしてタイムラインを水平方向または垂直方向に移動します。
  • ズームツール（3 または #）: ドラッグしてタイムラインの特定の領域にズームインします。
  • タイミング ツール（4 または $）: ドラッグして時間間隔をマークします。マークされた区間の時間が表示されます。また、[m] キーを使用して選択範囲をマークし、合計時間を特定することもできます。
• 選択したイベントにズームインする（f）: 1 つ以上のイベントを選択し、f キーを押すと、タイムラインのその部分にすばやくズームインできます。これは、特定のトレーニング ステップに焦点を当てる場合に便利です。

大量のトレース イベントを処理する場合、トレース ビューアはストリーミング モードで動作します。つまり、地図アプリと同様に、タイムラインをパンやズームするときにデータがオンデマンドで読み込まれます。データの読み込みよりも速くズームすると、読み込みが完了するまで低解像度のデータが表示されることがあります。

トレース ビューアのインターフェース コンポーネント

トレース ビューアの主な UI コンポーネントは次のとおりです。

• 時間軸は上部に水平方向に表示され、トレースの開始を基準にした時間が示されます。
• タイムラインはセクションとトラックごとに整理され、ラベルは左の縦軸に配置されます。各セクションは処理要素（デバイスノードやホスト スレッドなど）を表し、展開または閉じることができます。各セクションには、特定のアクティビティのタイムラインであるトラックがあります。
• イベントは、タイムライン トラック上の色付きの長方形のブロックで、オペレーションの所要時間や、トレーニング ステップなどのメタイベントを表します。イベントの色には特別な意味はありません。
• 詳細ペインには、タイムライン ペインで選択したイベントに関する追加情報（名前、開始時間、期間など）が表示されます。

一般的なセクションとトラック

トレース ビューアには、次のセクションとトラックがあります。

• TPU ノードごとに 1 つのセクション。次のトラックがあります。
  • XLA モジュール: 実行中の XLA プログラム。
  • XLA 演算: TPU コアで実行された XLA HLO 演算が表示されます。各上位レベルのフレームワーク演算（JAX、Tensorflow、PyTorch など）は 1 つ以上の XLA 演算に変換され、TPU で実行するようにコンパイルされます。XLA オペレーションをクリックすると、グラフ ビューアにオペレーションへのリンクと、オペレーションに関する追加情報（開始時間、停止時間、所要時間、ソース スタック トレースなど）が表示されます（フレームワークとコンパイラによって利用可能な場合）。
  • XLA TraceMe: トラッキングする論理的な作業単位を記述する、コード内のユーザー指定のアノテーション。アノテーションを追加しなかった場合でも、ここにデータが表示されることがあります。通常、これらのデータは XLA（バリア コアなど）または XProf 自体（ドロップされたトレース エントリなど）によって追加されます。
  • ステップ: ユーザー プログラムまたはフレームワークで適切にアノテーションが付けられている場合、その TPU コアで実行されているトレーニング ステップの所要時間が表示されます。
  • Framework Ops: ユーザー プログラムまたはフレームワークで適切にアノテーションが付けられている場合、TPU コアで実行されるフレームワーク オペレーション（JAX、Tensorflow、PyTorch など）が表示されます。
  • フレームワーク名スコープ: フレームワーク オペレーションごとに、スタック トレースのイメージ。簡潔にするため、このトラックは 1 台のデバイスにのみ表示されます。
  • ソースコード: 実行中のソースコードのパス（プロファイルで利用可能な場合）。
  • スカラー単位: TPU の場合、スカラー単位で実行されるイベント。プロファイルに存在する場合に表示されます。
  • TensorCore Sync Flags: TPU の同期メカニズム。プロファイルに存在する場合に表示されます。
  • ホスト オフロード: ホストメモリとアクセラレータ メモリ間でデータを非同期的に移動するオペレーション。通常、XLA Ops 行には、アクセラレータがデータ転送の準備をしていることを示している対応する開始オペレーションと停止オペレーションが表示されます（転送中は、送信元/宛先メモリ領域が「使用中」としてマークされます）。複数のオフロード オペレーションが並列に実行されている場合、複数のホスト オフロード行が存在することがあります。この場合、トレース ビューアは複数のイベントを同時に表示する必要があります。
• Sparsecore ノードごとに 1 つのセクション: 一部の TPU 世代（TPU v5p と TPU v6e）には、高密度コンピューティング MXU ユニットに加えて、1 つ以上の SparseCore ユニットが搭載されています。これらのコアに関連付けられたモジュール、オペレーション、TraceMes がこのセクションに表示されます。
• GPU ノードごとに 1 つのセクション。次のトラックが含まれます。
  • XLA モジュール、フレームワーク オペレーション、フレームワーク名スコープ、ステップ、ソースコード。これらはすべて TPU セクションに似ています。
  • ストリームごとに 1 つのトラック。ストリーム名には、ストリームで実行されるオペレーションの種類（Memcpy、Compute など）に関する情報も含まれます。
  • XLA TraceMe は GPU ではサポートされていません。
  • XLA オペレーションは GPU セクションに表示されますが、ストリームデータから派生するため、現在のところ常に正確であるとは限りません。したがって、異なるストリームで実行される実際のカーネルへの XLA オペレーションの N:M マッピングや、ハードウェア内の異なる SM への複数のストリームの動的スケジューリングが存在する可能性がある GPU の実行モデルを完全に考慮することはできません。
• ホストマシンの CPU 上で実行されるコンポーネント（スレッドプールなど）ごとに 1 つのセクション。スレッドプールの場合は、スレッドごとに 1 つのトラック。また、プロファイル収集中に Python トレースが有効になっている場合は、ここに表示されます。

収集されたプロファイルに直接関連付けられるのは、TPU の XLA オペレーションと GPU のストリーミング データのみです。他のすべての行は「派生行」であり、コンパイラから提供されるオプションの「サイドバンド」情報、オプションのユーザー アノテーション、XProf によって適用されるヒューリスティクスが含まれます。そのため、これらの派生ラインは、特定のプロファイルに表示されることもありますし、表示されないこともあります。

その他の機能

• [イベントを検索] 検索バーを使用して、特定のイベント名を検索できます。現在、これは完全なトレースではなく、画面上の表示時間枠内でのみ検索されます。
• フローイベント: 上部のバーにある [フローイベント] ボタンをクリックしてこのオプションを有効にすると、可視化が追加され、1 つのスレッドまたは行のイベントが別のスレッド行のイベントにリンクされます。たとえば、アクセラレータの作業をキューに登録または起動するホストのオペレーションから、その作業を実行するアクセラレータのオペレーションに矢印が引かれる場合があります。XProf は、ユーザー アノテーション、組み込みヒューリスティクス、およびさまざまなコンポーネント（CUPTI ドライバ、カーネル起動 ID、TPU ランタイム情報など）。
• XLA オペレーションをクリックすると、詳細ペインに追加情報が表示されます。たとえば、グラフ ビューア ツールの op にリンクします。また、ソースコードや Python スタック トレース、この XLA オペレーションの生成につながったフレームワーク オペレーションなどへのポインタも提供されます（プロファイルに存在する場合）。FLOPS（オペレーションによって実行される浮動小数点演算の数）とオペレーションによってアクセスされたバイト数も表示されます。この情報は、プロファイルのランタイム情報ではなく、コンパイル時に XLA から静的に取得されます。