0
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

@AISOL

CUDAのメモリ不足エラーは本当にメモリ不足？メモリフラグメンテーションの見つけ方と確実な解決策

0
Posted at

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

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)

「たった2MiBの確保に失敗？GPUの空きメモリはまだあるはずなのに…」と感じたあなたの直感は正しいかもしれません。このエラーの背後には、**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がテンソルなどのデータを格納する際には、物理的にバラバラな小領域ではなく、一続きの連続したメモリブロックを必要とする点です。 メモリフラグメンテーションが発生すると、空きメモリの総量は足りていても、要求されたサイズの「連続した空き領域」が見つからず、アロケーション（確保）に失敗します。

なぜフラグメンテーションが発生するのか？3つの主要な原因

1. メモリの頻繁な確保と解放（特に可変サイズ）

トレーニングループ内で、バッチサイズが可変だったり、中間テンソルを繰り返し作成・破棄したりするコードは典型的な原因です。時間の経過とともに、メモリ空間に大小さまざまな「穴」（解放された領域）が散在し、断片化が進行します。

2. フレームワークのメモリアロケータの挙動（PyTorch/TensorFlow）

パフォーマンス向上のため、PyTorchやTensorFlowは解放されたメモリを即座にGPUドライバに返さず、内部の「メモリプール（キャッシュ）」に保持します。このキャッシュはサイズごとに管理されるため、特定のサイズのブロックが枯渇すると、それより大きな連続ブロックの確保が困難になることがあります。

3. CUDAコンテキストとメモリの「予約」

CUDAを初期化する際や、フレームワークの起動時に、一定量のメモリが先行確保（予約）されます。この予約メモリの位置と、アプリケーションが使用するメモリの配置パターンが悪いと、大きな連続領域が分断される原因となります。

実践的解決ステップ：検出から解消まで

ステップ1：フラグメンテーションの存在を確認する

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

import torch

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

# より簡潔なスナップショット
print(torch.cuda.memory_snapshot())

# 基本的な統計量を個別に取得
print(f"Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"Cached:    {torch.cuda.memory_reserved() / 1024**3:.2f} GB")
# 「Cached」（予約済み）と「Allocated」（実際使用中）の差がフラグメンテーションの可能性を示す

より低レベルでGPUの生のメモリ情報を見たい場合は、pynvmlライブラリが便利です。

!pip install pynvml  # Colabなどでは必要
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")

ステップ2：メモリの動的挙動をプロファイリングする（根本原因の特定）

フラグメンテーションを引き起こしている具体的なコード部分を特定するには、メモリプロファイリングが必須です。PyTorch 1.10以降では、組み込みのプロファイラを使用できます。

# メモリ履歴の記録を開始（エントリ数を多めに設定）
torch.cuda.memory._record_memory_history(max_entries=100000)

# ここに、問題が発生するトレーニングループや推論コードを実行
# 例: for epoch in range(epochs): train_one_epoch(...)

# プロファイリングを停止し、スナップショットをファイルに保存
torch.cuda.memory._dump_snapshot("memory_snapshot.pickle")

print("スナップショットを保存しました。可視化ツールで開いてください。")

生成されたmemory_snapshot.pickleファイルは、Chromeブラウザのchrome://tracingページで読み込むことで、時間軸に沿ったメモリの確保・解放イベントを視覚的に分析できます。これにより、どのタイミングでどのサイズのメモリが断片化を引き起こしているのかが一目瞭然です。

ステップ3：即効性のある解消法を試す

プロファイリングの前に、または軽度のフラグメンテーションに対しては、以下の方法で解消できることがあります。

方法A: CUDAキャッシュのクリア
PyTorchが保持しているメモリプールを強制的に空にし、GPUドライバにメモリを返させます。

import torch
import gc

# Pythonのガベージコレクションを実行
gc.collect()
# PyTorchのCUDAキャッシュを全て解放
torch.cuda.empty_cache()

# メモリ状態を確認
print(f"クリア後 Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
print(f"クリア後 Cached:    {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

方法B: 環境変数によるアロケータの挙動変更
PyTorchのメモリアロケータの振る舞いを変える環境変数を設定します。特にPYTORCH_CUDA_ALLOC_CONFは効果的です。

# ターミナルで実行する場合
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は、アロケータがメモリブロックを分割する最大サイズを制限します。値を小さくしすぎるとオーバーヘッドが増えるので、128や64などから試すのがおすすめです。

ステップ4：根本的な対策で再発を防ぐ

一時的な解消ではなく、コードレベルで再発を防ぐための対策です。

対策1: メモリの確保・解放パターンの最適化

  • テンソルの使い回し: ループ内で同じサイズの中間テンソルを繰り返し作成する場合は、初回のみ作成して使い回す。
  • バッチサイズの固定: 可能な限りバッチサイズを固定し、可変サイズによる異なるサイズのメモリ要求を減らす。
  • 不要なテンソルの早期解放: delキーワードで明示的に参照を解除し、torch.cuda.empty_cache()を適宜呼び出す。
# 悪い例: ループごとに新しいテンソルを作成
for data in dataloader:
    intermediate = torch.randn(1000, 1000, device='cuda') # 毎回新規確保
    # ...処理...

# 良い例: テンソルを事前に確保して使い回す
intermediate = torch.empty(1000, 1000, device='cuda') # 一度だけ確保
for data in dataloader:
    intermediate.fill_(0) # 再利用
    # ...処理... (intermediateを使用)

対策2: より効率的なデータローダの使用
DataLoaderのpin_memory=TrueオプションはCPU-GPU間の転送を高速化しますが、ホストメモリを多く消費します。メモリが逼迫している環境では、このオプションをFalseに設定して様子を見るのも一手です。

対策3: モデルと処理の見直し

  • 勾配チェックポイント: メモリ使用量と計算時間のトレードオフですが、非常に大きなモデルでは有効です。 
    from torch.utils.checkpoint import checkpoint

# モデルのメモリ消費の大きい部分をチェックポイントでラップ
def custom_forward(module, input):
    # 順伝播を分割
    return checkpoint(module, input)
  • 混合精度訓練 (AMP): テンソルをFP16（半精度）で保持することで、メモリ使用量をほぼ半分に削減できます。 
    from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for data, target in dataloader:
    optimizer.zero_grad()
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

まとめ

CUDAのout of memoryエラーは、単純な物理メモリ不足だけが原因ではありません。メモリフラグメンテーションという「見えない敵」が原因であることが多々あります。

  1. 確認: torch.cuda.memory_summary()やnvidia-smiで、空き容量と確保失敗の矛盾を確認する。
  2. 特定: PyTorchのメモリプロファイラを使って、問題の起きているコード部分を特定する。
  3. 応急処置: torch.cuda.empty_cache()や環境変数設定で一時的に解消する。
  4. 根本対策: テンソルの使い回し、バッチサイズ固定、混合精度訓練など、コードレベルで再発を防ぐ。

これらのステップを踏むことで、謎のGPUメモリ不足エラーから解放され、より安定した深層学習の開発を行うことができるでしょう。

この記事の詳細版は元記事サイトで公開しています。

0
1
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
0
1

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?