Repository: huggingface/accelerate

Issue:
Question regarding `device_map="auto"`
In `diffusers` we essentially want to offload all the maintenance of efficient model loading (meaning using torch's meta tensors to not allocate any memory at init). 

We'd like to do so via just one logic/kwarg being `device_map="auto"`. This means we would like to **always** advocate to pass `device_map="auto"` stating that this is always in every environment the recommend way of loading a modedl.
This works great on GPU, but when loading the model on CPU with `from_pretrained(...., device_map="auto")` we're getting: 
```
tests/models/test_models_unet_2d.py::UNetLDMModelTests::test_from_pretrained_hub - IndexError: list index out of range
```

=> Is this expected? 


Discussion:
Reproducible:

```python
!pip install git+https://github.com/huggingface/accelerate.git
!pip install git+https://github.com/huggingface/diffusers.git

from diffusers import UNet2DModel

UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256", device_map="auto")
```

=> this works well on GPU but not on CPU. Intended or bug?
It seems it's only intended for GPU for now. In case we could allow `device_map="auto"` functionality also for CPU this could be quite helpful to speed up testing and overall UX in `diffusers`: see: https://github.com/huggingface/diffusers/pull/1018

Resolving patch:
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
--- a/src/accelerate/big_modeling.py
+++ b/src/accelerate/big_modeling.py
@@ -245,11 +245,15 @@ def dispatch_model(
     check_device_map(model, device_map)
 
     if main_device is None:
-        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
+        if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
+            main_device = "cpu"
+        else:
+            main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
 
-    cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
-    if state_dict is None and len(cpu_modules) > 0:
-        state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
+    if main_device != "cpu":
+        cpu_modules = [name for name, device in device_map.items() if device == "cpu"]
+        if state_dict is None and len(cpu_modules) > 0:
+            state_dict = extract_submodules_state_dict(model.state_dict(), cpu_modules)
 
     disk_modules = [name for name, device in device_map.items() if device == "disk"]
     if offload_dir is None and len(disk_modules) > 0:
@@ -266,7 +270,8 @@ def dispatch_model(
     execution_device = {
         name: main_device if device in ["cpu", "disk"] else device for name, device in device_map.items()
     }
-    offload = {name: device in ["cpu", "disk"] for name, device in device_map.items()}
+    offloaded_devices = ["disk"] if main_device == "cpu" else ["cpu", "disk"]
+    offload = {name: device in offloaded_devices for name, device in device_map.items()}
     save_folder = offload_dir if len(disk_modules) > 0 else None
     if state_dict is not None or save_folder is not None:
         weights_map = OffloadedWeightsLoader(state_dict=state_dict, save_folder=save_folder)


Test patch: