On Monday CompVis published the weights of Stable Diffusion. Their README says the model has “860M UNet and 123M text encoder” and “runs on a GPU with at least 10GB VRAM”. I multiplied in my head. About a billion numbers at 4 bytes each is 4 GB. The rest of the 10 I expected to be a reserve for CUDA and a habit of rounding up.
That was my whole model of GPU memory: the weights plus a bit. I never checked it, because everything I ran before fit anywhere. The machine where I do these experiments has no video card at all. It has 4 cores of an AMD EPYC and a bit less than 5 GiB that I can give to one process. This turned out useful, because a process that asks for too much is killed at once instead of slowing down. Everything below ran on PyTorch 1.12.1 on the CPU in float32.
Counting the weights
The files of CompVis/stable-diffusion-v1-4 hold 859.5 million numbers in the UNet alone. With the text encoder and the VAE, the part that turns the result into pixels, it comes to 1.07 billion. That is 4.26 GB in float32 and 2.13 GB in half precision. My estimate of the weights was right. It explained less than half of the 10.
The UNet alone is 3.2 GiB in float32. Loading it pushed the memory of the process up to 3.76 GiB, file cache included. Before I had clean numbers I spent a while on kills that made no sense. Then I found that my user session has its own ceiling of 5 GiB and everything else I had open counts against it.
Where the rest goes
The answer is in ldm/modules/attention.py. The UNet looks at the image through attention layers and the first of them works at 1/8 of the picture. At 512x512 it sees 64 by 64, which is 4096 positions. Every position is compared with every other one, a matrix of 4096 by 4096 numbers. There are 8 heads and the sampler runs the prompt and an empty prompt together, so there are 16 such matrices. The code computes all of them in one einsum and then softmax makes a second tensor of the same size.
16 matrices of 64 MiB give 1 GiB in float32 and with the softmax copy it is 2 GiB, which is 2048 MiB. I took this block out of the UNet and ran it alone with the real weights. Memory went up by 2241 MiB above the baseline, so the 2 tensors explain 91% of it. The rest is the other tensors of the block. Then I rewrote the same math as a loop over the heads, one matrix at a time. It took 387 MiB. A repeat run gave 405. The outputs were the same bit for bit.
I expected the loop to cost time. In 3 runs of each it took 2.04 to 2.16 seconds against 1.85 to 2.19 for all heads at once. On this CPU the difference is lost in the noise. On a card it may look different, but I would take this trade on any card that is short of memory.
The matrix grows with the square of the number of positions and this is the part I did not feel until I saw it. At 640x640 one layer needs 2.4 GiB and at 768x768 it needs 5.1 GiB, before the softmax copy. With all heads at once both runs were killed at a cap of 4 GiB. One head at a time they took 795 MiB and 1390 MiB.
The default command
scripts/txt2img.py by default makes 3 images per batch and the sampler doubles that for the empty prompt, so the UNet sees 6. It keeps the weights in float32 and runs under autocast. The queries and keys come out of the linear layers in half precision, so the scores do too. softmax is on the list of operations that autocast returns in float32. For one layer at 512x512 I count 1.5 GiB of scores in half precision and 3 GiB after the softmax. The code keeps the first one alive while it uses the second. The next einsum runs on bmm, which autocast casts to half precision, so that makes another copy of 1.5 GiB. Together it is 6 GiB for one layer, next to 3.97 GiB of weights.
That is close to 10 GiB, which fits the README surprisingly well. I did not run it on a card, so it is my reading of the code, but it can be off in both directions. Autocast may also keep half copies of the weights it has cast and I do not know how much the CUDA context takes. So I read most of the 10 GB as the default batch of 6 under autocast and less of it as the model.
One whole image
At the end I made one picture, to see that the pieces work together. The prompt was “a photograph of an astronaut riding a horse” with 50 steps of the PLMS sampler and attention one head at a time. As in the CompVis code, the prompt and the empty prompt went through the UNet as one batch of 2. Each of the 3 parts ran as a separate process, so they were never in memory together. The UNet process peaked at 3.76 GiB. Loading alone had already reached 3.70, so the 51 passes of the sampler added only 64 MiB on top. The other 2 stayed under 1.7 GiB.
On 4 cores the whole thing took 16 minutes, with 17.4 seconds for a typical step. The astronaut is there and so is the horse. How long the same work takes on a card I cannot say, I have no number of my own for it.
On paper the floor for 512x512 is now far below 10 GB. It is about 2 GiB of weights in half precision and about 0.1 GiB for the scores of one head, counting every copy autocast makes. I would want to see it on a real card before I believe it.
What I did not check
- Nothing here ran on a GPU. I have no measurement of the CUDA context or of speed on a card.
- The 10 GiB for the default command comes from the code and the documented autocast rules, not from a run.
- Every time is from 1 to 3 runs on a machine that had other work on it.
- With all heads at once I measured only the first attention block. In the whole UNet I ran only attention one head at a time. That run used the diffusers code of the UNet, not my loop.
- The VAE has its own attention at 64 by 64 and I did not look at it separately.