Repository: huggingface/diffusers
Pull request: add accelerate to load models with smaller memory footprint
Closes #281 

Problem statement:
Optimize model loading by natively using Accelerate, please.
**Is your feature request related to a problem? Please describe.**
Loading pretrained models, for example `UNet2DConditionModel` peaks at 8GB when loading a pretrained one, while it takes a lot less memory when it is fully loaded to the GPU. That is happens likely because there are some redundancies between loaded weights and the model.

When we start it with empty weights and load them using accelerate, it peaks on 4.91GB, which lets us deploy Diffusers in servers that have a lot less RAM and thus are cheaper.

**Describe the solution you'd like**
It would be great if Diffusers natively supported loading models from pretrained using `accelerate.init_empty_weights` and `accelerate.accelerate.load_checkpoint_and_dispatch`, so it has lower memory footprint.

**Describe alternatives you've considered**
- Loading the models with accelerate myself, but that requires model surgery and a lot of coding, which makes me prone to mistakes
- Do what accelerate does by hand

**Additional context**
Using accelerate to natively load models would also reduce the RAM memory footprint when loading pipelines composed of more than one model, by already allocating them on the proper device and thus enabling us to deploy diffusers on cheaper servers.

Hints from issue discussion:
Hey @piEsposito, happy to look into that in the coming week! If you have a code snippet that works for you, feel free to include it here as a reference :hugs: 

@muellerzr @patrickvonplaten wdyt? We can either do this, or bring the low-RAM loading code from transformers
cc @sgugger so it's on your radar since next week I'll be out for a bit :) 
Happy to include `accelerate`! Maybe we can follow the same logic we have in `transformers` with `device_map="auto"` :-) And definitely +1 to make use of torch's meta device to half the peak required memory usage!
Now this is great news, you just made my life a lot easier. Anyway I can help? 
@piEsposito - if you want you could give it a try to add `accelerate` and `torch`'s `meta` device logic to: https://github.com/huggingface/diffusers/blob/9b704f768880c732503fe46800986c6998ca205d/src/diffusers/modeling_utils.py#L203 more or less in the same spirit as it's done in `transformers`: https://github.com/huggingface/transformers/blob/65fb71bc762c46bb067306c1fd083b1cba87a095/src/transformers/modeling_utils.py#L1610

Happy to guide you through a PR! 
@patrickvonplaten is something on that direction you were expecting? https://github.com/huggingface/diffusers/pull/361

I would be very happy if you could guide me through a PR.
Thanks a lot for opening the PR! Think @patil-suraj wanted to guide you here a bit :-) 
@patrickvonplaten it is just a draft to see if the direction is right. It will of course be improved to production level code with guidance from you folks. Thanks :D 
Hey @piEsposito ! Looking through the PR and will get back to you soon about this.
@patil-suraj I saw that when we pass the `low_cpu_mem_usage` and `device_map` to the Pipeline abstraction, it will break on stable diffusion when loading the Safety Checker. Should we add an exception for that and let safety checker be loaded in full?

Commit message: Merge branch 'main' of github.com:huggingface/diffusers into main

Patch to review:
diff --git a/setup.py b/setup.py
index fffb8f26f09b..a3fb11e23c46 100644
--- a/setup.py
+++ b/setup.py
@@ -97,6 +97,7 @@
     "tensorboard",
     "torch>=1.4",
     "transformers>=4.21.0",
+    "accelerate>=0.12.0"
 ]
 
 # this is a lookup table with items like:
diff --git a/src/diffusers/configuration_utils.py b/src/diffusers/configuration_utils.py
index fbe75f3f1441..4aeacca75edb 100644
--- a/src/diffusers/configuration_utils.py
+++ b/src/diffusers/configuration_utils.py
@@ -24,6 +24,7 @@
 
 from huggingface_hub import hf_hub_download
 from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError
+import accelerate
 from requests import HTTPError
 
 from . import __version__
@@ -153,11 +154,17 @@ def from_config(cls, pretrained_model_name_or_path: Union[str, os.PathLike], ret
         </Tip>
 
         """
-        config_dict = cls.get_config_dict(pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
-
-        init_dict, unused_kwargs = cls.extract_init_dict(config_dict, **kwargs)
-
-        model = cls(**init_dict)
+        device_map = kwargs.pop("low_cpu_mem_usage", None)
+        if device_map is not None:
+            with accelerate.init_empty_weights():
+                config_dict = cls.get_config_dict(pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
+                init_dict, unused_kwargs = cls.extract_init_dict(config_dict, **kwargs)
+                model = cls(**init_dict)
+                    
+        else:
+            config_dict = cls.get_config_dict(pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
+            init_dict, unused_kwargs = cls.extract_init_dict(config_dict, **kwargs)
+            model = cls(**init_dict)
 
         if return_unused_kwargs:
             return model, unused_kwargs
diff --git a/src/diffusers/modeling_utils.py b/src/diffusers/modeling_utils.py
index fb613614a878..81d73a8be36c 100644
--- a/src/diffusers/modeling_utils.py
+++ b/src/diffusers/modeling_utils.py
@@ -22,6 +22,7 @@
 
 from huggingface_hub import hf_hub_download
 from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError, RevisionNotFoundError
+import accelerate
 from requests import HTTPError
 
 from .utils import CONFIG_NAME, DIFFUSERS_CACHE, HUGGINGFACE_CO_RESOLVE_ENDPOINT, logging
@@ -259,6 +260,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
         from_auto_class = kwargs.pop("_from_auto", False)
         torch_dtype = kwargs.pop("torch_dtype", None)
         subfolder = kwargs.pop("subfolder", None)
+        device_map = kwargs.pop("device_map", None)
 
         user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class}
 
@@ -275,6 +277,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             use_auth_token=use_auth_token,
             revision=revision,
             subfolder=subfolder,
+            device_map=device_map,
             **kwargs,
         )
 
@@ -357,25 +360,41 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 )
 
             # restore default dtype
-        state_dict = load_state_dict(model_file)
-        model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
-            model,
-            state_dict,
-            model_file,
-            pretrained_model_name_or_path,
-            ignore_mismatched_sizes=ignore_mismatched_sizes,
-        )
 
-        # Set model in evaluation mode to deactivate DropOut modules by default
-        model.eval()
+        if device_map is not None:            
+            accelerate.load_checkpoint_and_dispatch(
+                model,
+                model_file,
+                device_map
+            )
+            loading_info = {
+                "missing_keys": [],
+                "unexpected_keys": [],
+                "mismatched_keys": [],
+                "error_msgs": [],
+            }
+
+        else:
+            
+            state_dict = load_state_dict(model_file)
+            model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
+                model,
+                state_dict,
+                model_file,
+                pretrained_model_name_or_path,
+                ignore_mismatched_sizes=ignore_mismatched_sizes,
+            )
 
-        if output_loading_info:
             loading_info = {
                 "missing_keys": missing_keys,
                 "unexpected_keys": unexpected_keys,
                 "mismatched_keys": mismatched_keys,
                 "error_msgs": error_msgs,
             }
+
+        # Set model in evaluation mode to deactivate DropOut modules by default
+        model.eval()
+        if output_loading_info:
             return model, loading_info
 
         return model
diff --git a/tests/test_models_unet.py b/tests/test_models_unet.py
index d9c4967b57ed..8ea25560c80c 100644
--- a/tests/test_models_unet.py
+++ b/tests/test_models_unet.py
@@ -133,6 +133,13 @@ def test_from_pretrained_hub(self):
 
         assert image is not None, "Make sure output is not None"
 
+    def test_from_pretrained_accelerate(self):
+        model, _ = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update", output_loading_info=True, device_map="auto")
+        model.to(torch_device)
+        image = model(**self.dummy_input).sample
+
+        assert image is not None, "Make sure output is not None"
+
     def test_output_pretrained(self):
         model = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update")
         model.eval()


Code review comments:
diff_hunk: @@ -153,11 +154,17 @@ def from_config(cls, pretrained_model_name_or_path: Union[str, os.PathLike], ret
         </Tip>
 
         """
-        config_dict = cls.get_config_dict(pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
-
-        init_dict, unused_kwargs = cls.extract_init_dict(config_dict, **kwargs)
-
-        model = cls(**init_dict)
+        device_map = kwargs.pop("low_cpu_mem_usage", None)
path: src/diffusers/configuration_utils.py
text: @user1:
Ideally we would like to try to keep `configuration_utils.py` framework and component independent. Could we maybe try to set:

```
with accelerate.init_empty_weights():
            model, unused_kwargs = cls.from_config(
            config_path,
            cache_dir=cache_dir,
            return_unused_kwargs=True,
            force_download=force_download,
            resume_download=resume_download,
            proxies=proxies,
            local_files_only=local_files_only,
            use_auth_token=use_auth_token,
            revision=revision,
            subfolder=subfolder,
            device_map=device_map,
            **kwargs,
        )
```

in `modeling_utils.py` instead? 

@author:
Just did it.

diff_hunk: @@ -357,25 +360,41 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                 )
 
             # restore default dtype
-        state_dict = load_state_dict(model_file)
-        model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
-            model,
-            state_dict,
-            model_file,
-            pretrained_model_name_or_path,
-            ignore_mismatched_sizes=ignore_mismatched_sizes,
-        )
 
-        # Set model in evaluation mode to deactivate DropOut modules by default
-        model.eval()
+        if device_map is not None:            
path: src/diffusers/modeling_utils.py
text: @user1:
If possible it would be very nice if all `accelerate` logic would only be added here.

@author:
Just did it. Had to put the model creation and checkpoint loading after grabbing the weight and config files to avoid splitting the accelerate logic into two.