GPU・環境構築

【CUDA】GPUメモリ不足エラーの意外な原因！メモリフラグメンテーションの検出と効果的な解消法

· 14分で読めます · 5 閲覧

問題の概要：突然の「out of memory」エラーとパフォーマンス低下

深層学習モデルの学習や大規模なバッチ推論を実行中に、以下のようなエラーに遭遇したことはありませんか？

RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 23.69 GiB total capacity; 20.34 GiB already allocated; 0 bytes free; 22.00 GiB reserved in total by PyTorch)

一見すると、GPUメモリが物理的に不足しているように見えます。しかし、nvidia-smi コマンドで確認すると、利用可能な空きメモリが十分にあるにもかかわらず、このエラーが発生することがあります。

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 535.54.03    Driver Version: 535.54.03    CUDA Version: 12.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0  On |                  N/A |
| 30%   50C    P2    89W / 350W |   22000MiB / 24576MiB |     45%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

上記の例では、GPUメモリは約24.5GiB中22GiB使用されており、約2.5GiBの空きがあります。それでも「2.00 MiBの確保に失敗した」というエラーが発生するのです。この矛盾の背後にある主な原因の一つが、CUDAメモリフラグメンテーションです。これは、メモリの断片化により、連続した大きな空きメモリブロックが不足している状態を指します。

原因の解説：なぜフラグメンテーションが発生するのか？

CUDAメモリフラグメンテーションは、以下のような典型的な開発パターンによって引き起こされます。

1. メモリの頻繁な確保と解放

トレーニングループ内で可変サイズのテンソルを繰り返し作成・破棄したり、異なるサイズのバッチを処理したりすると、メモリ空間に「穴」が生じます。時間の経過とともに、これらの穴が散在し、合計の空き容量は十分でも、新しい大きなテンソルを配置できる連続した領域が見つからなくなります。

2. PyTorch/TensorFlowのメモリアロケータの挙動

これらのフレームワークはパフォーマンス向上のため、解放されたメモリを即座にGPUに返さず、内部のメモリプール（キャッシュ）に保持します。このキャッシュはサイズごとに管理されるため、特定のサイズのブロックが枯渇すると、より大きな連続ブロックを要求できなくなることがあります。

3. メモリの「予約」

CUDAコンテキストの初期化時や、フレームワークの起動時に、ドライバやランタイムが一定量のメモリを先行確保（予約）することがあります。この予約メモリと、アプリケーションが使用するメモリの配置が悪く、断片化を助長する場合があります。

解決方法：フラグメンテーションの検出と解消ステップ

ステップ1：フラグメンテーションの検出と確認

まず、本当にフラグメンテーションが問題なのかを確認します。PyTorchを使用している場合、以下のコードで詳細なメモリ統計を取得できます。

import torch
# 現在のメモリ統計を表示
print(torch.cuda.memory_summary())
# または、より簡潔なスナップショット
print(torch.cuda.memory_snapshot())

より視覚的で詳細な情報を得るには、pynvml ライブラリを使用する方法があります。

!pip install pynvml
import pynvml
pynvml.nvmlInit()
handle = pynvml.nvmlDeviceGetHandleByIndex(0) # GPU 0
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(f"Total memory: {info.total / 1024**3:.2f} GB")
print(f"Free memory: {info.free / 1024**3:.2f} GB")
print(f"Used memory: {info.used / 1024**3:.2f} GB")

根本的な原因を探るには、メモリの割り当て・解放のパターンを記録するメモリプロファイラを使用します。PyTorch 1.10以降では、以下のようにして有効化できます。

# メモリプロファイリングの開始
torch.cuda.memory._record_memory_history(max_entries=100000)
# ここに問題のコードを実行...
# プロファイリングの停止とレポート生成
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

生成された .pickle ファイルは、Chromeブラウザの chrome://tracing で開くことで、時間軸に沿ったメモリの使用状況を可視化できます。

ステップ2：即効性のある解消法 – GPUメモリのクリア

最も簡単な方法は、CUDAコンテキストをリセットし、すべてのキャッシュをクリアすることです。Jupyter Notebookや対話型環境で有効です。

import torch, gc
# 1. Pythonガベージコレクションの実行
gc.collect()
# 2. PyTorchのCUDAキャッシュをクリア
torch.cuda.empty_cache()
# 3. カーネルを再起動する（Jupyterの場合の最終手段）

これで一時的にフラグメンテーションは解消されますが、根本的なコードのパターンが変わらなければ、すぐに再発します。

ステップ3：根本的な解決 – コードと設定の最適化

以下のプラクティスをコードに取り入れることで、フラグメンテーションの発生を抑制できます。

1. 固定サイズのバッチとテンソルを使用する: 可能な限り、バッチサイズやテンソルサイズを固定し、可変長データの場合はパディングやマスキングを活用します。

# 悪い例：ループごとに異なるサイズのテンソルを作成
for data in dataloader:
    x = data.to(device) # dataのサイズが可変

# 良い例：最大長に合わせて固定サイズのバッチを準備
batch_size = 32
max_seq_len = 512
input_tensor = torch.zeros((batch_size, max_seq_len), device=device)

2. メモリプールの設定を調整する (PyTorch): 環境変数でアロケータの挙動を変更できます。

import os
# キャッシュされたメモリを積極的に解放するようにする（速度とトレードオフ）
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
# または、キャッシュ戦略を「常に拡張」から「ベストフィット」に変更（実験的）
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'backend:cudaMallocAsync'

3. In-place操作を活用する: 新しいメモリを割り当てず、既存のテンソルを書き換える操作を使用します。

# 新しいメモリを割り当てる
x = x + 1
# In-place操作でメモリ割り当てを防ぐ
x.add_(1)

4. 不要なテンソルの参照を早期に解放する: 大きなテンソルがスコープを外れても、他の変数が参照していると解放されません。

# 中間テンソルを明示的にNoneにする
intermediate = heavy_operation(x)
result = final_operation(intermediate)
intermediate = None # 参照を解除し、GCの対象に
gc.collect()

ステップ4：環境レベルの対策

アプリケーション起動時に、GPUメモリの大きなブロックを事前に確保（「占拠」）して、断片化が起こりにくい状態を作る方法もあります。

import torch
# 起動直後に大きな連続メモリブロックを確保
large_block = torch.cuda.caching_allocator_alloc(0.9 * torch.cuda.get_device_properties(0).total_memory)
# 必要に応じて手動で解放可能
# torch.cuda.caching_allocator_delete(large_block)

また、Dockerコンテナを使用している場合は、--gpus all ではなく特定のGPUを指定したり、CUDA Visible Devicesを設定することで、他のプロセスによる干渉を減らせます。

# Docker runの例
docker run --gpus '"device=0"' your_image
# またはコード内で
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

まとめ・補足情報

CUDAメモリフラグメンテーションは、GPUメモリの物理的な不足ではなく、その「使い方」に起因する問題です。症状は「out of memory」エラーとして現れますが、nvidia-smi では空き容量が確認できるという点が特徴です。

対処法は以下の流れで行うと効果的です。

  1. 確認: torch.cuda.memory_summary() やプロファイラでフラグメンテーションを疑う。
  2. 応急処置: torch.cuda.empty_cache() でキャッシュをクリア。
  3. 根本対策: コードを見直し、可変サイズテンソルの作成を減らし、In-place操作や固定サイズバッチを採用する。
  4. 環境調整: メモリアロケータの設定（PYTORCH_CUDA_ALLOC_CONF）やGPUの隔離を検討する。

最後に、本当にモデルやバッチサイズがGPUメモリ容量に対して大きすぎる場合もあるため、モデルの軽量化（量子化、プルーニング）、勾配累積、CPUオフローディングなどの他のメモリ最適化技術と組み合わせて検討することが、持続可能な開発につながります。フラグメンテーション対策は、GPUという貴重なリソースを効率的に長く使い続けるための重要なスキルです。

この記事は役に立ちましたか？