メモリビューア ツール
メモリビューアを使用すると、プログラムの存続期間全体のメモリ使用量を可視化できます。メモリ使用量がピークになった時点でのメモリの内容の詳細を確認できます。メモリ不足（OOM）の状況をデバッグすることもできます。メモリビューアを使用すると、グローバル メモリ使用量と、計算がメモリ不足に近づいている程度を確認できます。
Memory Viewer に表示される情報はすべて静的で、XLA コンパイラから取得されます。動的ランタイム情報は Memory Profile ツールに表示されます。
サポートされているプラットフォーム
TPU: サポート対象
GPU: サポート対象
メモリビューアのコンポーネント
メモリ ビューアは、次の主要コンポーネントで構成されています。
- 可視化しているデータをカスタマイズできるユーザー コントロール プルダウン:
- メモリタイプ: サポートされているメモリタイプはアクセラレータによって異なります。GPU では高帯域幅メモリ（HBM）に重点が置かれますが、TPU では、VMEM、SMEM、CMEM、同期フラグ（SFlag）、Sparsecore、ホストメモリなど、オンチップメモリの使用状況も確認できます。
- モジュール: 実行の一部であった XLA プログラムです。多くの場合、適切な出発点は、jit_train_step や jit_generate などのラベルが付いたトップレベル モジュールです。
- テキストの概要セクションには、プログラムに必要なピークメモリ割り当て、引数と一時変数の分割など、概要情報が表示されます。アクセラレータでサポートされているテンソルの形状の制限により、パディングによるオーバーヘッドが発生します。このパディングが割り当て全体の大部分を占めている場合は、最適化の余地がある可能性があります。
- [メモリ割り当てサイズとプログラム順序] 折れ線グラフには、コンパイラによってスケジュールされたプログラム ポイント（HLO シーケンス）とメモリ使用量がプロットされます。
- X 軸は時間ではありません。
- グラフでは、選択したモジュールのメモリ使用率がピークになるプログラムの箇所がハイライト表示されます。プロファイラとツールは、特定のチップ/メモリにモジュールが共存していることを認識していません。各モジュールがコンパイルされると、コンパイラは、割り当てられたメモリの合計量の「ベースライン」開始点（このモジュールの前にコンパイルされたモジュールによる）をプロファイルにアノテーションし、その後の割り当てと割り当て解除が発生するたびにそこから増減をカウントします。ただし、実行前にコンパイルされた他のモジュールによる将来の割り当てはキャプチャされません。OOM の状況をデバッグする際には、この点に注意してください。
ページの下部にあるバッファグラフには、ピーク時のメモリ使用量（メモリ使用量の折れ線グラフの縦線）の詳細が表示されます。3 つのグラフがあり、すべてプログラムによって割り振られたバッファのセット全体を示していますが、3 つの異なる方法で並べ替えられています。
- プログラム順序: プログラムの実行中に作成された順序で、最も古いものが左側に表示されます。
- サイズ別: ピーク時のメモリ使用量に最も大きな影響を与えたバッファが左側に表示されます。
パディングのオーバーヘッド: ハードウェアの観点から最も「非効率的」なものが左側に表示されます。
バッファの色には特別な意味はありません。
3 つのグラフのいずれかのバッファにカーソルを合わせると、次の 2 つのディスプレイが追加で表示されます。
- メモリ使用量の折れ線グラフにバッファカードと同じ色でオーバーレイされ、そのバッファの存続期間を示します。つまり、左端と右端がプログラム順序内の割り当てポイントと割り当て解除ポイントを示す水平バーです。横棒の高さは、ピーク割り当てと比較した選択したバッファの相対サイズを表します。
- 別のバッファ詳細カード（通常は左側）に、該当する場合は特定の演算の詳細が表示されます。一般的なカードには次の情報が含まれます。
- 名前: XLA オペレーションの名前。グラフ ビューアまたはトレース ビューアで検索できます。
- サイズ: バッファ割り当てのサイズ（パディングありとパディングなし）。
- シェイプ: N 次元配列のランク、サイズ、データ型。
- フレームワーク演算名: この割り当てに関連付けられたフレームワーク演算名が表示されます。
- 割り当てタイプ: バッファ割り当てを、パラメータ、出力、スレッド ローカル、一時（fusion 内のバッファ割り当てなど）のタイプに分類します。