huggingface / huggingface/diffusers

Freeing GPU memory after `torch.compile` StableDiffusionXLPipeline UNet

Ouverte
#9,530 4 commentaires 0 réactions 0 personnes assignées Voir sur GitHub
stale torch-compile
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

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.