Repository: huggingface/accelerate

Issue:
`get_balanced_memory` raises a `ZeroDivisionError` when `max_memory` of the single GPU available is set to 0
### System Info

```Shell
- `Accelerate` version: 0.22.0
- Platform: Linux-6.2.0-1009-aws-x86_64-with-glibc2.17
- Python version: 3.8.17
- Numpy version: 1.24.4
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- System RAM: 62.01 GB
- GPU type: Tesla T4
- `Accelerate` default config:
        Not found
```


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] 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

On a machine with a single GPU available, run:

```python
import psutil

from transformers import AutoModel

flant5 = AutoModel.from_pretrained(
    "google/flan-t5-xl",
    device_map="auto",
    max_memory={
        0:0, 
        "cpu": psutil.virtual_memory().available
    }
)
```

### Output
```
Traceback (most recent call last):
  File "test.py", line 5, in <module>
    flant5 = AutoModel.from_pretrained(
  File "/home/ubuntu/miniconda3/envs/venv/lib/python3.8/site-packages/transformers/models/auto/auto_factory.py", line 516, in from_pretrained
    return model_class.from_pretrained(
  File "/home/ubuntu/miniconda3/envs/venv/lib/python3.8/site-packages/transformers/modeling_utils.py", line 3012, in from_pretrained
    max_memory = get_balanced_memory(
  File "/home/ubuntu/miniconda3/envs/venv/lib/python3.8/site-packages/accelerate/utils/modeling.py", line 802, in get_balanced_memory
    per_gpu = module_sizes[""] // (num_devices - 1 if low_zero else num_devices)
ZeroDivisionError: integer division or modulo by zero
```

### Expected behavior

`get_balanced_memory` should not raise a `ZeroDivisionError` when `max_memory` configurations renders the `num_devices` to be 0. A check should be made if `num_devices` is 0 and immediately return the `max_memory` dictionary without attempting to divide the modules sizes on the GPUs.


Discussion:
cc @SunMarc 

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
@@ -766,9 +766,6 @@ def get_balanced_memory(
     user_not_set_max_memory = max_memory is None
     max_memory = get_max_memory(max_memory)
 
-    if not (torch.cuda.is_available() or is_xpu_available()) or is_mps_available():
-        return max_memory
-
     if not is_xpu_available():
         num_devices = len([d for d in max_memory if torch.device(d).type == "cuda" and max_memory[d] > 0])
     else:
@@ -784,6 +781,9 @@ def get_balanced_memory(
             ]
         )
 
+    if num_devices == 0:
+        return max_memory
+
     if num_devices == 1:
         # We cannot do low_zero on just one GPU, but we will still reserve some memory for the buffer
         low_zero = False


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
@@ -547,6 +547,10 @@ def test_get_balanced_memory(self):
         max_memory = get_balanced_memory(model, max_memory={0: 0, 1: 300, 2: 300})
         self.assertDictEqual({0: 0, 1: 215, 2: 300}, max_memory)
 
+        # If we set a device to 0, it's not counted.
+        max_memory = get_balanced_memory(model, max_memory={0: 0, "cpu": 100})
+        self.assertDictEqual({0: 0, "cpu": 100}, max_memory)
+
     @require_cuda
     @require_safetensors
     def test_load_state_dict(self):