[FEAT] Model loading refactor #10604

SunMarc · 2025-01-17T23:48:24Z

What does this PR do?

Fixes #10013 . This PR refactors model loading in diffusers. Here's a list of major changes in this PR.

only two loading paths (low_cpu_mem_usage=True and low_cpu_mem_usage = False). We don't rely on load_checkpoint_and_dispatch anymore and we don't merge sharded checkpoint also.
support for sharded checkpoints for both loading paths
keep_module_in_fp32 support for sharded checkpoints
better support for displaying warning due to error/unexpected/missing/mismatched keys

For low_cpu_mem_usage = False:

Faster initialization (thanks to skipping the init + assign_to_params_buffers). I didn't benchmarked it but it should be as fast as low_cpu_mem_usage=True or maybe even faster. We did a similar PR in transformers thanks to @muellerzr.
Better torch_dtype support We don't initialize anymore the model in fp32 then cast the model to a specific dtype after finishing to load the weights.

For low_cpu_mem_usage = True or device_map!=None:

one path, we don't rely anymore on load_checkpoint_and_dispatch
device_map support for quantization
non persistance buffer support through dispatch_model ( the test you added is passing cc @hlky )

Single format file:

Simplified the single file format loading through from_pretrained. This way we have the same features as this function (device_map, quantization ...). Feel free to share your opinion @DN6, I didn't expect to touch this but I felt that we could simplify a bit

TODO (some items can be done in follow-up PRs):

Check if we have any regression / tests issues
Add more tests
Deal with missing keys in the model for both paths (before, it only worked when low_cpu_mem_usage=False since we are initializing the whole model)
Fix typing
Better support for offload with safetensors (like in transformers)

Please let me know your thoughts on the PR !

cc @sayakpaul, @DN6 , @yiyixuxu , @hlky , @a-r-r-o-w

SunMarc · 2025-01-18T10:54:55Z

FLAX CPU failing test is unrelated, failing in other PRs too

src/diffusers/loaders/single_file_model.py

sayakpaul

Thanks for starting this! Left some comments from a first pass.

I think we will need to also add tests for seeing if device_map works as expected for quantization. Okay to not test that a bit later once there is consensus about the design changes. Maybe we could add that as a TODO.

Other tests could include checking if we can do low_cpu_mem_usage=True along with some changed config values. This will ensure we're well tested for cases like #9343.

src/diffusers/models/model_loading_utils.py

src/diffusers/models/modeling_utils.py

src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py

sayakpaul · 2025-01-20T06:37:59Z

@SunMarc,

Additionally, I ran some tests on audace (two RTX 4090s). Some tests that are failing (they fail on main too):

Failures

FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_load_sharded_checkpoint_from_hub_0_hf_internal_testing_unet2d_sharded_dummy - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argumen...
FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_load_sharded_checkpoint_from_hub_1_hf_internal_testing_tiny_sd_unet_sharded_latest_format - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argumen...
FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_load_sharded_checkpoint_from_hub_local - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argumen...
FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_load_sharded_checkpoint_from_hub_local_subfolder - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argumen...
FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_load_sharded_checkpoint_from_hub_subfolder_0_hf_internal_testing_unet2d_sharded_dummy_subfolder - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argumen...
FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_load_sharded_checkpoint_from_hub_subfolder_1_hf_internal_testing_tiny_sd_unet_sharded_latest_format_subfolder - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argumen...
FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_sharded_checkpoints - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argumen...
FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_sharded_checkpoints_device_map - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argumen...
FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_sharded_checkpoints_with_variant - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1! (when checking argument...

^^ passes when using with CUDA_VISIBLE_DEVICES=0 (same with main). Expected?

Same for following:

FAILED tests/models/unets/test_models_unet_2d_condition.py::UNet2DConditionModelTests::test_sharded_checkpoints_device_map - RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:1!

And then I also ran:

RUN_SLOW=1 pytest tests/pipelines/stable_diffusion/test_stable_diffusion.py::StableDiffusionPipelineDeviceMapTests

Everything passes.

hlky · 2025-01-20T07:58:35Z

src/diffusers/models/model_loading_utils.py


-    for param_name, param in named_buffers:


We need to keep this or equivalent elsewhere, context: #10523

The changes I did should also cover this use case. The test you added should pass with my PR. The is mainly due to adding the dispatch_model function at the end.

src/diffusers/models/modeling_utils.py

Co-authored-by: Sayak Paul <[email protected]>

…odel-loading-refactor

src/diffusers/models/modeling_utils.py

Co-authored-by: YiYi Xu <[email protected]>

…odel-loading-refactor

src/diffusers/models/modeling_utils.py

DN6 · 2025-01-21T03:10:16Z

src/diffusers/models/modeling_utils.py

 logger = logging.get_logger(__name__)

 _REGEX_SHARD = re.compile(r"(.*?)-\d{5}-of-\d{5}")

+TORCH_INIT_FUNCTIONS = {


Not a merge blocker, but is it possible to dynamically create this mapping? Then we could avoid having to make manual updates in case new inits are added to torch.

Although I suppose that doesn't happen too often.

Something like this could work:

import torch.nn.init as init init_functions = { name: getattr(init, name) for name in dir(init) if callable(getattr(init, name)) and name.endswith("_") and not name.startswith("_") } print("Available initialization functions:") for name in init_functions: print(name)

Prints:

Available initialization functions: constant_ dirac_ eye_ kaiming_normal_ kaiming_uniform_ normal_ ones_ orthogonal_ sparse_ trunc_normal_ uniform_ xavier_normal_ xavier_uniform_ zeros_

WDYT @DN6 ? I'm fine with either. Also there are some missing function since you only choose the one finishing with "_", though I don't think these are deprecated now.

"uniform": nn.init.uniform, "normal": nn.init.normal, "xavier_uniform": nn.init.xavier_uniform, "xavier_normal": nn.init.xavier_normal, "kaiming_uniform": nn.init.kaiming_uniform, "kaiming_normal": nn.init.kaiming_normal,

src/diffusers/models/modeling_utils.py

sayakpaul

Some more comments.

I am running the 4bit quantization tests currently. And so far things are looking nice! Some tests that might be worth including/consdering:

Device map with quantization
Effectiveness of keep_modules_in_fp32 when not using quantization.

WDYT?

Edit: 4bit and 8bit tests (bitsandbytes) are passing.

sayakpaul · 2025-01-27T13:07:53Z

src/diffusers/loaders/single_file_model.py

@@ -362,17 +362,18 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: Optional[str] =

        if is_accelerate_available():
            param_device = torch.device(device) if device else torch.device("cpu")
-            named_buffers = model.named_buffers()
-            unexpected_keys = load_model_dict_into_meta(
+            unexpected_keys = [


Are the single-file related changes to uniformize the use of load_model_dict_into_meta() (with the new signature)?

yeah that's right !

src/diffusers/models/model_loading_utils.py

sayakpaul · 2025-01-27T13:25:03Z

src/diffusers/models/modeling_utils.py

 logger = logging.get_logger(__name__)

 _REGEX_SHARD = re.compile(r"(.*?)-\d{5}-of-\d{5}")

+TORCH_INIT_FUNCTIONS = {


Something like this could work:

import torch.nn.init as init init_functions = { name: getattr(init, name) for name in dir(init) if callable(getattr(init, name)) and name.endswith("_") and not name.startswith("_") } print("Available initialization functions:") for name in init_functions: print(name)

Prints:

Available initialization functions: constant_ dirac_ eye_ kaiming_normal_ kaiming_uniform_ normal_ ones_ orthogonal_ sparse_ trunc_normal_ uniform_ xavier_normal_ xavier_uniform_ zeros_

sayakpaul · 2025-01-27T13:26:18Z

src/diffusers/models/modeling_utils.py

+
+
+@contextmanager
+def no_init_weights():


Could you briefly then elaborate what happens in this codepath?

src/diffusers/models/modeling_utils.py

tests/models/test_modeling_common.py

HuggingFaceDocBuilderDev · 2025-01-27T13:42:41Z

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.

src/diffusers/models/modeling_utils.py

…odel-loading-refactor

Co-authored-by: Aryan <[email protected]>

SunMarc · 2025-02-14T16:50:22Z

I am running the 4bit quantization tests currently. And so far things are looking nice! Some tests that might be worth including/consdering:

Device map with quantization
Effectiveness of keep_modules_in_fp32 when not using quantization.

Done ! Please check

SunMarc added 3 commits January 17, 2025 22:36

first draft model loading refactor

e54c540

revert name change

645abc9

fix bnb

bd81f50

SunMarc changed the title ~~[FEAT ] Model loading refactor~~ [FEAT] Model loading refactor Jan 17, 2025

SunMarc added 4 commits January 18, 2025 10:57

revert name

17c1be2

fix dduf

72b6259

fix huanyan

b4e4f3b

style

5a00dc6

SunMarc requested review from sayakpaul, DN6, yiyixuxu and a-r-r-o-w January 18, 2025 10:48

Merge branch 'main' into model-loading-refactor

3bcd6cc

sayakpaul reviewed Jan 20, 2025

View reviewed changes

src/diffusers/loaders/single_file_model.py Outdated Show resolved Hide resolved

sayakpaul reviewed Jan 20, 2025

View reviewed changes

sayakpaul mentioned this pull request Jan 20, 2025

FLUX error when loading with low_cpu_mem_usage=False and ignore_mismatched_sizes=True #9343

Open

hlky reviewed Jan 20, 2025

View reviewed changes

SunMarc and others added 3 commits January 20, 2025 16:23

Update src/diffusers/models/model_loading_utils.py

2f671af

Co-authored-by: Sayak Paul <[email protected]>

suggestions from reviews

7273a94

Merge remote-tracking branch 'upstream/model-loading-refactor' into m…

00f0bd1

…odel-loading-refactor

yiyixuxu reviewed Jan 21, 2025

View reviewed changes

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

yiyixuxu reviewed Jan 21, 2025

View reviewed changes

src/diffusers/models/modeling_utils.py Show resolved Hide resolved

yiyixuxu reviewed Jan 21, 2025

View reviewed changes

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

SunMarc and others added 3 commits January 21, 2025 10:56

Update src/diffusers/models/modeling_utils.py

c5da192

Co-authored-by: YiYi Xu <[email protected]>

remove safetensors check

039eef5

Merge remote-tracking branch 'upstream/model-loading-refactor' into m…

21f94a1

…odel-loading-refactor

DN6 reviewed Jan 21, 2025

View reviewed changes

yiyixuxu reviewed Jan 21, 2025

View reviewed changes

src/diffusers/models/modeling_utils.py Show resolved Hide resolved

SunMarc added 6 commits January 23, 2025 15:47

fix default value

337b2fc

Merge remote-tracking branch 'upstream/main' into model-loading-refactor

aedf6af

more fix from suggestions

0df7010

revert logic for single file

d3a7dc8

style

fc4af16

Merge remote-tracking branch 'upstream/main' into model-loading-refactor

18d61bb

sayakpaul reviewed Jan 27, 2025

View reviewed changes

Merge branch 'main' into model-loading-refactor

26228eb

typing + fix couple of issues

592c878

yiyixuxu mentioned this pull request Jan 29, 2025

possibly to avoid from_single_file loading in fp32 to save RAM #10679

Open

a-r-r-o-w reviewed Jan 30, 2025

View reviewed changes

src/diffusers/models/modeling_utils.py Show resolved Hide resolved

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

SunMarc and others added 6 commits February 4, 2025 18:51

improve speed

31c7d95

Merge remote-tracking branch 'upstream/model-loading-refactor' into m…

1634362

…odel-loading-refactor

Update src/diffusers/models/modeling_utils.py

37d59f6

Co-authored-by: Aryan <[email protected]>

Merge remote-tracking branch 'upstream/main' into model-loading-refactor

b48fedc

fp8 dtype

2753abe

add tests

abd3a91

SunMarc requested review from sayakpaul, yiyixuxu and DN6 February 14, 2025 16:50

SunMarc added 2 commits February 14, 2025 17:53

rename resolved_archive_file to resolved_model_file

d1c4a61

format

bcbd493

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[FEAT] Model loading refactor #10604

[FEAT] Model loading refactor #10604

SunMarc commented Jan 17, 2025 •

edited

Loading

SunMarc commented Jan 18, 2025

sayakpaul left a comment

sayakpaul commented Jan 20, 2025

hlky Jan 20, 2025

SunMarc Jan 20, 2025

DN6 Jan 21, 2025

sayakpaul Jan 27, 2025

SunMarc Jan 27, 2025 •

edited

Loading

sayakpaul left a comment •

edited

Loading

sayakpaul Jan 27, 2025

SunMarc Jan 27, 2025

sayakpaul Jan 27, 2025

sayakpaul Jan 27, 2025

HuggingFaceDocBuilderDev commented Jan 27, 2025

SunMarc commented Feb 14, 2025

[FEAT] Model loading refactor #10604

Are you sure you want to change the base?

[FEAT] Model loading refactor #10604

Conversation

SunMarc commented Jan 17, 2025 • edited Loading

What does this PR do?

SunMarc commented Jan 18, 2025

sayakpaul left a comment

Choose a reason for hiding this comment

sayakpaul commented Jan 20, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

SunMarc Jan 27, 2025 • edited Loading

Choose a reason for hiding this comment

sayakpaul left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

HuggingFaceDocBuilderDev commented Jan 27, 2025

SunMarc commented Feb 14, 2025

SunMarc commented Jan 17, 2025 •

edited

Loading

SunMarc Jan 27, 2025 •

edited

Loading

sayakpaul left a comment •

edited

Loading