Memory usage by later pipeline stages
vRAM usage during my pipeline run looks like this:It uses a bunch of memory when I load the pipeline withfrom_pretrained, which is expected. Then more memory when I call the pipeline. Then it holds steady for a while as the pipeline is iterating through its steps, before making one smaller jump later.I assume that last one is when the VAE is invoked to decode the output.I didn’t number this chart, but the total range of the Y-axis is 8 GB, which makes that last jump something like half a gig.While that seems like a lot for a measly 0.75 MB worth of pixel data, it’s not so much the amount I’m concerned with. It’s modest in comparison to the overall needs of the pipeline. My question iswhythe allocator grabs more memory at that time.The diffusion model is done. Shouldn’t it be able to reclaim more than enough memory from that?I guess the answer is that the allocator is grabbing more memory because itcan— there’s no memory pressure yet. Throwing atorch.cuda.empty_cache()in there before that stage seems to confirm this: then memory goesdowna tad at that point, notup;what’s more is that itstaysthere for the next pipeline run.So maybe I’ve answered my own question, but I’m still confused about that big chunk of memory that isnotallocated during pipeline init, but seems to be on our firstforwardcall, and I’m not sure if or when it is reclaimed.
2023-10-18T13:44:27Z
Same problem here. TheStableDiffusionPipelinegot incremental VRAM, especially at the end of a batch inference. A simple solution is just to add this line:# This line below
pipe.final_offload_hook = TrueExplainationThe code offloads the modules before loading the VAE decoder. Concretely, I foundthis code:if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
    self.unet.to("cpu")
    self.controlnet.to("cpu")
    torch.cuda.empty_cache()So just set this final_offload_hook should work.
2023-12-12T09:24:25Z
i have a same problem, that my virtual memory reach 28Gimage974×732 25.6 KBis there any solution to get back my memory like before?