Running into out of memory issues

KVS_Moudgalya: <p>Hi,</p>
<ul>
<li>I’m currently trying to train diffusion model for 2D image generation task with images as input.</li>
<li>Training on AWS G5 instances i.e., A10G GPU’s with 24GB GPU memory.</li>
<li>I’m running into out of memory issues when I go beyond image size of 256x256 and batch size of 8.</li>
<li>Results with Image size = 256 and batch size = 8 is unacceptable</li>
<li>I did use gradient accumulation and mixed precision training.</li>
<li>Using 1 attention block only.</li>
</ul>
<p>Trying to understand is it a genuine memory issue or can it be solved by some other approaches?<br>
Is diffusion models so heavy that even a 24GB memory GPU is insufficient?<br>
What’s the typical memory requirement for running image to image diffusion models, for generating images resolution higher than 512x512</p>
<p>Thanks and regards<br>
KVS Moudgalya</p>

J_Johnson: <p>Welcome to the Pytorch Forums!</p>
<p>Couple of questions/comments:</p>
<ol>
<li>Can you provide a model summary?</li>
<li>Are you using self attention and what type?</li>
<li>Typically, UNet diffusion models are trained on 512x512 images but can be used to extend to larger images, given it’s entire structure involves convolutions.</li>
<li>What float are you using? You may find mixed precision or bfloat16 to be sufficient, but that will typically require ~half the memory as float32.</li>
<li>What optimizer are you using? Different optimizers require more or less elements in the graph per parameter.</li>
</ol>

KVS_Moudgalya: <p>Hey, thanks Johnson</p>
<ol>
<li>Employing DDPM architecture with UNet model.</li>
<li>I’m using one Q,K,V attention block.<br>
Optimizer - torch adamw<br>
Yes, I have tried gradient accumulation and mixed precision training. It did not improve.</li>
</ol>
<p>So, I am trying to understand<br>
Is diffusion models so heavy that even a 24GB memory GPU is insufficient?<br>
What’s the typical GPU memory requirement for running image to image diffusion models, for generating images of resolution higher than 512x512</p>

J_Johnson: <p>There are some additional memory-efficient methods you can make within the UNet model architecture as found in this code here:</p>
<aside class="onebox githubblob" data-onebox-src="https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py">
  <header class="source">

      <a href="https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py" target="_blank" rel="noopener nofollow ugc">lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py</a></h4>


      <pre><code class="lang-py">import math
import copy
from random import random
from beartype.typing import List, Union
from beartype import beartype
from tqdm.auto import tqdm
from functools import partial, wraps
from contextlib import contextmanager, nullcontext
from collections import namedtuple
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel
from torch import nn, einsum
from torch.cuda.amp import autocast
from torch.special import expm1
import torchvision.transforms as T

import kornia.augmentation as K
</code></pre>



  This file has been truncated. <a href="https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py" target="_blank" rel="noopener nofollow ugc">show original</a>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>And additionally, a more efficient self-attention architecture, located here:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/cmsflash/efficient-attention">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/cmsflash/efficient-attention" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/33a4d3aa69ac2d1c03f8c0a1c4c94771c457281af83b786970b87641316996e8/cmsflash/efficient-attention" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/cmsflash/efficient-attention" target="_blank" rel="noopener nofollow ugc">GitHub - cmsflash/efficient-attention: An implementation of the efficient...</a></h3>

  <p>An implementation of the efficient attention module. - GitHub - cmsflash/efficient-attention: An implementation of the efficient attention module.</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>Just note that vanilla self-attention can be a bit resource hungry.</p>
<p>Lastly, UNets are quite slow when it comes to training, even with the most recent GPUs. Especially when you get into training them on larger sizes, like 512x512. Which is why they are usually trained on multiple TPUs. However, GigaGAN looks promising, and I see Phil Wang is almost done with his Pytorch version of it, here:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/lucidrains/gigagan-pytorch">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/lucidrains/gigagan-pytorch" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/1952d59a3b599226238a260dfe4f9c723e727eb46c05f1334c173c47b321d7e2/lucidrains/gigagan-pytorch" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/lucidrains/gigagan-pytorch" target="_blank" rel="noopener nofollow ugc">GitHub - lucidrains/gigagan-pytorch: Implementation of GigaGAN, new SOTA GAN...</a></h3>

  <p>Implementation of GigaGAN, new SOTA GAN out of Adobe - GitHub - lucidrains/gigagan-pytorch: Implementation of GigaGAN, new SOTA GAN out of Adobe</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


KVS_Moudgalya: <p>Thanks <a class="mention" href="/u/j_johnson">@J_Johnson</a>  Will try this out and revert back to you the outcomes.</p>