Fix get_balanced_memory for MPS (#3464)

This also fixes a failure in test_get_balanced_memory:

```
assert {0: 215, 1: 300} == {0: 300, 1: 300}
[...]
tests/test_modeling_utils.py:871: AssertionError
```

Signed-off-by: Ihar Hrachyshka <[email protected]>

Author: Ihar Hrachyshka

src/accelerate/utils/modeling.py

@@ -970,6 +970,8 @@ def get_balanced_memory(
         expected_device_type = "xpu"
     elif is_hpu_available():
         expected_device_type = "hpu"
+    elif is_mps_available():
+        expected_device_type = "mps"
     else:
         expected_device_type = "cuda"
     num_devices = len([d for d in max_memory if torch.device(d).type == expected_device_type and max_memory[d] > 0])