Speedup model loading by 4-5x ⚡ by a-r-r-o-w · Pull Request #11904 · huggingface/diffusers

a-r-r-o-w · 2025-07-10T02:54:41Z

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 · 2025-07-10T03:01:36Z

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.

src/diffusers/loaders/single_file_model.py

src/diffusers/loaders/single_file_utils.py

sayakpaul · 2025-07-10T04:06:45Z

src/diffusers/models/model_loading_utils.py

    return parsed_parameters
+
+
+def _find_mismatched_keys(


Taken out of here:

diffusers/src/diffusers/models/modeling_utils.py

Line 1509 in 9f4d997

def _find_mismatched_keys(

src/diffusers/models/model_loading_utils.py

src/diffusers/models/modeling_utils.py

sayakpaul · 2025-07-10T04:18:46Z

src/diffusers/utils/torch_utils.py

+    if device_type is None:
+        device_type = get_device()
+    device_mod = getattr(torch, device_type, torch.cuda)
+    device_mod.synchronize()


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

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

src/diffusers/models/model_loading_utils.py

Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>

src/diffusers/models/model_loading_utils.py

* update * update * update * pin accelerate version * add comment explanations * update docstring * make style * non_blocking does not matter for dtype cast * _empty_cache -> clear_cache * update * Update src/diffusers/models/model_loading_utils.py Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com> * Update src/diffusers/models/model_loading_utils.py --------- Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>

a-r-r-o-w added 3 commits July 9, 2025 12:06

update

7a9c448

update

8385f45

update

9e4873b

a-r-r-o-w requested review from DN6, SunMarc and sayakpaul July 10, 2025 02:54

a-r-r-o-w mentioned this pull request Jul 10, 2025

Speedup model loading by 4-5x in Diffusers ⚡ huggingface/accelerate#3674

Merged

a-r-r-o-w added 3 commits July 10, 2025 08:33

Merge branch 'main' into speedup-model-loading

20b1155

pin accelerate version

b776aaa

Merge branch 'main' into speedup-model-loading

e364dfd

sayakpaul reviewed Jul 10, 2025

View reviewed changes

a-r-r-o-w added 5 commits July 10, 2025 11:22

add comment explanations

ea446b1

update docstring

e736b09

make style

4c81c96

Merge branch 'main' into speedup-model-loading

591655e

non_blocking does not matter for dtype cast

582af9b

sayakpaul approved these changes Jul 10, 2025

View reviewed changes

nitinmukesh mentioned this pull request Jul 10, 2025

[Feature] Speedup model loading by 4-5x nunchaku-ai/nunchaku#525

Closed

2 tasks

a-r-r-o-w added 3 commits July 11, 2025 06:02

_empty_cache -> clear_cache

a6ee660

Merge branch 'main' into speedup-model-loading

bbbc4c0

update

39f0850

SunMarc approved these changes Jul 11, 2025

View reviewed changes

src/diffusers/models/model_loading_utils.py Outdated Show resolved Hide resolved

a-r-r-o-w and others added 2 commits July 11, 2025 20:33

Update src/diffusers/models/model_loading_utils.py

58fcfdc

Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>

Merge branch 'main' into speedup-model-loading

b0552bb

a-r-r-o-w commented Jul 11, 2025

View reviewed changes

src/diffusers/models/model_loading_utils.py Outdated Show resolved Hide resolved

Update src/diffusers/models/model_loading_utils.py

275e470

a-r-r-o-w merged commit c903527 into main Jul 11, 2025
32 checks passed

a-r-r-o-w deleted the speedup-model-loading branch July 11, 2025 16:13

sayakpaul mentioned this pull request Aug 1, 2025

support hf_quantizer in cache warmup. #12043

Merged

This was referenced Aug 11, 2025

feat: cuda device_map for pipelines. #12122

Merged

[core] parallel loading of shards #12028

Merged

stevhliu mentioned this pull request Aug 13, 2025

[docs] Parallel loading of shards #12135

Merged

sayakpaul mentioned this pull request Aug 18, 2025

fix: caching allocator behaviour for quantization. #12172

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Speedup model loading by 4-5x ⚡#11904

Speedup model loading by 4-5x ⚡#11904
a-r-r-o-w merged 17 commits intomainfrom
speedup-model-loading

a-r-r-o-w commented Jul 10, 2025 •

edited

Loading

Uh oh!

HuggingFaceDocBuilderDev commented Jul 10, 2025

Uh oh!

sayakpaul left a comment

Uh oh!

Uh oh!

Uh oh!

sayakpaul Jul 10, 2025

Uh oh!

Uh oh!

Uh oh!

sayakpaul Jul 10, 2025

Uh oh!

a-r-r-o-w Jul 10, 2025 •

edited

Loading

Uh oh!

sayakpaul left a comment

Uh oh!

SunMarc left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

a-r-r-o-w commented Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Jul 10, 2025

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

sayakpaul Jul 10, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

sayakpaul Jul 10, 2025

Choose a reason for hiding this comment

Uh oh!

a-r-r-o-w Jul 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

SunMarc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

a-r-r-o-w commented Jul 10, 2025 •

edited

Loading

a-r-r-o-w Jul 10, 2025 •

edited

Loading