Repository: huggingface/accelerate

Issue:
device_map error
### System Info

```Shell
- `Accelerate` version: 0.12.0
- Platform: Linux-4.18.0-305.25.1.el8_4.x86_64-x86_64-with-glibc2.17
- Python version: 3.8.13
- Numpy version: 1.23.1
- PyTorch version (GPU?): 1.11.0a0+gitbc2c6ed (True)
- `Accelerate` default config:
      Not found
```


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [X] My own task or dataset (give details below)

### Reproduction

```
File "/net/llm-shared-nfs/nfs/mayank/transformers-bloom-inference/bloom-inference-server/models/hf_accelerate.py", line 33, in __init__
    self.model = AutoModelForCausalLM.from_pretrained(**kwargs)
  File "/net/llm-shared-nfs/nfs/yelkurdi/conda/miniconda3/envs/bloom-server/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py", line 463, in from_pretrained
    return model_class.from_pretrained(
  File "/net/llm-shared-nfs/nfs/yelkurdi/conda/miniconda3/envs/bloom-server/lib/python3.8/site-packages/transformers/modeling_utils.py", line 2324, in from_pretrained
    model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_pretrained_model(
  File "/net/llm-shared-nfs/nfs/yelkurdi/conda/miniconda3/envs/bloom-server/lib/python3.8/site-packages/transformers/modeling_utils.py", line 2386, in _load_pretrained_model
    raise ValueError(
ValueError: The current `device_map` had weights offloaded to the disk. Please provide an `offload_folder` for them.
```
This happend with max_memory = {0: '0GIB', 1: '51GIB', 2: '51GIB', 3: '51GIB', 4: '51GIB', 5: '51GIB', 6: '51GIB', 7: '51GIB'}
and device_map = "auto" in kwargs for BLOOM-176 B.

This issue was not seen in transformers 4.21.2

@sgugger fyi

### Expected behavior

```Shell
related: https://github.com/huggingface/transformers-bloom-inference/issues/17
```



Discussion:
Without seeing the code you run, it's hard to see the root. Are you sure you passed along the `torch_dtype`? I don't see any problem on my side when this argument is passed.

Also note that now you only need to pass `device_map="balanced_low_0"` and not `memory_map` to achieve the the same results (fill all the GPUs to the max and put as little as possible on the first one).
ok, ill try to replicate and share a minimal script
@sgugger 
max_memory_per_gpu = {
                0: "0GIB",
                1: "51GIB",
                2: "51GIB",
                3: "51GIB",
                4: "51GIB",
                5: "51GIB",
                6: "51GIB",
                7: "51GIB",
            }
and device_map = "auto" fails.

but your suggestion is working well.
You can close this issue if this doesn't need a fix.

Resolving patch:
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
--- a/src/accelerate/utils/modeling.py
+++ b/src/accelerate/utils/modeling.py
@@ -365,7 +365,7 @@ def get_balanced_memory(
     if not torch.cuda.is_available():
         return max_memory
 
-    num_devices = len([d for d in max_memory if torch.device(d).type == "cuda"])
+    num_devices = len([d for d in max_memory if torch.device(d).type == "cuda" and max_memory[d] > 0])
     module_sizes = compute_module_sizes(model, dtype=dtype)
     per_gpu = module_sizes[""] // (num_devices - 1 if low_zero else num_devices)
 
@@ -409,7 +409,7 @@ def get_balanced_memory(
 
     max_memory = get_max_memory(max_memory)
     # The last device is left with max_memory just in case the buffer is not enough.
-    for i in range(num_devices - 1):
+    for i in range(len(max_memory) - 1):
         max_memory[i] = min(0 if low_zero and i == 0 else per_gpu, max_memory[i])
 
     if low_zero:


Test patch:
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -375,3 +375,7 @@ def test_get_balanced_memory(self):
         # Last device always get max memory to give more buffer and avoid accidental CPU offload
         max_memory = get_balanced_memory(model, max_memory={0: 300, 1: 500})
         self.assertDictEqual({0: 215, 1: 500}, max_memory)
+
+        # If we set a device to 0, it's not counted.
+        max_memory = get_balanced_memory(model, max_memory={0: 0, 1: 300, 2: 300})
+        self.assertDictEqual({0: 0, 1: 215, 2: 300}, max_memory)