キャプチャからのJAX OOMのデバッグ
概要
JITコンパイルされたJAXの各ステップは、すべてのスクラッチメモリを1つの計画されたテンポラリブロックとして取得します。したがって、ステップのメモリ使用量はXLAの計画であり、.xplane.pb ファイルにはその計画とアロケータのランタイムイベントの両方が記録されています。scripts/oom_report.py スクリプトは、標準ライブラリのみを使用して、両方を2〜5秒で読み取ります。
- 計画: 各モジュールのバッファ割り当て（
/host:metadataから取得）。パラメータ（donateされた量を含む）、非エイリアス出力、テンポラリブロック、ヒープピークが示されます。ピーク時に生存しているバッファについては、サイズ、形状、役割、Pythonの行番号が表示されます。役割はJAXの名前スタックから取得されます:jvp(はフォワード（バックワードパスのために保持）、transpose(はバックワード、s8カスタムコールの出力はライブラリのワークスペースです。 - ランタイム: BFCアロケータのイベント。容量、各実行開始時の使用バイト数、ピーク、
AllocatorRetry待機、カーネルを一度も起動しなかった実行、最後の2回の実行間のホストイベントが示されます。また、最後の実行開始時に生存していたバッファが、それを割り当てたホストイベントごとにグループ化されて表示されます:outputs of jit_step execution 3、BufferFromHostBuffer(s32[…])（device_put）、または「allocated before the capture started」。
レポートの最後には診断行が出力され、それぞれに measured（測定値）または inferred（推測値）のタグが付与されます。まず以下のコマンドで実行してください。
python3 -P <this skill's base directory>/scripts/oom_report.py <capture dir or .xplane.pb> [--module 'jit_train_step(123)'] [--top 25] [--json out.json]
python3 -P <this skill's base directory>/scripts/oom_report.py --self-test
デフォルトでは、最後の実行でカーネルを起動しなかったモジュールを解析します。該当するモジュールがない場合は、実行された最大のモジュールを解析します。--module オプションで別のモジュール（例えば評価ステップ）を指定することもできます。
診断行から原因と修正へ
| 診断行 | 原因 | 提案する修正；節約量はレポートのバイト数から判断する |
|---|---|---|
the step itself doesn't fit | 計画単独でプール容量を超えている | 以下のピーク時の役割分割に従う |
↳ ほとんどが forward, kept for the backward pass | アクティベーション：バッチ×シーケンス長に比例して増加；単純なアテンションではレイヤーごとに B·H·S·S | jax.checkpoint / 再計算（remat）をブロックに適用（フォワード保持バイトの大半を節約、約1回分のフォワードパスが余計に必要）；より小さいマイクロバッチと勾配累積；より短いシーケンス長（S²のバッファはSが半減すると4倍縮小）。融合アテンションは「B·H·S·S バッファが計画から外れることを確認する」場合のみ提案する |
↳ ほとんどが backward | 勾配とそのテンポラリ | シャーディング（FSDP）または累積；bf16勾配のf32コピーがないか確認 |
↳ workspace | GEMM / conv ライブラリのスクラッチ領域 | 通常は小さい；ROCmではgfx950上のhipBLASLt GEMMあたり64 MiB |
grow by X MiB per execution, = 1.00x the non-aliased outputs | プログラムが各ステップの出力を生存させ続けている（ステップ間リーク）：生存リストに outputs of … execution 1, 2, … と表示される | 実行間のホストイベントで保持者を見つける（例：$<unknown> append = list.append）。その後、参照を削除する、ステップ内で削減する（ノルムであってツリー全体ではない）、または jax.device_get する。プールを大きくしてもOOMを遅らせるだけ |
no parameter is donated … can save up to X | 古いパラメータ/最適化器の状態と新しいものが同時に生存している | 状態引数に donate_argnums を指定する |
other live arrays take the room; the largest … came from X | ステップ外の配列：データセットや評価コピーがGPU上にある、プリフェッチ、他のJIT出力 | ホスト上に保持してストリーミングする、解放する、またはデータ型を小さくする。「Allocated before the capture started」は、より早い時点でプロファイルを取得して名前を特定する必要があることを示す |
AllocatorRetry … longest N s, launched none | キャプチャ内でOOMが発生した | 証拠のみ。失敗したリクエストは通常、そのモジュールのテンポラリブロック、または1つの出力 |
XLA already rematerialized N instructions | XLA独自の再計算が既に制限に対して実行されている | 修正はXLAではなくプログラム側から期待する |
各修正には推定節約量とコストを付記し、inferred とラベル付けします。見積もりが容量を超えていないか確認してください。「XLA_PYTHON_CLIENT_MEM_FRACTION を上げる」は、GPUに空き容量がある場合にのみ、収まらない計画を修正します。容量がHBMを大幅に下回っている場合は、フラクションが設定されている可能性があります。キャプチャには環境変数は記録されません。
プレイブックのメモリ構造をJAX用語で
Ultra-Scale Playbook（「メモリ使用量のプロファイリング」）では、トレーニングメモリを重み、勾配、オプティマイザ状態、アクティベーションに分割しています。レポートでは以下のように対応します。
- 重み + オプティマイザ状態 = ステップのパラメータ。Adam fp32ではパラメータあたり12 B（p, m, v）。Donateされていればコピーは発生しません。
- 勾配 = ピーク時の
backward役割、およびステップが返す場合の出力。 - アクティベーション =
forward, kept for the backward pass。Scan /nn.scanモデルでは、whileループ上にスタックされた[layers, …]バッファとして保持されます。 - 「ステップ1は収まるが、ステップNで失敗する」: PyTorchでは最初のステップでオプティマイザ状態が作成されることが原因です。JAXでは状態はステップ1の前に存在するため、代わりに成長ラインを探します。最初の呼び出しはコンパイルとオートチューニングも行うため、最初の実行前のピークはコンパイル時のものです。
ルール
- キャプチャのみを使用すること。 修正をテストするためにモデルを再構築したりワークロードを再実行したりしないでください。計画のバイト数から見積もり、ユーザー側で確認する方法（
jit(...).lower(...).compile().memory_analysis()）を提案してください。 - キャプチャディレクトリに書き込まないこと。 XProfのコンバータ（
xspace_to_tool_names、memory_viewer、xprofCLI）は、.xplane.pbの隣に*.hlo_proto.pbキャッシュファイルを書き込みます。XProfビューが必要な場合は、まず.xplane.pbをスクラッチディレクトリにコピーしてください。 - すべての数値に measured（測定値）、read（読み取り値）、inferred（推測値）のラベルを付けること。
測定リファレンス（単一MI355X、小プール、2026-11-06）
| キャプチャ | レポートの表示内容 |
|---|---|
| アテンション計画が大きすぎる | 3.72 GiB 必要 > 2.79 GiB；ヒープピークの81.6%がフォワード保持；softmax行に f32[8,8,1024,1024] ×5；テンポラリブロック 3,495 MiB = エラーの「3.41GiB」；AllocatorRetry 10 秒；89個の再計算命令 |
| 各ステップで勾配が保持される | ステップ開始時の使用量 316 → 1,252 MiB、ステップあたり +104.0 MiB = 非エイリアス出力の1.00倍；計画 788 MiB は収まる；ステップ間に $<unknown> append |
| 引数のdonateなし | 3.88 GiB 必要 > 2.79 GiB；パラメータ 1,902 MiB + 出力 1,902 MiB、donate 0；テンポラリはわずか 172 MiB |
| MaxText llama2-7b、8 GPU、FSDP（収まる場合） | テンポラリ 61.9 GiB；ヒープピークの57%がフォワード保持、スキャンスタックは while.50 による bf16[32,4,4096,11008] ×2 |
制限
- ファイルあたり1ホスト: マルチホスト実行ではホストごとに1つの
.xplane.pbが生成されるため、各ホストで実行する必要があります。 - HBMのみ: メモリ空間0のみが詳細分析されます。ホストオフロードは
memory space 5として表示されます。 - 連続性: 「収まる」はバイト数を比較しています。BFCは連続したチャンクも必要とするため、バイト数が近くてもフラグメンテーションで失敗する可能性があります。
- 「other」役割: JAXの名前スタックを持たないXLA作成のオペレーションによって定義されたバッファ、およびそのような名前スタックを持つエイリアスがないバッファ。
- CUDA: 未検証。同じBFCおよびXLAイベントを使用しますが、名前は未確認です。
- TPU / Eager: TPUキャプチャは計画の合計のみを提供し、ヒープピークや実行はありません。Eager（オペレーション単位）のプログラムは多数の小さなモジュールとして表示されるため、その場合はランタイムセクションを参照してください。