analyze-memory-snapshot

メモリスナップショットを分析する

アロケーターの状態とイベント履歴を調べて、PyTorch の CUDA メモリスナップショットを分析します。一般的なメモリに関するアドバイスより、スナップショットの証拠を常に優先してください。

ワークフロー

1. ユーザーが指定していない場合は、スナップショットのパスを確認します。.pickle、.pkl、.json、またはランクごとのスナップショットを含むディレクトリを受け付けます。ユーザーが memory_viz のスクリーンショットしか持っていない場合は、元のスナップショットファイルを依頼してください。

2. スナップショットをローカルで読み込みます。PyTorch の _dump_snapshot() ファイルは pickle 形式なので、ユーザーが提供したファイル、またはすでにワークスペースにあるファイルのみをアンピクルしてください。JSON スナップショットはそのまま読み込めます。

3. 利用できる場合は、このスキルのヘルパーを実行して初期概要を取得します。

   python scripts/analyze_snapshot.py /path/to/snapshot.pickle
   

   複数のランクがある場合は、ランクごとに実行し、ピーク値、アクティブバイト数、非アクティブバイト数、OOM イベント、上位のスタックトレースを比較します。

   同じ実行の異なる時点で取得したスナップショットが 2 つある（または取得できる）場合は、それらの差分を取ってください。これにより、リークの調査を解釈から、増加した割り当て箇所の特定に切り替えられます。

   python scripts/analyze_snapshot.py step100.pickle --diff step200.pickle
   

   非常に大きなスナップショット（トレース履歴が数百 MB）の場合、ヘルパーはトレース分析を --max-trace-events（デフォルトは 1M）までに制限し、削除した古いイベントの数を報告します。これにより、分析に必要なデータをすべて読み込まずに利用できます。

4. ヘルパーの出力だけでは不十分な場合は、未加工のスナップショットを読み取ります。以下を確認します。

   • segments：cudaMalloc によって確保されたアロケーターのセグメント
   • blocks：各セグメント内の使用中、解放待ち、非アクティブな領域
   • device_traces：メモリ履歴が有効な場合の、時系列順の割り当てイベント
   • frames：割り当て箇所のスタックトレース

5. アロケーターの用語に基づいてメモリの状況を報告します。

   • 予約済みメモリ：セグメントの total_size の合計
   • 使用中のテンソルメモリ：active_allocated にあるブロック
   • 遅延解放：active_awaiting_free にあるブロック。多くの場合、ストリームに関連しています
   • 非アクティブなキャッシュ済みメモリ：PyTorch が再利用できる inactive ブロック
   • 要求サイズと丸め後サイズ：requested_size とブロックの size の比較

診断する内容

メモリの使用先

使用中の active_allocated ブロックを、最も有用なユーザー側のフレームごとに分類します。最大のスタックトレースから順に、バイト数、使用中メモリに占める割合、ブロック数、代表的なファイル／関数／行番号を報告します。スタック名から判別できる場合は、パラメーター、オプティマイザーの状態、活性化、テンポラリーバッファー、通信バッファー、ユーザーコンテナーを分けて示します。

メモリの割り当て元は、アロケーターやフレームワークの内部ではなく、常にユーザーのモデルコードとして特定します。スナップショットのスタックは torch/nn/modules/*、torch/autograd/*、torch/optim/* 内のフレームが大半を占めます。これらをたどって、ユーザー自身のファイル内にある最初のフレームを見つけ、メモリを所有するモジュールまたは行を示してください（例：「使用中メモリの 38% は modeling.py:214 で割り当てられた model.decoder.layers[*].mlp の活性化です」）。スタック内にユーザー側のフレームがない場合に限り、torch 内部のフレームを使用してください。

メモリが解放されるタイミング

device_traces を使って、free_requested と free_completed を区別します。解放待ちまたは遅延解放は、別のストリームがまだテンソルを必要としているか、処理がまだ解放地点に到達していない可能性を示します。複数ステップのトレースでは、繰り返し現れる割り当て箇所を対応する解放イベントと比較してください。毎ステップ現れる一方で対応する解放がない割り当て箇所は、一度限りのモデルやオプティマイザーの割り当てよりも、保持の強い証拠になります。トレースに OOM イベントが含まれる場合は、失敗した割り当てのサイズ、device_free、OOM が発生したスタックフレーム、OOM の直前にメモリを保持していた使用中の割り当て箇所を報告してください。

リークとメモリ保持のパターン 手順をまたいで増加するアクティブなバイト数、解放されない同一スタックからのアクティブな割り当ての繰り返し、または多数のテンソルが蓄積して一斉に解放される鋸歯状のパターンを探してください。最も有力な証拠は、2 つのスナップショットの差分（ヘルパーの --diff）です。k ステップ離して取得したスナップショット間で、ライブバイト数とブロック数の両方が増加している割り当て箇所は、リークしているか保持されています。

次の名前付きパターンをスナップショットの特徴から特定し、該当する修正方法を提示してください。

• 保持された autograd グラフ：ステップ N のライブなアクティベーションブロックがステップ N+k にも残っており、多くの場合、ステップごとに損失テンソルまたは出力テンソルも存在します。典型的な原因は、グラフを保持しているテンソルをログに記録したり蓄積したりすることです。追加する前に loss.detach()（または .item()）を使って修正するか、グラフを保持する retain_graph=True を探してください。
• 増加する Python コンテナ：同じユーザーコードの割り当て箇所で、ブロック数が手順ごとに線形に増加します。追加する前にデタッチして CPU に移動するか、コンテナのサイズに上限を設けてください。
• 初回ステップでのオプティマイザー状態の急増：optimizer.step() の初回に torch/optim/* に起因する大きな割り当てが一度だけ発生し、その後は安定します。これは想定される保持であり、リークではないことを明示してください。この場合、OOM の解決策はモデルやバッチの縮小、メモリ使用量の少ないオプティマイザー、またはシャーディングであり、リークの調査ではありません。
• アクティベーションメモリとチェックポイント：ライブメモリの大半をフォワードのアクティベーションが占めており、これらはすべてバックワード中に解放されます（device_traces で鋸歯状になります）。アクティベーションチェックポイント（torch.utils.checkpoint）を推奨し、解放されるアクティベーションのバイト数と再計算コストのトレードオフをスナップショットから定量化してください。
• DDP と FSDP における勾配メモリ：各ランクでフルサイズの勾配ブロックとパラメータブロックが見られる場合は、DDP による複製が考えられます。ユーザーがメモリ不足に直面しており、勾配＋パラメータ＋オプティマイザー状態がライブメモリの大半を占めている場合は、構造的な解決策として FSDP／ZeRO 形式のシャーディングに言及してください。
• Python の参照サイクル：GC が実行されるときだけ、解放がまとまって発生します。CUDA テンソルを保持するオブジェクトサイクルを見つけるため、初期化時に from torch.utils.viz._cycles import warn_tensor_cycles をインポートし、warn_tensor_cycles() を呼び出すことを推奨してください。

断片化とアロケーターの挙動

nvidia-smi の使用量をライブテンソルと同一視しないでください。PyTorch はキャッシュ型の CUDA アロケーターを使用するため、予約済みメモリは割り当て済みテンソルメモリを大幅に上回ることがあります。大きな非アクティブブロックは PyTorch で再利用できます。empty_cache() は完全に非アクティブなセグメントを CUDA に返す場合がありますが、ライブテンソルが使用中のメモリは解放しません。

よくある誤解を招くケースは、十分な予約済みメモリがあるのに発生する OOM です。非アクティブメモリの合計が十分でも、要求されたサイズが最大の連続した非アクティブブロックを超えていると、割り当ては失敗します。スナップショットから診断してください。OOM 時の要求サイズと最大の非アクティブブロック（ヘルパーの「largest single reusable inactive block」）を比較し、非アクティブメモリが多数のセグメントに分割されていないか確認してください。このパターンに該当する場合は、次の順で推奨してください。

1. PYTORCH_ALLOC_CONF=expandable_segments:True — 新しい固定サイズのセグメントを割り当てる代わりに、セグメントを拡張できるようにします。通常、断片化による OOM に最も効果的な修正です。
2. PYTORCH_ALLOC_CONF=max_split_size_mb:<N> — N MiB を超えるブロックをアロケーターが分割しないようにし、再利用しにくい断片を減らします。N には失敗した要求サイズに近い値を提案してください。
3. ユーザーコードで割り当てサイズの変動を減らす（バッチサイズやシーケンスサイズの変動はよくある原因です。形状を固定するためのパディングが役立ちます）。

断片化が問題ではない場合は、その旨を明示してください。ライブテンソルメモリだけでデバイス容量に迫っているなら、アロケーターの調整では解決しません。モデルレベルでの対策（バッチサイズ、チェックポイント、シャーディング、精度の調整）が必要です。

PyTorch 以外のメモリ スナップショットで確認できるのは、PyTorch アロケーターが管理する CUDA メモリのみです。NCCL、カスタム CUDA 拡張、cuDNN/cuBLAS ワークスペース、または CUDA の直接割り当ては表示されないことがあります。スナップショットに含まれないデバイスメモリが疑われる場合は、torch.cuda.device_memory_used(device) または nvidia-smi の値と、スナップショットの予約済みバイト数を比較してください。

ユーザーがスナップショットの取得を必要としている場合

最小限の取得パターンを提示してください：

import torch

torch.cuda.memory._record_memory_history(max_entries=100000)
try:
    run_your_code()
    torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")
finally:
    torch.cuda.memory._record_memory_history(enabled=None)

長時間実行するジョブでは max_entries を大きくしてください。ただし、トレース履歴にはメモリのオーバーヘッドがあることを伝えてください。

OOM のデバッグでは、OOM オブザーバーを登録して、失敗した瞬間にスナップショットを自動的にダンプしてください。ユーザーがクラッシュを 2 回再現する必要がないようにします：

import torch


def dump_on_oom(device, alloc, device_alloc, device_free):
    torch.cuda.memory._dump_snapshot(f"oom_snapshot_dev{device}.pickle")


torch.cuda.memory._record_memory_history(max_entries=100000)
torch._C._cuda_attach_out_of_memory_observer(dump_on_oom)

トレーニングステップをまたぐリークが疑われる場合は、スナップショットを2 つ取得します。1 つは早い時点（オプティマイザーの状態がすでに割り当てられた、最初のオプティマイザーステップの後）、もう 1 つはかなり後のステップで取得し、analyze_snapshot.py earlier.pickle --diff later.pickle で差分を確認してください。同じステップ位置（たとえば optimizer.step() の直後）でダンプすると、ステップ内のノイズを差分から除外できます。

出力形式

主な調査結果を 1 文で最初に示してください。その後、以下を提示します：

1. 合計値：予約済み、ライブ割り当て済み、アクティブだが解放待ち、非アクティブなキャッシュ、およびトレース履歴があればピーク値。
2. ライブメモリ使用量の多い割り当て元の上位項目。スタックフレームと割合も示します。
3. 解放のタイムライン：完了した解放、保留中の解放、遅延した解放、および OOM イベント。
4. 診断：該当するパターン（保持されたグラフ、増大するコンテナー、オプティマイザー状態の急増、アクティベーションのメモリ圧迫、断片化／キャッシュ、参照サイクル、または PyTorch 外部メモリとの差）を特定し、それを裏付けるスナップショットの証拠を説明します。想定される保持とリークは明確に区別してください。
5. 具体的な次の手順：変更すべきユーザーコードの箇所（スナップショットのスタックにある file:line）、該当パターンに対する具体的な修正、断片化が関係する場合に試すアロケーター設定、そして診断の確認に差分が必要な場合に取得する次のスナップショットを示します。

不確実な点は明確にしてください。スナップショットに履歴がない場合は、そのデータからはどの解放やピークに関する疑問に答えられないのかを説明してください。