Speedup model loading by 4-5x ⚡

[a-r-r-o-w]

All thanks to @Cyrilvallez's PR: huggingface/transformers#36380

The accelerate PR is required because we end up calling clear_device_cache in a loop (over the sharded files). This is bad. Without this, you'll see no speedup.

Another small optimization is using non_blocking everywhere and syncing just before returning control to the user. This is slightly faster.

import time
t_ini = time.time()

import torch
from diffusers import FluxPipeline, FluxTransformer2DModel
print(f"import time: {time.time() - t_ini:.3f}s")

model_id = "black-forest-labs/FLUX.1-dev"

t0 = time.time()
torch.cuda.synchronize()
print(f"CUDA sync time: {time.time() - t0:.3f}s")

print("starting model load")
t1 = time.time()
transformer = FluxTransformer2DModel.from_pretrained(model_id, subfolder="transformer", torch_dtype=torch.bfloat16, device_map="cuda")
torch.cuda.synchronize()
t2 = time.time()

diff = t2 - t1
print(f"time: {diff:.3f}s")

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", transformer=transformer, torch_dtype=torch.bfloat16)
pipe.text_encoder.to("cuda")
pipe.text_encoder_2.to("cuda")
pipe.vae.to("cuda")
prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=28, guidance_scale=4.0).images[0]
image.save("flux.png")

Sister PR in accelerate required to obtain speedup: huggingface/accelerate#3674

• On main: 16.765s
• On this branch: 4.521s

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[sayakpaul]

Thanks a lot for quickly getting this up 🔥

My comments are mostly minor, the major one being adding hf_quantizer to the allocator function.

Additionally, for a potentially better user-experience, if we could try to rethink the to() method of DiffusionPipeline, it would be helpful. I mean the following.

Currently, from what I understand, we have to first initialize the denoiser using device_map and then the rest of the components. If a user is calling .to() on a DiffusionPipeline, we could consider using device_map="cuda" for dispatching the model-level components to CUDA. I don't immediately see a downside to it.

[sayakpaul]

Taken out of here:

diffusers/src/diffusers/models/modeling_utils.py

Line 1509 in 9f4d997

def _find_mismatched_keys(

[sayakpaul]

I guess all different backends ought to have this method. Just flagging.

[a-r-r-o-w]

afaik, synchronize should be available on all devices. Just the empty_cache function required a special check because it would fail if device was cpu

[sayakpaul]

Ship

[SunMarc]

Thanks for this ! Just a nit