PyTorchのCUDA out of memory対処法|GPUメモリ不足の原因を切り分ける

PyTorchのCUDA out of memoryを原因別に確認する

PyTorchのCUDA out of memoryを原因別に確認する

PyTorchでCUDA out of memoryが出たら、エラーが出た処理と、使っているGPUの空き容量を確認し、バッチサイズや入力サイズを小さくして再実行するところから始めます。何回か動いた後で止まる場合は、テンソル(多次元配列)や計算履歴を残していないかも調べます。

torch.cuda.empty_cache()は、使用中のテンソルを消す操作ではありません。これを追加するだけで、モデルや入力に必要なGPUメモリが減るわけではないため、原因に合う対策を選びます。

この記事は、NVIDIA GPUを使うPyTorchの学習・推論を対象にします。CUDA(NVIDIA GPUで計算するための仕組み)が認識されていない場合は、先にtorch.cuda.is_available()がFalseになる原因と確認手順を確認してください。

参照資料は2026年9月11日に確認したPyTorch 2.14の公式ドキュメントです。GPUでのOOM再現や削減率は今回測定していません。CPUで確認した挙動と、公式資料に基づくGPU向け手順を区別して説明します。

どの処理でメモリ不足になったかを記録する

エラー全文を保存し、モデル、入力の形状、バッチサイズ、データ型、PyTorchの版を残します。Tried to allocateに出る容量は、その時点で追加確保しようとした量です。モデル全体の必要容量と同じではありません。

VRAM(GPUのメモリ)は、モデルの重みだけでなく、学習中の中間結果、勾配、最適化アルゴリズムの状態などにも使われます。推論でも入力や作業領域が必要なので、モデルファイルのサイズだけで動作可否は決められません。

失敗するタイミング 最初に確認すること 対処の方向
モデルをGPUへ移すとき 他プロセス、重みの容量、同じモデルの重複 不要な自分の処理を終了、小さいモデルや対応する軽量化方式を検討
最初のforward(順伝播) バッチ、画像解像度、系列長 同時処理量や入力サイズを減らす
backward(逆伝播) 学習用の中間結果や勾配 バッチ縮小、AMP、checkpointを検討
最初のoptimizer.step 最適化用の状態の追加確保 optimizerを含む学習全体の必要量を見直す
反復するほど使用量が増える Tensorの保存、計算履歴、モデルの作り直し 参照を残す場所を探す
大きい入力や形状変更時だけ失敗 入力の最大値、確保サイズの変動 入力上限やバッチ分けを調整し、必要ならallocatorを調査

これは原因を絞るための表です。CUDA処理は非同期に進むため、表示されたPython行だけを根拠に原因を断定しないでください。

他プロセスの使用量とPyTorch内の使用量を分ける

ターミナルで、GPU全体の状態を確認します。

nvidia-smi

対象GPUに別の学習処理やNotebookが残っていないか確認します。表示できるプロセス情報はOSや実行環境に依存します。自分の不要な処理を正常終了することから始め、他ユーザーの処理を停止しないでください。

続けて、OOMが起きるPythonプロセス内でメモリを記録します。新しいPythonを別に起動しても、元のプロセスのTensor使用量は測れません。

以下はCUDAが利用可能な環境向けの診断例です。

import logging
import torch

logger: logging.Logger = logging.getLogger("cuda_memory")
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")


def log_cuda_memory(label: str, device: int = 0) -> None:
    """指定GPUのメモリ指標をログへ出す。

    Args:
        label: 比較する処理地点の名前。
        device: プロセスから見えるCUDAデバイス番号。
    Returns:
        None。
    Raises:
        RuntimeError: CUDAが利用不可、またはCUDA処理が失敗した場合。
    Example:
        log_cuda_memory("after_forward", device=0)
    """
    if not torch.cuda.is_available():
        raise RuntimeError("CUDAが利用できる環境で実行してください。")

    # 診断時だけ同期し、処理地点ごとの値を比較しやすくする。
    torch.cuda.synchronize(device)
    free_bytes: int
    total_bytes: int
    free_bytes, total_bytes = torch.cuda.mem_get_info(device)
    gib: int = 1024 ** 3
    logger.info(
        "%s device=%d allocated=%.3f GiB reserved=%.3f GiB "
        "peak_allocated=%.3f GiB free=%.3f GiB total=%.3f GiB",
        label,
        device,
        torch.cuda.memory_allocated(device) / gib,
        torch.cuda.memory_reserved(device) / gib,
        torch.cuda.max_memory_allocated(device) / gib,
        free_bytes / gib,
        total_bytes / gib,
    )


log_cuda_memory("before_work")

モデルの読み込み後、forward後、backward後などへ呼び出しを置き、同じ処理地点の値を比較します。 実際の処理中にOOMになった場合は、例外で残った参照を保持したまま何度も再試行せず、実行状態を整理してから測り直します。device=0はプロセスから見た番号で、GPUの可視範囲を変更している場合はnvidia-smiの番号と単純に対応しないことがあります。

指標 対象 読み方
memory_allocated() このプロセスのTensorが占めるメモリ Tensor保持や実処理量を調べる
memory_reserved() このプロセスのキャッシュ用allocatorが管理するメモリ 使用中の領域と再利用用の領域を含む
max_memory_allocated() 記録区間内のTensor使用量の最大値 処理中の一時的な増加を比べる
mem_get_info()のfree / total 指定GPU全体 他プロセスも含めた空きと総量を見る

定義は公式のmemory_allocatedmemory_reservedmax_memory_allocatedmem_get_infoで確認できます。

allocatedreservedに含まれるので、2つを足して使用量としません。また、nvidia-smiにはCUDAの管理領域や他の割当も含まれるため、PyTorch内の値と一致するとは限りません。PyTorchのメモリ可視化でも、allocatorの外側で確保されたメモリには見えないものがあります。Understanding CUDA Memory Usageがこの範囲を説明しています。

allocatedはreservedの内側にあり、GPU全体の使用量とは範囲が異なる

図:公式のメモリ指標を基に独自作成した概念図。面積は実測値ではなく、未使用領域がすべて直ちに返却できることも意味しません。

開始直後に足りないなら、処理量を減らす

バッチサイズと入力サイズを小さくする

バッチサイズ(一度に処理するサンプル数)を、たとえば8から4、4から2へ減らして比較します。これは調整例であり、特定の値なら必ず動くという意味ではありません。

画像なら解像度、テキストなら系列長、生成処理なら同時生成数なども確認します。平均的な入力では動いても、最大サイズの入力だけでOOMになる場合があります。

バッチを小さくしても、モデルの重みやoptimizerの状態は同じ量のまま残ることがあります。バッチ1でもモデルの転送時に失敗するなら、バッチ縮小だけで解決する余地は小さくなります。PyTorch公式FAQも、バッチ削減と不要な参照の見直しを案内しています。

推論ではevalと勾配記録の無効化を併用する

学習しない処理では、計算履歴を作らないようにします。次は、modelinputsを事前に同じGPUへ配置した推論処理の抜粋です。

model.eval()

# 後でbackwardしない推論なので、勾配用の履歴を作らない。
with torch.inference_mode():
    outputs: torch.Tensor = model(inputs)

model.eval()はDropoutやBatchNormなどを評価時の動作に切り替える操作です。それだけでは勾配計算用の記録を止めません。逆に、inference_mode()だけでもモデルは自動的にevalへ切り替わりません。inference_modeの公式説明に両者の役割が記載されています。

このモードで作ったTensorを後からautograd(自動微分)の計算へ渡す場合には制約があるため、学習の一部を切り出す用途ではno_grad()なども含めて設計を確認します。推論モードでも、重みや出力そのものの容量は必要です。

反復するほど増えるなら、保存しているTensorを調べる

ログや可視化用のリストにTensorを追加し続けると、必要以上にデータや計算履歴を保持することがあります。次はlossがスカラーTensorである学習ループに追加する、ログ保存部分の例です。

loss_history: list[float] = []

# 各反復で実行する。ログにはTensorではなく数値だけを残す。
loss_history.append(loss.item())

リストの初期化はループの前、appendは各反復へ置きます。loss.item()はPythonの数値を返すので、ログを通じてTensorへの参照を残すことを避けられます。GPUから値を取り出す際には同期のコストがあるため、必要ならログの頻度も調整します。

出力Tensorを後で使いたいなら、目的に合わせて保存方法を変えます。

保存方法 残るもの 注意点
outputsをそのまま保持 データと、状況によって計算履歴 不要な履歴を保持し得る
outputs.detach() 履歴を切ったTensor GPU上のデータ領域は残る
outputs.detach().cpu() CPU側のTensor CPUメモリを使い、転送コストもある
loss.item() スカラーのPython数値 多要素Tensorにはそのまま使えない

detach()は履歴を切り離しますが、元のTensorとストレージ(データの保存領域)を共有します。したがって、detach()した出力をGPU上にため続ければ、そのデータ量は増えます。Tensor.detachの公式説明を参照してください。

不要なretain_graph=True、毎反復でのモデル生成、クラスの属性や辞書への出力保存も確認します。del outputsで変数名を消しても、別のリストなどから同じTensorを参照していれば、それだけでは解放されません。

Notebookでは、以前のセルの変数や例外の履歴が残る場合もあります。保存が必要な結果を退避したうえでカーネルを再起動し、最小限のセルから再現するか確認すると、持ち越した状態を切り分けられます。

detachしたTensorが同じGPUデータ領域を参照し続ける例

図:Tensor.detachのストレージ共有を基に独自作成。参照関係を示す図で、GPUの測定結果ではありません。

学習時のピークを下げる対策を比較する

バッチや入力を減らしても必要な学習条件を満たせない場合は、次の手段を検討します。

対策 減らす対象・仕組み トレードオフ
AMP(自動混合精度) 一部演算を低精度で実行する 数値安定性・対応演算・GPUの確認が必要
勾配蓄積 小さいバッチごとに逆伝播し、複数回分をためて更新する 更新1回あたりの時間が伸びる。重みやoptimizer状態は減らない
Activation checkpointing 一部中間結果を保存せず、逆伝播時に再計算する メモリと引き換えに再計算が増える
小さいモデルや対応する量子化方式 重みなどの容量を減らす 品質・対応機能・学習可能範囲が変わり得る

AMPでは、公式例にあるtorch.autocasttorch.amp.GradScalerを参照し、使用するデータ型と学習方法に合わせます。FP16とBF16では扱いが異なり、すべての重み・状態・作業領域が一律に半分になるとは限りません。

勾配蓄積を使うときは、各小バッチでbackwardしてから次へ進みます。複数のforwardの結果を全部保持して最後にまとめてbackwardすると、中間結果をため込み、期待した削減につながらない場合があります。lossの重み付け、最後の端数バッチ、optimizerの更新頻度も調整が必要です。BatchNormなどがあるため、大バッチを一度に処理した結果と常に同じになるわけでもありません。

実装の基準はAutomatic Mixed Precision examples、再計算を使う場合はtorch.utils.checkpointを参照してください。checkpointは公式が推奨するuse_reentrant=Falseを含め、利用版の条件を確認します。

empty_cacheで解放できるのは未使用のキャッシュ

PyTorchは確保済みのGPUメモリを再利用するため、不要になったTensorの領域をキャッシュとして保持することがあります。

import torch

# 不要なTensorへの参照を取り除いた後、未使用キャッシュを返したい場合。
torch.cuda.empty_cache()

この操作は、使用中のTensorや別プロセスのメモリを削除しません。キャッシュはもともとPyTorchが再利用できる領域なので、呼び出せば処理に使える容量がそのまま増えるとも限りません。empty_cacheの公式説明を確認してください。

毎反復で呼ぶ方法を基本対策にするより、まず実際の使用量と参照保持を調べます。同じモデル・入力を載せるだけで容量を超えるなら、キャッシュの整理より必要量の削減が先です。

reservedが大きいだけで断片化と決めない

断片化とは、空き領域の分かれ方などによって、要求された大きさの領域をうまく再利用できない状態です。ただし、reserved - allocatedが大きいことはキャッシュの存在も示すため、その差分だけで断片化が原因だとは言えません。

必要なら、対象プロセス内で次を確認します。

import torch

print(torch.cuda.memory_summary(device=0, abbreviated=False))

さらに調査する場合は、公式のCUDAメモリ可視化にある履歴・snapshotを利用します。取得方法や利用可能な指標は、PyTorchの版やallocatorのbackend(割当処理の実装)に依存します。

環境変数の設定は、その後に検討します。PyTorch 2.14のCUDA semanticsでは、次の条件が示されています。

設定 想定する状況 注意点
expandable_segments:True バッチなどの確保サイズが変動する処理 実験的機能。必要容量そのものを減らす設定ではない
max_split_size_mb 大量のinactive split blocksがありOOMになる処理 nativeでの最終手段。性能コストがあり、cudaMallocAsyncでは無視される

現行資料では環境変数名はPYTORCH_ALLOC_CONFで、PYTORCH_CUDA_ALLOC_CONFは後方互換の別名です。古いPyTorchへ新しい設定をそのまま適用せず、使用版の資料を確認します。

設定を試すなら新しいPythonプロセスの起動前に1項目だけ適用し、同じ入力で使用量と実行時間を比較します。Notebookでは、起動済みカーネル内で変数を書き換えただけで反映されたと判断しないでください。効果がなければ設定を外して起動し直し、原因の切り分けへ戻ります。

変更前後は同じ条件で確認する

一度動いたことに加えて、失敗した入力サイズと反復回数を再現できるか確認します。確認中は次の記録をそろえると比較しやすくなります。

  • モデル・バッチ・入力形状・データ型・PyTorchの版
  • 対象GPUと他プロセスの使用状況
  • 同じ地点のallocatedreservedと、処理中のピーク
  • 実行時間、出力の妥当性、学習ならlossの推移

ピークの比較区間を区切るには、対象処理の前でtorch.cuda.reset_peak_memory_stats(device)を使い、終了時のmax_memory_allocated(device)を読みます。このリセットは統計値を区切る操作で、Tensorを解放しません。診断用の同期やログを常時入れると実行時間へ影響するため、計測条件もそろえます。

今回確認した範囲

CPU版PyTorch 2.6.0、Python 3.12.3で、小さなLinearモデルを使い、次の挙動を確認しました。

確認項目 実測結果
eval()だけでforward 出力のrequires_gradTrue
inference_mode()内でforward 出力のrequires_gradFalse
detach()後のデータ領域 元の出力と同じアドレス
スカラーlossのitem() Pythonのfloat

これは勾配記録や参照の性質の確認です。CUDA上のメモリ使用量、OOMの解消、AMPやcheckpointによる削減率は実測していません。GPU向けの診断コードと対策は、利用環境で確認するための手順として掲載しています。

Notebookで別のPython環境を選んでいた場合は、Jupyterでインストール済みなのにModuleNotFoundErrorになる原因と対処法も参考になります。

コメント

タイトルとURLをコピーしました