この記事は公開から2年以上経過しています。
NVIDIA RTX3060(12GB)環境でHugging Face Stable DiffusionをCUDAで実行したところVRAMでメモリ不足が発生してしまったので、その対応についての備忘録。
問題
十分なVRAMが搭載されているGPU環境でStableDiffusionを実行しても
Exception has occurred: OutOfMemoryError
CUDA out of memory. Tried to allocate 1024.00 MiB (GPU 0; 11.73 GiB total capacity; 7.46 GiB already allocated; 1013.94 MiB free; 9.04 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
のようにVRAMメモリ不足エラーが発生する。
原因
PyTorchのメモリアロケーターがメモリを確保する際、メモリブロックサイズが大きすぎて確保に失敗している(VRAMが断片化している)。
対応
PyTorchのメモリアロケーター設定でメモリブロックの最大分割サイズを指定する。
import torch
from diffusers import StableDiffusionPipeline
# ここでメモリアロケーターの最大分割サイズを小さめに設定する
torch.cuda.memory._set_allocator_settings("max_split_size_mb:100")
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", revision="fp16", torch_dtype=torch.float16, use_auth_token=YOUR_TOKEN)もし、
CachingAllocator option max_split_size_mb too small, must be > XXX
のようなエラーが発生する場合は、このエラーが発生しないサイズでトライしてみてください。
また、上記対応を行ってもVRAM不足が解消できない場合は他のプログラムがGPUを使用している場合もあるため、nvidia-smiコマンドでVRAMを使用しているプログラムを確認して、そのプログラムを終了してみてください。
上図では既に他のプログラムによって1.4GB位が使用済であることが分かります。
参考ウェブサイトなど
以上です。