問題の概要:突然の「out of memory」エラーとパフォーマンス低下
深層学習モデルの学習や大規模なバッチ推論を実行中、以下のようなエラーに遭遇したことはありませんか?
RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 24.00 GiB total capacity; 18.50 GiB already allocated; 0 bytes free; 19.20 GiB reserved in total by PyTorch)
一見、単純なメモリ不足に見えますが、GPUの全メモリ容量(例:24GB)に対して、確保済みメモリ(18.5GB)と空きメモリ(0バイト)の合計が、PyTorchが「予約済み」と報告する総メモリ(19.2GB)よりも少ないという矛盾が生じています。この「行方不明のメモリ」の正体が、CUDAメモリフラグメンテーションです。
この現象が発生すると、以下の症状が見られます。
- 理論上は足りているはずのメモリ容量で「out of memory」エラーが発生する。
- モデル学習中、エポックやバッチによってメモリ使用量が不安定に変動する。
- 長時間実行を続けると、徐々にパフォーマンスが低下する。
- 小さなテンソルの割り当てに予想外に時間がかかる。
原因の解説:GPUメモリの「断片化」とは?
CUDAメモリフラグメンテーションは、メインメモリ(RAM)の断片化と概念的に似ています。CUDAランタイム(PyTorchやTensorFlowを通じて)がGPUメモリを動的に確保・解放する過程で発生します。
フラグメンテーションが発生する仕組み
1. 様々なサイズのメモリブロックの割り当てと解放: 学習プロセスでは、大きな重みテンソル、中間活性化、勾配、オプティマイザの状態など、様々なサイズのメモリブロックが絶え間なく確保・解放されます。
2. 隙間(フラグメント)の発生: 解放されたメモリブロックのサイズと、次に要求されるメモリブロックのサイズが一致しない場合、解放された領域を再利用できません。例えば、1GBと2GBのブロックが解放された後、2.5GBの連続したメモリを要求されると、合計3GBの空き領域があるにもかかわらず、それが2つの断片に分かれているため、要求を満たせず「メモリ不足」エラーとなります。
3. キャッシュによる悪化: PyTorchのCUDAメモリアロケータはパフォーマンス向上のため、解放されたメモリブロックを一定期間プール(キャッシュ)します。これは小さな割り当てを高速化しますが、キャッシュされたブロックが再利用されないまま長時間滞留すると、実質的に使用不能な「予約済みメモリ」として存在し続け、断片化を加速させます。
この状態を図で表すと以下のようになります。
GPUメモリレイアウト(理想):
[ 空き: 5GB ]
GPUメモリレイアウト(断片化後):
[ 使用中1GB | 空き1GB | 使用中2GB | 空き2GB | 使用中... ]
↑ ↑ ↑
| | この2GBの空きブロックはあるが...
| この1GBの空きブロックはあるが...
連続した3GBの空き領域がない!
→ 3GBのテンソルを割り当てられず「out of memory」
解決方法:フラグメンテーションの検出と解消ステップ
ステップ1:フラグメンテーションの検出と状況確認
まず、PyTorchの組み込み関数を使用して、メモリの詳細なスナップショットを取得します。
import torch
# 現在のメモリ統計を表示
print(torch.cuda.memory_summary(device=None, abbreviated=False))
# より簡潔なスナップショット
print(torch.cuda.memory_snapshot())
# 割り当て済みメモリとキャッシュメモリを個別に確認
allocated = torch.cuda.memory_allocated(0) / 1024**3 # GB単位
cached = torch.cuda.memory_reserved(0) / 1024**3 # GB単位
print(f"Allocated: {allocated:.2f} GB")
print(f"Cached (Reserved): {cached:.2f} GB")
print(f"Potential Fragmentation (Cached - Allocated): {cached - allocated:.2f} GB")
memory_summaryの出力の「Active Memory」と「GPU reserved Memory」の差が大きい場合、またはキャッシュメモリが異常に大きい場合、フラグメンテーションが発生している可能性が高いです。
ステップ2:即時解消法 – CUDAメモリのキャッシュクリア
最も簡単な解消法は、PyTorchが保持しているCUDAメモリのキャッシュを空にすることです。これにより、すべての「予約済み」メモリが解放され、フラグメントが解消されます。
import torch
import gc
# Pythonのガベージコレクションを実行(参照カウント0のオブジェクトを解放)
gc.collect()
# PyTorchのCUDAメモリキャッシュを完全に解放
torch.cuda.empty_cache()
# 解放後のメモリ状態を確認
allocated_after = torch.cuda.memory_allocated(0) / 1024**3
cached_after = torch.cuda.memory_reserved(0) / 1024**3
print(f"After empty_cache -> Allocated: {allocated_after:.2f} GB, Cached: {cached_after:.2f} GB")
注意点: empty_cache()は、現在計算グラフで使用中のテンソル(例えば、モデルパラメータ、オプティマイザの状態)は解放しません。あくまで「未使用のキャッシュ」を解放します。スクリプトの区切り目(エポック間、異なるモデルのロード前など)で実行するのが効果的です。
ステップ3:根本的対策 – メモリアロケーションファクトリの設定(PyTorch 1.10以降)
PyTorch 1.10以降では、より積極的にフラグメンテーションを抑制するアロケータを設定できます。
import torch
# 環境変数で設定する方法(スクリプト実行前)
# ターミナルで: export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
# Pythonスクリプト内で設定する方法(最初のCUDA呼び出し前に行う)
# この設定は、大きなメモリブロックを分割する際の最大サイズを制限し、
# 小さな断片が大量に発生するのを防ぎます。
torch.cuda.set_per_process_memory_fraction(0.9) # 必要に応じて
# アロケータ設定はスクリプト起動時の環境変数での設定が推奨されます。
主要な環境変数設定:
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128: 分割ブロックの最大サイズを制限。PYTORCH_CUDA_ALLOC_CONF=backend:native(PyTorch 2.0以降): 場合によってはパフォーマンスが向上。PYTORCH_CUDA_ALLOC_CONF=garbage_collection_threshold:0.8: ガベージコレクションの頻度を上げる。
ステップ4:予防的コーディングプラクティス
コードレベルでフラグメンテーションの発生を抑制します。
# 1. テンソルをインプレース操作する
x = torch.randn(1024, 1024, device='cuda')
# 非推奨: x = x * 2 (新しいメモリを割り当てる)
# 推奨: x.mul_(2) (インプレース操作でメモリ割り当てを防ぐ)
# 2. 不要なテンソル参照を早期に解放
def process_data(data):
intermediate = heavy_operation(data) # 大きな中間テンソル
result = final_operation(intermediate)
del intermediate # 明示的に参照を削除
torch.cuda.empty_cache() # 必要に応じて
return result
# 3. データローダのワーカー数とピンメモリのバランスを取る
# ワーカー数が多いとCPU側のメモリ圧迫がGPUメモリ確保に影響する場合がある。
dataloader = DataLoader(dataset, batch_size=32, shuffle=True,
num_workers=4, pin_memory=True) # num_workersは調整を
# 4. 大きなメモリブロックを最初に確保する(可能な場合)
# モデル、オプティマイザを最初にGPUに移し、大きなワークスペースを確保してから学習を開始する。
ステップ5:高度なツールによるモニタリング
NVIDIAが提供するnvtop(Linux)やnvidia-smiの詳細モードで、メモリ使用の変動をリアルタイムで観察できます。
# nvidia-smi の定期的な監視(1秒間隔)
nvidia-smi -l 1
# より詳細なメモリイベントのトレース(デバッグ用)
# 環境変数を設定してPyTorchを起動
# PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,record_context:True
# 実行後、メモリの履歴を分析可能
まとめ・補足情報
CUDAメモリフラグメンテーションは、特に長時間の学習ジョブや、メモリ確保/解放パターンが複雑なマルチタスク環境で顕在化する「静かな性能劣化」の原因です。
要点のまとめ:
- 検出:
torch.cuda.memory_summary()で「Reserved」と「Allocated」の差を確認する。 - 即時解決:
gc.collect()とtorch.cuda.empty_cache()の組み合わせでキャッシュをクリア。 - 根本対策:
PYTORCH_CUDA_ALLOC_CONF環境変数でメモリアロケータの挙動を最適化する。 - 予防: インプレース操作、明示的な参照削除(
del)、メモリ確保パターンの最適化を心がける。
補足:
- TensorFlowでは、
tf.config.experimental.set_memory_growthを使用してメモリの増分確保を有効にすることで、ある程度フラグメンテーションを緩和できます。 - どうしても解決しない場合、最終手段としてプロセス自体を再起動するのが最も確実な方法です。これは、すべてのメモリ状態が初期化されるためです。
- PyTorchのバージョンアップに伴い、CUDAアロケータは継続的に改良されています。最新の安定版を使用することも問題解決の一助となります。
メモリフラグメンテーションへの理解と適切な対処法を身につけることで、貴重なGPUリソースを最大限に活用し、安定したAI開発を行うことができるでしょう。