[Core] refactor model loading

[sayakpaul]

Currently, we have got two codepaths:

1. For non-sharded checkpoints we do:

   diffusers/src/diffusers/models/modeling_utils.py

   Line 855 in 047bf49

   unexpected_keys = load_model_dict_into_meta(

2. For sharded checkpoints we do:

   diffusers/src/diffusers/models/modeling_utils.py

   Line 886 in 047bf49

   accelerate.load_checkpoint_and_dispatch(

And then for the (bnb) quantized checkpoints, we merge a sharded checkpoint:

diffusers/src/diffusers/models/modeling_utils.py

Line 775 in 047bf49

model_file = _merge_sharded_checkpoints(sharded_ckpt_cached_folder, sharded_metadata)

Essentially, we shouldn't have to merge sharded checkpoints even if it's quantized.

This will also allow us to more generally use keep_module_in_fp32 for sharded checkpoints. Currently, we have this logic for casting a model (which is tested thoroughly):

diffusers/src/diffusers/models/modeling_utils.py

Line 997 in 43534a8

elif torch_dtype is not None and hf_quantizer is None and not use_keep_in_fp32_modules:

When using load_model_dict_into_meta(), we do consider keep_module_in_fp32:

diffusers/src/diffusers/models/model_loading_utils.py

Line 177 in 43534a8

keep_in_fp32_modules=None,

But since for sharded checkpoints, we use load_checkpoint_and_dispatch(), there is no way to pass keep_module_in_fp32:
https://huggingface.co/docs/accelerate/main/en/package_reference/big_modeling#accelerate.load_checkpoint_and_dispatch

As discussed with @SunMarc, it's better to uniformize this so that we don't have to maintain two different codepaths and rely completely on load_model_dict_into_meta(). Marc has kindly agreed to open a PR to attempt this (this could be done in a series of PRs if needed). But I will join if any help is needed.

[sayakpaul]

@huggingface/diffusers Marc has started working on this 🥳

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.