Pipelines fail if enable_model_cpu_offload is called twice

[antoche]

Describe the bug

enable_model_cpu_offload seem to be tripping over its feet if it's called a second time. See repro case.

From what I can see, pipe._execution_device is cuda:0 after the first call, but switches to cpu after the second call. I'd expect the second call to be pretty much a no-op.

Reproduction

>>> from diffusers import StableDiffusionPipeline
>>> pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
>>> pipe.enable_model_cpu_offload()
>>> result = pipe(prompt="foo")
>>> result = pipe(prompt="bar")
>>> pipe.enable_model_cpu_offload()
>>> result = pipe(prompt="baz")
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Logs

No response

System Info

• diffusers version: 0.14.0
• Platform: Linux-4.14.240
• Python version: 3.9.10
• PyTorch version (GPU?): 1.13.0a0+git49444c3 (True)
• Huggingface_hub version: 0.13.1
• Transformers version: 4.27.2
• Accelerate version: 0.17.1
• xFormers version: 0.0.16+6f3c20f.d20230309

[patrickvonplaten]

Any reason why you would call it twice? I'm wondering whether it's better to have it just silently not do anything like done in #2921 or whether we should throw an error

@pcuenca @sayakpaul @williamberman wdyt?

[antoche]

In my use case, I have a helper function which constructs and prepares pipeline instances for callers to use. The function caches the models in memory to speed up pipeline instantiation. So any model a pipeline is using may have already been touched by cpu_offload or moved to any arbitrary device. Therefore enable_cpu_offload needs to be robust to that. I'd also argue there should be a way to check offload hooks have already been installed on a model, and a way to remove them

[pcuenca]

I think the simplest and most general solution could be to provide a method to remove all hooks, just as we recently did to disable the custom attention processor. Then the use case described by @antoche would be to always disable hooks before returning a new pipeline from his helper function.

[sayakpaul]

I think the simplest and most general solution could be to provide a method to remove all hooks, just as we recently did to disable the custom attention processor.

+1 on favoring the simplest solution.

[patrickvonplaten]

That works for me :-)

[williamberman]

We should note that adding cpu offload with hook is a stateful change to the models that assumes they are being run in the same pipeline. I.e. the order of installation of the hooks matters and that they are run with the same models that have the relevant hooks for model movement already installed. It also requires adding state to the pipeline for the final offload hook to ensure it's manually removed.

In your case @antoche, I would recommend caching the entire pipeline instead of the individual model components. I don't think it's that big of an issue to add something that removes all installed hooks, but I would be hesitant about commiting to making models dynamically swappable between pipeline instances as pipelines and models can be stateful.

[antoche]

Yes, I do cache entire pipeline objects, not individual models, and create pipelines by re-using components from cached pipelines, as recommended in the documentation itself (e.g., https://huggingface.co/docs/diffusers/v0.14.0/en/api/pipelines/alt_diffusion#tips).

It's probably worth noting that sharing components in the way that's described in the docs is actually not really compatible with the cpu offloading workflow. For example, if we take this example from the AltDiffusion docs:

from diffusers import (
    AltDiffusionPipeline,
    AltDiffusionImg2ImgPipeline,
)

text2img = AltDiffusionPipeline.from_pretrained("BAAI/AltDiffusion-m9")
img2img = AltDiffusionImg2ImgPipeline(**text2img.components)

# now you can use text2img(...) and img2img(...) just like the call methods of each respective pipeline

If enable_cpu_offload is called on the first pipeline, and I read the code correctly, then the second pipeline will not actually call the final offload hook, even though it is using models that have been configured for CPU offloading. That's probably an unexpected and confusing behaviour for users.

If the final offload hook was stored on the model rather than on the pipeline, then we could effectively have pipelines that share models and all support the same offloading behaviour.

[williamberman]

Yes, I think the best change would be to for now update the docs that model cpu offloading and not clearing hooks isn't strictly compatible with re-using models between pipelines.

Note that also due to how accelerate re-wraps hooks if they are installed twice can also break how diffusers inspects them for the execution device 😁

I think it's reasonable to consider offloading something that's tied to the pipeline and the models rather than the models themselves, and that once cpu offloading has been enabled on a pipeline, it shouldn't be expected to carry over to another pipeline if the components are re-used.

If you need a snippet to clear the offloading state of a pipeline, feel free to adapt this method

diffusers/src/diffusers/pipelines/deepfloyd_if/pipeline_if.py

Lines 211 to 223 in e0a2bd1

	def remove_all_hooks(self):
	if is_accelerate_available():
	from accelerate.hooks import remove_hook_from_module
	else:
	raise ImportError("Please install accelerate via `pip install accelerate`")
	for model in [self.text_encoder, self.unet, self.safety_checker]:
	if model is not None:
	remove_hook_from_module(model, recurse=True)
	self.unet_offload_hook = None
	self.text_encoder_offload_hook = None
	self.final_offload_hook = None

tl;dr

pipe = Pipe.from_pretrained("...")
pipe.enable_cpu_offloading()

pipe("foo")

remove_all_hooks(pipe) # or pipe.remove_all_hooks()

pipe = Pipe2(*pipe.components)
pipe.enable_cpu_offloading()

pipe("bar")

should be the expected way to use cpu offloading and re-using pipeline components for now

[pcuenca]

@williamberman so do you think it'd be worthwhile to expand the use of remove_all_hooks to other pipelines?

[williamberman]

@pcuenca yep! I think we also want to take a look and see if there's anything else we want to solve around hooks i.e. duplicate calls or what not. patrick and I've talked about it adhoc a few times but haven't formalized anything yet

[williamberman]

closing issue for now, we'll probably open a separate issue once we take a longer look at these smol clean up things around hooks