GPU・環境構築

【CUDA】GPUメモリ不足の真犯人？CUDAメモリフラグメンテーションの検出・解消法

· 15分で読めます · 1 閲覧

問題の概要：メモリは十分なのに「out of memory」が発生する

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

RuntimeError: CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 23.70 GiB total capacity; 10.34 GiB already allocated; 17.51 GiB free; 10.88 GiB reserved in total by PyTorch)

一見すると、GPUメモリの空き容量（上記例では17.51 GiB）は、確保しようとしているメモリ（2.00 GiB）よりも十分に大きいにもかかわらず、「CUDA out of memory」エラーが発生します。この矛盾の根底にある一般的な原因の一つが、CUDAメモリフラグメンテーション（メモリ断片化）です。これは、メモリの確保と解放を繰り返す過程で、空きメモリが小さな断片に分断され、連続した大きなメモリブロックを確保できなくなる現象です。

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

CUDAメモリの管理（特にPyTorchなどの深層学習フレームワークにおける）は、多くの場合「メモリアリーナ」と呼ばれる仕組みに基づいています。フレームワークは、効率化のために大きなメモリブロックを事前に確保（キャッシュ）し、そこから必要なサイズを割り当てます。

問題は、以下のような操作パターンで発生します。

  • 異なるサイズのテンソルを頻繁に作成・破棄する
  • 可変長シーケンスをバッチ処理する際のパディング操作
  • モデルの途中で大きなワークメモリを一時的に使用する
  • メモリ解放を明示的に制御しないコードの実行

これらの操作により、メモリアリーナ内の空き領域が「スイスチーズ」のように穴だらけになり、合計の空き容量は足りていても、要求されたサイズの連続した一つの空きブロックが見つからなくなります。これが「CUDA out of memory」エラーの原因となります。

フラグメンテーションを引き起こす具体的なコード例

import torch

# 様々なサイズのテンソルを繰り返し作成・破棄するループ
for i in range(1000):
    # サイズが毎回変わるテンソルを作成
    size = torch.randint(100, 10000, (1,)).item()
    temp_tensor = torch.randn(size, size, device='cuda')
    # 何らかの計算...
    result = temp_tensor * 2
    # temp_tensorはスコープを外れるが、メモリの解放タイミングはフレームワーク任せ
    # 明示的な del と gc.collect() がない場合、フラグメンテーションが進行しやすい

解決方法：検出と解消のステップバイステップ

ステップ1：メモリフラグメンテーションの検出

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

import torch
# 現在のメモリ統計を表示
print(torch.cuda.memory_summary(device=None, abbreviated=False))

出力の「Allocated memory」と「Free memory」の関係を確認します。また、より視覚的に確認するために、pynvmlライブラリを使用する方法もあります。

# pynvmlのインストール: pip install pynvml
from pynvml import *
nvmlInit()
handle = nvmlDeviceGetHandleByIndex(0) # GPU 0
info = 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の内部メモリアリーナの状態をダンプする方法が有効です。

# PyTorch 1.10以降で利用可能
import torch
# メモリアリーナの全ブロック状態をダンプ（非常に詳細な出力）
print(torch.cuda.memory_snapshot())

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

最も簡単な方法は、CUDAメモリのキャッシュを完全にクリアすることです。これは開発・デバッグ中に有効です。

import torch
import gc

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

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

# クリア後のメモリ状態を確認
print(f"Memory allocated after clear: {torch.cuda.memory_allocated(0) / 1024**3:.2f} GB")
print(f"Memory cached after clear: {torch.cuda.memory_reserved(0) / 1024**3:.2f} GB")

注意点: empty_cache()は、カーネル実行中やグラフ内で使用中のテンソルは解放しません。また、頻繁に実行するとパフォーマンスオーバーヘッドの原因となります。

ステップ3：根本的な対策 – メモリ使用パターンの改善

1. テンソルの事前確保と再利用: 可能な限り、同じサイズのテンソルを使い回す設計にします。

# 非推奨: ループ内で都度確保
for data in dataloader:
    temp = torch.zeros(data.shape, device='cuda') # 毎回新しいメモリ確保

# 推奨: 最大サイズを事前に確保して再利用
max_shape = (batch_size, max_seq_len, hidden_dim)
buffer = torch.zeros(max_shape, device='cuda')
for data in dataloader:
    actual_size = data.shape[1] # 実際のシーケンス長
    temp = buffer[:, :actual_size, :] # ビューを作成し、メモリは再利用
    # ... 処理

2. 固定サイズバッチの使用: 可変長バッチの代わりに、パディングを用いた固定サイズバッチを使用します。

3. メモリ解放の明示化: 大きな一時テンソルが不要になった時点で明示的に解放します。

large_tensor = torch.randn(10000, 10000, device='cuda')
# 処理...
del large_tensor  # 参照を削除
gc.collect()      # ガベージコレクションを促す
torch.cuda.empty_cache() # 状況に応じて

ステップ4：環境設定による緩和

PyTorchのメモリアリーナの動作を環境変数で調整できます。

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

max_split_size_mbを設定すると、指定サイズ（MB）以上の空きブロックの分割を防ぎ、大きな連続ブロックを維持しやすくなります。値はハードウェアとワークロードに応じて調整が必要です（32, 64, 128など）。

また、キャッシュサイズを制限することで、フラグメンテーションをある程度抑制できます。

# キャッシュされる最大メモリ量を制限（例：10GB）
torch.cuda.set_per_process_memory_fraction(0.8) # 全メモリの80%まで使用

コード例・コマンド例：総合的なトラブルシューティングスクリプト

以下は、フラグメンテーションの調査から一時的な解消までを行う実用的なスクリプト例です。

import torch
import gc
import sys

def diagnose_cuda_memory():
    """CUDAメモリの状態を診断"""
    print("=== CUDA Memory Diagnosis ===")
    print(f"Allocated: {torch.cuda.memory_allocated(0) / 1024**3:.2f} GB")
    print(f"Cached: {torch.cuda.memory_reserved(0) / 1024**3:.2f} GB")
    if torch.cuda.memory_reserved(0) > 0:
        frag_ratio = (torch.cuda.memory_reserved(0) - torch.cuda.memory_allocated(0)) / torch.cuda.memory_reserved(0)
        print(f"Fragmentation ratio (1 is worst): {frag_ratio:.3f}")

def clear_cuda_memory():
    """CUDAメモリを可能な限りクリア"""
    print("nClearing CUDA memory...")
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    print("Clear completed.")
    diagnose_cuda_memory()

def check_large_contiguous_block(min_size_gb=1.0):
    """指定サイズ以上の連続メモリブロックが存在するかシミュレーション"""
    try:
        # 指定サイズのテンソルを確保してみる
        test_tensor = torch.zeros(int(min_size_gb * 1024**3 // 4), dtype=torch.float32, device='cuda')
        del test_tensor
        print(f"n✓ Contiguous block of at least {min_size_gb} GB is available.")
        return True
    except RuntimeError as e:
        print(f"n✗ Cannot allocate contiguous block of {min_size_gb} GB.")
        print(f"Error: {e}")
        return False

if __name__ == "__main__":
    # 1. 現在の状態を診断
    diagnose_cuda_memory()
    
    # 2. 大きな連続ブロックが確保できるかテスト
    check_large_contiguous_block(min_size_gb=1.0)
    
    # 3. ユーザーにクリアの可否を確認
    if input("nClear CUDA memory? (y/n): ").lower() == 'y':
        clear_cuda_memory()
        # クリア後に再テスト
        check_large_contiguous_block(min_size_gb=1.0)

まとめ・補足情報

CUDAメモリフラグメンテーションは、合計空き容量があるのにメモリ確保に失敗するという紛らわしいエラーの主要原因です。対策は以下のように段階的に行うことが効果的です。

  1. 検出: torch.cuda.memory_summary()やmemory_snapshot()で現状を把握。
  2. 応急処置: torch.cuda.empty_cache()でキャッシュをクリア。
  3. 根本対策: テンソルの再利用、固定サイズバッチ、明示的解放によるメモリ使用パターンの改善。
  4. 環境調整: PYTORCH_CUDA_ALLOC_CONF環境変数によるアロケータの挙動最適化。

補足:

  • PyTorchのバージョンアップによりメモリアロケータは改善されるため、最新版の使用を検討してください。
  • マルチGPU環境では、プロセスごとのメモリ管理が独立しているため、状況はさらに複雑になります。
  • 根本的な解決には、メモリプロファイラ（torch.profiler.profileやNsight Systems）を使用した詳細な分析が不可欠です。
  • メモリフラグメンテーションはCUDAに限らず、長時間動作するサーバーアプリケーションでは一般的な問題です。定期的な再起動をプロセスに組み込むことも現実的な解決策の一つです。

これらの手法を組み合わせることで、CUDAメモリフラグメンテーションによる「out of memory」エラーを効果的に回避・解消し、GPUリソースを最大限に活用できるようになります。

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