running out of memory on introduction_to_stablediffusion.ipynb

[neelsoumya]

Thanks for this notebook @rkdan ! This is great!

In introduction_to_stablediffusion.ipynb running on Google Colab with an instance of a T4 GPU, I am running out of memory.

https://github.com/acceleratescience/diffusion-models/blob/main/notebooks/introduction_to_stablediffusion.ipynb

My code is here:

https://github.com/neelsoumya/diffusion-models/blob/main/notebooks/introduction_to_stablediffusion.ipynb

I can run until here (where I get an error):

# run both experts
image_base = base(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_end=high_noise_frac,
    output_type="latent",
).images

image_final = refiner(
    prompt=prompt,
    num_inference_steps=n_steps,
    denoising_start=high_noise_frac,
    image=image_base,
).images[0]

Error message is

---------------------------------------------------------------------------
OutOfMemoryError                          Traceback (most recent call last)
<ipython-input-8-410c6a79b4f8> in <cell line: 0>()
      7 ).images
      8 
----> 9 image_final = refiner(
     10     prompt=prompt,
     11     num_inference_steps=n_steps,

17 frames
/usr/local/lib/python3.11/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    547                 self.groups,
    548             )
--> 549         return F.conv2d(
    550             input, weight, bias, self.stride, self.padding, self.dilation, self.groups
    551         )

OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB. GPU 0 has a total capacity of 14.74 GiB of which 306.12 MiB is free. Process 6006 has 14.44 GiB memory in use. Of the allocated memory 14.17 GiB is allocated by PyTorch, and 145.94 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

My code is here:

https://github.com/neelsoumya/diffusion-models/blob/main/notebooks/introduction_to_stablediffusion.ipynb

[neelsoumya]

Gemini gives the following suggestions:

---------------------------------------------------------------------------
OutOfMemoryError                          Traceback (most recent call last)
<ipython-input-8-410c6a79b4f8> in <cell line: 0>()
      7 ).images
      8 
----> 9 image_final = refiner(
     10     prompt=prompt,
     11     num_inference_steps=n_steps,

17 frames
/usr/local/lib/python3.11/dist-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    547                 self.groups,
    548             )
--> 549         return F.conv2d(
    550             input, weight, bias, self.stride, self.padding, self.dilation, self.groups
    551         )

OutOfMemoryError: CUDA out of memory. Tried to allocate 512.00 MiB. GPU 0 has a total capacity of 14.74 GiB of which 306.12 MiB is free. Process 6006 has 14.44 GiB memory in use. Of the allocated memory 14.17 GiB is allocated by PyTorch, and 145.94 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

# Reduce image size during inference
    image_base = base(
        prompt=prompt,
        num_inference_steps=n_steps,
        denoising_end=high_noise_frac,
        output_type="latent",
        height=512,  # Reduce height
        width=512,  # Reduce width
    ).images

3. Gradient Accumulation:

Gradient accumulation simulates a larger batch size by accumulating gradients over multiple steps before updating model weights. This allows you to effectively increase batch size without increasing memory usage per step. However, this can be slightly more complex and may require adjusting the training loop.
4. Gradient Checkpointing:

Gradient checkpointing trades computation time for memory by recomputing parts of the computational graph during the backward pass. This significantly reduces memory usage but increases training time. You can enable gradient checkpointing using the torch.utils.checkpoint module.

import torch.utils.checkpoint as checkpoint

    # Wrap a layer with checkpointing
    def forward(x):
        return checkpoint.checkpoint(layer, x)

5. Use fp16 or bf16 precision:

You are already using fp16 precision. Using lower precision (fp16 or bf16) can significantly reduce memory usage with minimal impact on accuracy in many cases. Ensure you have enabled mixed precision training with torch.amp if not already done.
6. Restart Runtime and Clear Memory:

Try restarting the runtime and clearing any unnecessary variables or data in your notebook to free up memory before running the cell again.
7. Enable Memory Growth:

You can enable memory growth in PyTorch to allow it to dynamically allocate memory as needed.

import torch

torch.cuda.empty_cache()  # Clear any cached memory
torch.backends.cudnn.benchmark = True # Enable cuDNN auto-tuner to optimize performance

# Enable memory growth for GPU 0
torch.cuda.set_device(0) 
torch.cuda.set_per_process_memory_fraction(0.9, 0)

# Or:
# for device in range(torch.cuda.device_count()):
#     torch.cuda.set_device(device)
#     torch.cuda.empty_cache()
#     torch.backends.cudnn.benchmark = True

Try applying these suggestions to your code to resolve the OutOfMemoryError. I would recommend trying to reduce image resolution first.