huggingface / huggingface/diffusers
Freeing GPU memory after `torch.compile` StableDiffusionXLPipeline UNet
- Langage dominant
- Python
- Étoiles
- 34.5k
- Forks
- 7.3k
- Merge moyen
- 3 j 3 h
- PR mergées (30 j)
- 91
Description
While exploring optimizations listed in the [documentation](https://huggingface.co/docs/diffusers/optimization/torch2.0), I find myself unable to free GPU memory after using `torch.compile` on a StableDiffusionXLPipeline UNet.
```Python
from diffusers import StableDiffusionXLPipeline
pipe = StableDiffusionXLPipeline.from_pretrained(
'stabilityai/stable-diffusion-xl-base-1.0',
torch_dtype=torch.float16,
variant="fp16",
use_safetensors=True
).to('cuda')
# Compile UNet
pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
generator = torch.Generator(device="cuda").manual_seed(42)
prompt = "a photo of an astronaut riding a horse on mars"
image = pipe(prompt=prompt, num_inference_steps=20, generator=generator).images[0]
del pipe
gc.collect()
torch._dynamo.reset()
torch.cuda.empty_cache()
torch.cuda.synchronize()
# GPU memory is still in use, but it's not the case when we do not compile the pipeline unet.
```
It can sometimes be useful to free the GPU memory, especially if you want to load and compile another pipeline checkpoint to perform another large number of generations.
I made a [code reproduction](https://colab.research.google.com/drive/191XutMhBarF0MFXOmPbsOQH-a8CB9vmE?usp=drive_link) in collab for testing.
Am I missing something? Could it be a [memory leak](https://dev-discuss.pytorch.org/t/fixing-torch-compile-reference-leaks-automatic-deletion-of-dynamo-code-objects/2197) on the compilation backend side, in which case it might be better to turn to PyTorch to discuss about this?
### System Info
python: 3.10.12
diffusers: 0.30.3
torch: 2.4.1+cu121
Running on Google Colab?: Yes
Guide de contribution
Ouvrir le guide de contribution
Piste de recherche
Commencez par la reproduction Colab liée et comparez le nettoyage de la mémoire GPU avec et sans torch.compile sur le UNet de StableDiffusionXLPipeline. Consultez la discussion référencée sur les fuites de références dans torch.compile et déterminez si le comportement relève de diffusers ou de PyTorch. La tâche est terminée lorsque la cause a été identifiée et qu’un fix ou workaround confirmé a été consigné.
Rédigé par le modèle d'indexation à partir du texte de l'issue.
Évaluation
- Stack technique
- python, pytorch
- Domaine
- machine-learning, performance
- Type d'issue
- Bug
- Difficulté
- 4/5
- Temps estimé
- 3-5 jours
- Activité
- À l'abandon
- Clarté
- À clarifier
- Accessibilité débutants
- 35/100