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

問題の概要：突然の「out of memory」エラーとCUDAメモリフラグメンテーション

PyTorchやTensorFlowを使用した深層学習のトレーニングや推論中、以下のようなエラーに遭遇したことはありませんか？

RuntimeError: CUDA out of memory. Tried to allocate 1.50 GiB (GPU 0; 23.65 GiB total capacity; 20.34 GiB already allocated; 0 bytes free; 21.12 GiB reserved in total by PyTorch)

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

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.0    |
|-------------------------------+----------------------+----------------------+
| 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 ...  Off  | 00000000:01:00.0  On |                  N/A |
| 30%   50C    P2   120W / 250W |   18000MiB / 24576MiB |     50%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

上記の例では、約6GBの空きメモリがあるように見えますが、アプリケーションは1.5GBの確保に失敗しています。この矛盾の背後にある主な原因の一つが「CUDAメモリフラグメンテーション」です。本記事では、このメモリフラグメンテーションの仕組み、検出方法、そして効果的な解消手法について詳しく解説します。

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

CUDAメモリフラグメンテーションは、物理的なハードディスクのフラグメンテーションと概念的に似ています。GPU上でTensor（テンソル）などのメモリ領域を繰り返し確保（Allocate）と解放（Free）する過程で発生します。

フラグメンテーション発生のメカニズム

1. メモリの断片化: 様々なサイズのTensorがランダムな順序で確保・解放されると、メモリ空間に「隙間」が生じます。例えば、大きな連続したメモリブロックが必要な場合、全体の空き容量は足りていても、それが小さな断片に分かれているため、確保に失敗します。

2. CUDAメモリアロケータの挙動: PyTorchなどのフレームワークは、パフォーマンス向上のため、OSにメモリを返さずに内部でキャッシュ（「メモリプール」）します。これにより確保・解放のオーバーヘッドは減りますが、キャッシュされたメモリブロックのサイズが最適化されないと、フラグメンテーションが蓄積されやすくなります。

3. 長期実行時の蓄積: 長時間のトレーニングや、異なるサイズのバッチを処理する推論サーバーなどでは、フラグメンテーションが少しずつ進行し、最終的に致命的なメモリ不足エラーを引き起こします。

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

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

まず、本当にフラグメンテーションが原因なのかを確認します。PyTorchには便利なメモリプロファイリング機能があります。

import torch
# 現在のメモリ統計を詳細に表示
print(torch.cuda.memory_summary())

出力には、「largest block」のサイズがキーとなります。必要なメモリ量がこの「最大連続空きブロック」サイズを上回っている場合、フラグメンテーションが原因である可能性が高いです。

より視覚的に確認するには、メモリスナップショットを取得・分析します。

# メモリスナップショットの取得（PyTorch 1.10以降）
from torch.profiler import profile, record_function, ProfilerActivity
import torch

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
             record_shapes=True,
             profile_memory=True,
             schedule=torch.profiler.schedule(wait=0, warmup=0, active=1)) as prof:
    with record_function("model_inference"):
        # ここに問題を引き起こす推論やトレーニングのコードを記述
        pass

# スナップショットをファイルに保存（可視化ツールで開ける）
prof.export_memory_timeline("memory_timeline.html")

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

最も簡単な方法は、CUDAのキャッシュされたメモリを強制的にクリアすることです。開発中や、サービスを一時停止できる場合に有効です。

import torch
import gc

# Pythonのガベージコレクションを実行
gc.collect()

# PyTorchのCUDAキャッシュメモリを全て解放
torch.cuda.empty_cache()

# メモリ状態を確認
print(f"空きメモリ: {torch.cuda.memory_reserved() - torch.cuda.memory_allocated()} bytes")

注意点: empty_cache()はパフォーマンスに影響します。トレーニングループ内で頻繁に呼び出すと、かえって実行速度が大幅に低下するので避けてください。

ステップ3: 根本的な対策 – メモリアロケーション戦略の最適化

1. 固定サイズのバッチ処理: 可能な限り、入力データのサイズ（バッチサイズ、画像解像度など）を統一します。可変長シーケンスを扱う場合は、パディングやマスクを活用し、GPUメモリ上では固定サイズのTensorとして扱うようにします。

2. 最大メモリ使用量の制限: PyTorchのメモリアロケータに使用量の上限を設定し、大きなブロックの確保を早い段階で行わせます。

# 利用可能なGPUメモリの80%までしか使わないように設定
torch.cuda.set_per_process_memory_fraction(0.8)

3. 環境変数によるアロケータのチューニング: PyTorchのバックエンドであるCUDAのアロケータ挙動を環境変数で調整できます。

# シェルで実行（Linux/macOS）
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

# Pythonスクリプト内で設定
import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

max_split_size_mbは、アロケータがメモリブロックを分割する際の最大サイズを制御します。値を小さくするとフラグメンテーションは減りやすくなりますが、大きなTensorの確保に失敗するリスクがあります。デフォルトは非常に大きい値に設定されているため、適切な値（例: 128, 256, 512）に下げることで改善が見られる場合があります。

ステップ4: 高度な対策 – カスタムメモリ管理とツールの活用

1. トレーニングスクリプトの構造見直し: 不要な中間Tensorを早期に.detach()や.cpu()で解放したり、with torch.no_grad():ブロックを適切に使用して勾配計算用のメモリを確保しないようにします。

# 非効率的な例
for data in dataloader:
    output = model(data)
    loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    # この時点で、output, lossなどはまだGPUメモリを占有

# 改善例
for data in dataloader:
    optimizer.zero_grad()
    with torch.set_grad_enabled(True): # 明示的に勾配計算をON
        output = model(data)
        loss = criterion(output, target)
    loss.backward()
    optimizer.step()
    # 必要ない中間変数を明示的に削除
    del output, loss

2. 専用ツールの使用: NVIDIAが提供するNsight SystemsやPyTorch Profilerの詳細なタイムライン分析機能を使用すると、メモリ確保/解放のパターンを可視化し、ボトルネックを特定できます。

コード例・コマンド例：総合的なチェックスクリプト

以下は、フラグメンテーションの可能性をチェックし、基本的なクリア操作を行うユーティリティスクリプトの例です。

import torch
import gc

def check_and_clear_cuda_fragmentation(threshold_gb=1.0):
    """
    CUDAメモリのフラグメンテーションをチェックし、必要に応じてクリアする。
    threshold_gb: 最大連続ブロックがこの値(GB)より小さい場合、警告&クリアを試みる。
    """
    if not torch.cuda.is_available():
        print("CUDA is not available.")
        return

    # 現在のメモリ状態を取得
    reserved = torch.cuda.memory_reserved(0)
    allocated = torch.cuda.memory_allocated(0)
    free = reserved - allocated

    # メモリ統計から最大ブロックサイズを推測（簡易版）
    # 注: 正確な値はmemory_stats()やmemory_summary()で確認が必要
    stats = torch.cuda.memory_stats(0)
    # 'largest_block'が利用可能な場合（PyTorchバージョンに依存）
    largest_block_key = 'largest_block_size' if 'largest_block_size' in stats else None
    if largest_block_key:
        largest_block = stats[largest_block_key]
    else:
        # 簡易的な推定: 空きメモリの一定割合を最大ブロックと仮定
        largest_block = free * 0.5

    largest_block_gb = largest_block / (1024**3)
    free_gb = free / (1024**3)

    print(f"確保済みメモリ: {allocated / 1024**3:.2f} GB")
    print(f"予約済みメモリ: {reserved / 1024**3:.2f} GB")
    print(f"空きメモリ (予約内): {free_gb:.2f} GB")
    print(f"推定最大連続ブロック: {largest_block_gb:.2f} GB")

    if largest_block_gb  threshold_gb:
        print(f"⚠️  警告: フラグメンテーションの可能性があります。")
        print("キャッシュをクリアします...")
        gc.collect()
        torch.cuda.empty_cache()
        print("クリア完了。")
    else:
        print("✅ メモリ状態は正常です。")

if __name__ == "__main__":
    # 1GB以上の連続ブロックがなければ警告
    check_and_clear_cuda_fragmentation(threshold_gb=1.0)

まとめ・補足情報

CUDAメモリフラグメンテーションは、特に長時間動作するAIアプリケーションや、メモリ使用パターンが複雑なモデルで顕在化する問題です。単純な「メモリ不足」と誤解されがちですが、その実態はメモリの「断片化」にあります。

要点まとめ:

• 症状: 全体の空き容量は足りているのに「CUDA out of memory」エラーが発生する。
• 検出: torch.cuda.memory_summary()で「largest block」サイズを確認する。
• 応急処置: torch.cuda.empty_cache()とgc.collect()の実行。
• 根本対策: バッチサイズの固定、環境変数PYTORCH_CUDA_ALLOC_CONFの調整、不要なメモリ保持を避けるコード設計。
• 予防: プロファイリングツールを用いた定期的なメモリ使用パターンの監視。

最後に、CUDA 11以降の新しい機能である「メモリ管理API」や、PyTorchの実験的機能として提供されている「分割可能なアロケータ」にも注目してください。フレームワークやドライバのバージョンアップによって、より高度なメモリ管理オプションが利用可能になる可能性があります。問題に直面した際は、使用しているPyTorch/TensorFlow、CUDAドライバ、GPUハードウェアの組み合わせに応じた最新のベストプラクティスを常に確認することをお勧めします。