search

openai / consistency_models

Official repo for consistency models.
MIT License
6.02k stars 409 forks source link

CUDA out of memory #28

Open xhsongstanford opened 1 year ago

xhsongstanford commented 1 year ago

Hi,

Anyone knows the GPU requirement to train this network (Consistency Distillation)? I'm always getting CUDA out of memory errors:

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 4.00 GiB (GPU 0; 39.42 GiB total capacity; 34.41 GiB already allocated; 2.96 GiB free; 34.90 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

I've tried to decrease the batch size, the memory requirement in "Tried to allocate 4.00 GiB" do decrease, but such error still occurs as the "already allocated" data usage increases.

nekoshadow1 commented 1 year ago

Have you tried reducing num_channels and num_head_channels?