cpu_offload vRAM memory consumption large than 4GB

[Sanster]

Describe the bug

I am using the code from https://huggingface.co/docs/diffusers/optimization/fp16#offloading-to-cpu-with-accelerate-for-memory-savings to test cpu_offload, but the vRAM memory consumption is large than 4GB

GPU	cpu_offload enabled	vRAM cost
1080	Yes	4539MB
1080	No	5101MB
TITAN RTX	Yes	5134MB
TITAN RTX	No	5668MB

Reproduction

I am using the code from https://huggingface.co/docs/diffusers/optimization/fp16#offloading-to-cpu-with-accelerate-for-memory-savings

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    
    torch_dtype=torch.float16,
)
pipe = pipe.to("cuda")

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_sequential_cpu_offload()
image = pipe(prompt).images[0]

Logs

No response

System Info

test on 1080/TITAN RTX

• diffusers version: 0.11.1
• accelerate version: 0.15.0
• Platform: Linux-4.15.0-142-generic-x86_64-with-glibc2.29
• Python version: 3.8.10
• PyTorch version (GPU?): 1.10.1+cu111 (True)
• Huggingface_hub version: 0.11.1
• Transformers version: 4.25.1
• Using GPU in script?: Yes
• Using distributed or parallel set-up in script?: No

[patrickvonplaten]

Hey @Sanster,

Thanks a lot for the super clean bug report. When running your code-snippet in combination with:

nvidia-smi

I can observe the same numbers as in your table.

I think the problem is that we move all modules to GPU in the very beginning. We shouldn't do this. When setting:

pipe.enable_sequential_cpu_offload()

It's important to previously not run .to("cuda"). E.g. the following should give much better memory numbers:

import torch
from diffusers import StableDiffusionPipeline

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    safety_checker=None,
)

prompt = "a photo of an astronaut riding a horse on mars"
pipe.enable_sequential_cpu_offload()
image = pipe(prompt, num_inference_steps=4).images[0]

When running the above I'm getting <3GB memory usage.

In this example here, I'm not using the safety_checker since there is a bug with cpu_offload + safety checker at the moment that should however be corrected in this PR: #1968 along-side the incorrect documentation as spotted by you. Thanks a lot ❤️

It's an interesting use case here since by it might not be super intuitive to have to remove .to("cuda") when using enable_sequential_cpu_offload(...).

cc @pcuenca @patil-suraj @anton-l and maybe also @sgugger @muellerzr just FYI since it might be a common problem people run into.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.