深層学習の開発や実験をしていると、誰もが一度は遭遇する「CUDA out of memory」エラー。特に、「GPUの空きメモリは十分にあるのに、なぜかエラーが出る」という不可解な状況に頭を抱えた経験はありませんか?
この記事では、その一見矛盾したエラーの背後に潜む「CUDAメモリフラグメンテーション(断片化)」について、その原因から検出方法、そして根本的な解決策までを、実際のコード例と共に詳しく解説します。
はじめに:問題の本質は「連続した空き領域」の不足
まずは典型的なエラーメッセージを見てみましょう。
RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 10.00 GiB total capacity; 5.50 GiB already allocated; 0 bytes free; 5.80 GiB reserved in total by PyTorch)
このエラーは、「合計2GiBのメモリを確保しようとしたが失敗した」と報告しています。しかし、nvidia-smiコマンドで確認すると、利用可能な空きメモリが2GiB以上あることがしばしばあります。
この矛盾を生むのがメモリフラグメンテーションです。これは、メモリの確保と解放を繰り返す過程で、空きメモリが小さな断片として散らばってしまう現象です。結果として、合計の空き容量は足りていても、要求されたサイズの「連続した」メモリブロックが見つからず、メモリ確保に失敗するのです。
なぜフラグメンテーションが発生するのか? 〜PyTorchアロケータの挙動〜
フラグメンテーションを理解するには、PyTorch(やTensorFlow)がどのようにGPUメモリを管理しているかを知る必要があります。
これらのフレームワークは、パフォーマンス向上のため、独自のメモリアロケータを使用しています。このアロケータは、一度確保したメモリをすぐにGPUドライバに返さず、内部の「メモリプール(キャッシュ)」に保持します。次回同じサイズ付近のメモリ要求があった時に、プールから素早く割り当てるためです。
しかし、この最適化が仇となる場合があります。以下のようなコードパターンは、フラグメンテーションを引き起こす典型例です。
import torch
# フラグメンテーションを引き起こしやすい悪い例
for i in range(1000):
# 毎回異なるサイズのテンソルを生成
size = torch.randint(100, 1000, (1,)).item() # 100〜1000の乱数
x = torch.randn(size, size, device='cuda') # サイズが可変
y = x * 2 # 何らかの計算
# ループを抜けるとxはスコープ外になるが、メモリはアロケータのキャッシュに残る
このループでは、毎回異なるサイズのメモリブロックが要求され、解放されます。アロケータは様々なサイズのメモリ断片をキャッシュに保持するため、時間と共にメモリ空間が「穴だらけ」の状態になります。
フラグメンテーションを悪化させる実践的な要因
- 可変サイズバッチ処理: 画像の解像度や文長がサンプルごとに異なるデータセットを扱う時、バッチ内のテンソルサイズが可変になりがちです。
- モデルの分割実行: 非常に大きなモデルを部分ごとにGPUにロードして実行する場合、メモリの確保・解放パターンが複雑化します。
- メモリプールの永続化: キャッシュされたメモリ断片はスクリプト実行中ずっと保持され続けるため、断片化した状態が固定化されてしまいます。
実践的解決ガイド:検出 → 応急処置 → 根本治療
Step 1: フラグメンテーションを検出する
まず、本当にフラグメンテーションが起きているのかを確認しましょう。PyTorchでは以下の方法で詳細なメモリ状態を確認できます。
方法A: PyTorchのメモリサマリーを表示
import torch
# 詳細なメモリ統計を表示(非常に情報量が多い)
print(torch.cuda.memory_summary(device=None, abbreviated=False))
出力のAllocated memory(割り当て済み)とFree memory(空き)の関係を注視します。Freeが十分にあるのに大きな確保に失敗する場合は、フラグメンテーションが強く疑われます。
方法B: シンプルなメモリ情報の取得
import torch
def print_memory_stats():
allocated = torch.cuda.memory_allocated(0) / 1024**3 # GiB
reserved = torch.cuda.memory_reserved(0) / 1024**3 # GiB
# reserved - allocated が、キャッシュされているが未使用の断片化メモリ
print(f"Allocated: {allocated:.2f} GiB")
print(f"Reserved: {reserved:.2f} GiB")
print(f"Fragmented (approx): {reserved - allocated:.2f} GiB")
方法C: nvidia-smiによる監視
ターミナルで以下のコマンドを実行し、空きメモリの変動を観察します。
# 1秒間隔で監視
watch -n 1 nvidia-smi
Step 2: 応急処置 - 実行中のメモリをクリアする
実験中や推論中に突然エラーが出た場合の即効性のある対処法です。
import torch
import gc
def clear_cuda_memory():
"""CUDAメモリのキャッシュを強制クリアする"""
gc.collect() # Pythonのガベージコレクションを明示的に実行
torch.cuda.empty_cache() # PyTorchのCUDAキャッシュを全て解放
print("CUDA cache cleared.")
# 大きなメモリ確保が必要な処理の直前に実行
clear_cuda_memory()
large_tensor = torch.randn(10000, 10000, device='cuda') # 大きな確保
注意点: torch.cuda.empty_cache()は、カーネルによって使用中のメモリは解放しません。あくまで「未使用でキャッシュされている」メモリ断片を解放するものです。また、頻繁に呼び出すとパフォーマンス低下を招くため、あくまで応急処置として使いましょう。
Step 3: 根本的解決 - コードと環境を改善する
応急処置だけでは根本解決になりません。以下の方法で、フラグメンテーションの発生自体を抑制しましょう。
対策1: メモリ確保パターンの最適化
可変サイズバッチの固定化
# 改善前: 可変サイズ(フラグメンテーションの原因)
for images in dataloader: # imagesのサイズがバッチごとに異なる
output = model(images.to('cuda'))
# 改善後: 固定サイズ(パディングなどでサイズを統一)
from torch.nn.utils.rnn import pad_sequence
# または画像の場合は、リサイズ/クロップでバッチ内の画像サイズを統一する
メモリのプリアロケーション(事前確保)
# 改善前: ループ内で都度確保
results = []
for data in dataloader:
output = model(data.to('cuda')) # 毎回新しいメモリ確保
results.append(output.cpu())
# 改善後: 最大必要サイズを事前に確保
batch_size = dataloader.batch_size
feature_dim = 1000 # モデルの出力次元数
# バッチ全体の出力を格納できるメモリを最初に確保
preallocated_output = torch.empty((batch_size, feature_dim), device='cuda')
for i, data in enumerate(dataloader):
# 事前確保したメモリの一部を使用
actual_batch_size = data.size(0)
output_slice = model(data.to('cuda'))
preallocated_output[:actual_batch_size].copy_(output_slice)
results.append(preallocated_output[:actual_batch_size].cpu())
対策2: PyTorchアロケータの挙動をチューニングする
PyTorchには、フラグメンテーションを軽減するための実験的なアロケータ設定があります。
方法A: 環境変数で設定(推奨)
スクリプトを実行する前に、ターミナルで設定します。
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
python your_script.py
方法B: コード内で設定
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
# 注意: この設定はtorchをimportする前に行う必要があります
import torch
max_split_size_mbのチューニングについて
このパラメータは、アロケータがメモリブロックを分割する最大サイズを指定します。
- 値を小さくする(例: 32): メモリの分割が細かくなり、断片化しにくくなるが、アロケータの管理オーバーヘッドが増える。
- 値を大きくする(例: 512): オーバーヘッドは減るが、大きな断片が生まれやすく、逆にフラグメンテーションが起こりやすくなる。
- デフォルト: 非常に大きい値(実質無制限)に設定されていることが多い。
ベストな値はワークロードに依存します。128や256から始めて、メモリ使用効率と速度のバランスを見ながら調整するのが良いでしょう。
対策3: プログラミングプラクティスの見直し
-
テンソルのin-place操作を活用する: 新しいメモリを確保せず、既存のテンソルを書き換えます。
# 改善前 x = x * 2 # 新しいメモリが確保される # 改善後 x.mul_(2) # in-place操作。メモリ確保が発生しない注意: in-place操作は勾配計算グラフに影響を与える可能性があるため、学習時には注意が必要です。
-
不要なメモリ保持を避ける: 中間テンソルを不必要に保持しない。
# 改善前 intermediate = layer1(x) result = layer2(intermediate) # intermediateがスコープ内に残り続ける # 改善後(スコープを切る) result = layer2(layer1(x)) # または明示的に削除 del intermediate
まとめ:予防と早期発見が鍵
CUDAメモリフラグメンテーションは、深層学習開発において見落とされがちなパフォーマンスと安定性の課題です。
-
予防: 固定サイズバッチの採用、メモリプリアロケーションの検討、適切なアロケータ設定(
PYTORCH_CUDA_ALLOC_CONF)が有効です。 -
早期発見:
torch.cuda.memory_summary()や定期的なメモリ監視で、フラグメンテーションの兆候をキャッチしましょう。 -
応急処置:
torch.cuda.empty_cache()は便利ですが、根本解決にはなりません。あくまで一時的な手段として使いましょう。
メモリ管理への意識的な取り組みは、大規模モデルの学習やリソース制約のある環境での開発を成功させる重要な要素です。本記事の手法を参考に、より安定で効率的なGPUプログラミングを実現してください。
この記事の詳細版は元記事サイトで公開しています。