Repository: huggingface/transformers
Pull request: [device_map] Fix device_map computation by correctly adjusting memory available

# What does this PR do?

As per the title. The unused memory is taken into account too late, which leads to different device_map for the same hardware and models, and even random cuda OOM!! Basically, the max memory needs to be adjusted BEFORE being used in `get_balanced_memory`. Consider the following example:

```python
from transformers import AutoModelForCausalLM
import gc

model_id = "mistralai/Mixtral-8x7B-v0.1"
device = "auto"

for _ in range(5):
    model = AutoModelForCausalLM.from_pretrained(model_id, dtype="auto", device_map=device)
    del model
    gc.collect()
```

On main, we have the following value computed at each iteration:

```sh
>>> Inferred max memory: {0: '14.25 GiB', 1: '14.25 GiB', 2: '14.25 GiB', 3: '14.25 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.69 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 1, 'model.layers.5': 1, 'model.layers.6': 1, 'model.layers.7': 1, 'model.layers.8': 1, 'model.layers.9': 2, 'model.layers.10': 2, 'model.layers.11': 2, 'model.layers.12': 2, 'model.layers.13': 2, 'model.layers.14': 3, 'model.layers.15': 3, 'model.layers.16': 3, 'model.layers.17': 3, 'model.layers.18': 3, 'model.layers.19': 4, 'model.layers.20': 4, 'model.layers.21': 4, 'model.layers.22': 4, 'model.layers.23': 4, 'model.layers.24': 5, 'model.layers.25': 5, 'model.layers.26': 5, 'model.layers.27': 5, 'model.layers.28': 5, 'model.layers.29': 6, 'model.layers.30': 6, 'model.layers.31': 6, 'model.norm': 6, 'model.rotary_emb': 6, 'lm_head': 6})

>>> Inferred max memory: {0: '29.69 GiB', 1: '32.15 GiB', 2: '32.15 GiB', 3: '32.15 GiB', 4: '32.15 GiB', 5: '32.15 GiB', 6: '26.98 GiB', 7: '78.92 GiB', 'cpu': '1748.62 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 0, 'model.layers.5': 0, 'model.layers.6': 0, 'model.layers.7': 0, 'model.layers.8': 0, 'model.layers.9': 1, 'model.layers.10': 1, 'model.layers.11': 1, 'model.layers.12': 1, 'model.layers.13': 1, 'model.layers.14': 1, 'model.layers.15': 1, 'model.layers.16': 1, 'model.layers.17': 1, 'model.layers.18': 1, 'model.layers.19': 1, 'model.layers.20': 2, 'model.layers.21': 2, 'model.layers.22': 2, 'model.layers.23': 2, 'model.layers.24': 2, 'model.layers.25': 2, 'model.layers.26': 2, 'model.layers.27': 2, 'model.layers.28': 2, 'model.layers.29': 2, 'model.layers.30': 2, 'model.layers.31': 3, 'model.norm': 3, 'model.rotary_emb': 3, 'lm_head': 3})

>>> Inferred max memory: {0: '45.12 GiB', 1: '50.04 GiB', 2: '50.04 GiB', 3: '28.65 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.62 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 0, 'model.layers.5': 0, 'model.layers.6': 0, 'model.layers.7': 0, 'model.layers.8': 0, 'model.layers.9': 0, 'model.layers.10': 0, 'model.layers.11': 0, 'model.layers.12': 0, 'model.layers.13': 0, 'model.layers.14': 0, 'model.layers.15': 1, 'model.layers.16': 1, 'model.layers.17': 1, 'model.layers.18': 1, 'model.layers.19': 1, 'model.layers.20': 1, 'model.layers.21': 1, 'model.layers.22': 1, 'model.layers.23': 1, 'model.layers.24': 1, 'model.layers.25': 1, 'model.layers.26': 1, 'model.layers.27': 1, 'model.layers.28': 1, 'model.layers.29': 1, 'model.layers.30': 1, 'model.layers.31': 1, 'model.norm': 1, 'model.rotary_emb': 1, 'lm_head': 1})

>>> Inferred max memory: {0: '75.99 GiB', 1: '78.92 GiB', 2: '14.25 GiB', 3: '14.25 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.59 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 0, 'model.layers.5': 0, 'model.layers.6': 0, 'model.layers.7': 0, 'model.layers.8': 0, 'model.layers.9': 0, 'model.layers.10': 0, 'model.layers.11': 0, 'model.layers.12': 0, 'model.layers.13': 0, 'model.layers.14': 0, 'model.layers.15': 0, 'model.layers.16': 0, 'model.layers.17': 0, 'model.layers.18': 0, 'model.layers.19': 0, 'model.layers.20': 0, 'model.layers.21': 0, 'model.layers.22': 0, 'model.layers.23': 0, 'model.layers.24': 0, 'model.layers.25': 0, 'model.layers.26': 0, 'model.layers.27': 1, 'model.layers.28': 1, 'model.layers.29': 1, 'model.layers.30': 1, 'model.layers.31': 1, 'model.norm': 1, 'model.rotary_emb': 1, 'lm_head': 1})

>>> Inferred max memory: {0: '78.92 GiB', 1: '32.15 GiB', 2: '14.25 GiB', 3: '14.25 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.60 GiB'}
>>>Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 0, 'model.layers.5': 0, 'model.layers.6': 0, 'model.layers.7': 0, 'model.layers.8': 0, 'model.layers.9': 0, 'model.layers.10': 0, 'model.layers.11': 0, 'model.layers.12': 0, 'model.layers.13': 0, 'model.layers.14': 0, 'model.layers.15': 0, 'model.layers.16': 0, 'model.layers.17': 0, 'model.layers.18': 0, 'model.layers.19': 0, 'model.layers.20': 0, 'model.layers.21': 0, 'model.layers.22': 0, 'model.layers.23': 0, 'model.layers.24': 0, 'model.layers.25': 0, 'model.layers.26': 0, 'model.layers.27': 0, 'model.layers.28': 1, 'model.layers.29': 1, 'model.layers.30': 1, 'model.layers.31': 1, 'model.norm': 1, 'model.rotary_emb': 1, 'lm_head': 1})
>>> OutOfMemoryError: CUDA out of memory. Tried to allocate 1.75 GiB. GPU 0 has a total capacity of 79.44 GiB of which 572.69 MiB is free. Including non-PyTorch memory, this process has 78.87 GiB memory in use. Of the allocated memory 75.93 GiB is allocated by PyTorch, and 2.43 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
```

As you can see, because *unused* gpu memory is adjusted too late in the `max_memory`, the device_map keeps changing at each iteration, and putting more and more layers on GPU 0!! To the point that it makes GPU 0 OOM in the last iteration!!!

Now, on this PR, you can see that everything is solved and each iteration computes the same values:

```sh
>>> Inferred max memory: {0: '14.25 GiB', 1: '14.25 GiB', 2: '14.25 GiB', 3: '14.25 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.76 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 1, 'model.layers.5': 1, 'model.layers.6': 1, 'model.layers.7': 1, 'model.layers.8': 1, 'model.layers.9': 2, 'model.layers.10': 2, 'model.layers.11': 2, 'model.layers.12': 2, 'model.layers.13': 2, 'model.layers.14': 3, 'model.layers.15': 3, 'model.layers.16': 3, 'model.layers.17': 3, 'model.layers.18': 3, 'model.layers.19': 4, 'model.layers.20': 4, 'model.layers.21': 4, 'model.layers.22': 4, 'model.layers.23': 4, 'model.layers.24': 5, 'model.layers.25': 5, 'model.layers.26': 5, 'model.layers.27': 5, 'model.layers.28': 5, 'model.layers.29': 6, 'model.layers.30': 6, 'model.layers.31': 6, 'model.norm': 6, 'model.rotary_emb': 6, 'lm_head': 6})

>>> Inferred max memory: {0: '14.25 GiB', 1: '14.25 GiB', 2: '14.25 GiB', 3: '14.25 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.76 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 1, 'model.layers.5': 1, 'model.layers.6': 1, 'model.layers.7': 1, 'model.layers.8': 1, 'model.layers.9': 2, 'model.layers.10': 2, 'model.layers.11': 2, 'model.layers.12': 2, 'model.layers.13': 2, 'model.layers.14': 3, 'model.layers.15': 3, 'model.layers.16': 3, 'model.layers.17': 3, 'model.layers.18': 3, 'model.layers.19': 4, 'model.layers.20': 4, 'model.layers.21': 4, 'model.layers.22': 4, 'model.layers.23': 4, 'model.layers.24': 5, 'model.layers.25': 5, 'model.layers.26': 5, 'model.layers.27': 5, 'model.layers.28': 5, 'model.layers.29': 6, 'model.layers.30': 6, 'model.layers.31': 6, 'model.norm': 6, 'model.rotary_emb': 6, 'lm_head': 6})

>>> Inferred max memory: {0: '14.25 GiB', 1: '14.25 GiB', 2: '14.25 GiB', 3: '14.25 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.76 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 1, 'model.layers.5': 1, 'model.layers.6': 1, 'model.layers.7': 1, 'model.layers.8': 1, 'model.layers.9': 2, 'model.layers.10': 2, 'model.layers.11': 2, 'model.layers.12': 2, 'model.layers.13': 2, 'model.layers.14': 3, 'model.layers.15': 3, 'model.layers.16': 3, 'model.layers.17': 3, 'model.layers.18': 3, 'model.layers.19': 4, 'model.layers.20': 4, 'model.layers.21': 4, 'model.layers.22': 4, 'model.layers.23': 4, 'model.layers.24': 5, 'model.layers.25': 5, 'model.layers.26': 5, 'model.layers.27': 5, 'model.layers.28': 5, 'model.layers.29': 6, 'model.layers.30': 6, 'model.layers.31': 6, 'model.norm': 6, 'model.rotary_emb': 6, 'lm_head': 6})

>>> Inferred max memory: {0: '14.25 GiB', 1: '14.25 GiB', 2: '14.25 GiB', 3: '14.25 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.75 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 1, 'model.layers.5': 1, 'model.layers.6': 1, 'model.layers.7': 1, 'model.layers.8': 1, 'model.layers.9': 2, 'model.layers.10': 2, 'model.layers.11': 2, 'model.layers.12': 2, 'model.layers.13': 2, 'model.layers.14': 3, 'model.layers.15': 3, 'model.layers.16': 3, 'model.layers.17': 3, 'model.layers.18': 3, 'model.layers.19': 4, 'model.layers.20': 4, 'model.layers.21': 4, 'model.layers.22': 4, 'model.layers.23': 4, 'model.layers.24': 5, 'model.layers.25': 5, 'model.layers.26': 5, 'model.layers.27': 5, 'model.layers.28': 5, 'model.layers.29': 6, 'model.layers.30': 6, 'model.layers.31': 6, 'model.norm': 6, 'model.rotary_emb': 6, 'lm_head': 6})

>>> Inferred max memory: {0: '14.25 GiB', 1: '14.25 GiB', 2: '14.25 GiB', 3: '14.25 GiB', 4: '14.25 GiB', 5: '14.25 GiB', 6: '14.25 GiB', 7: '78.92 GiB', 'cpu': '1748.77 GiB'}
>>> Inferred device_map: OrderedDict({'model.embed_tokens': 0, 'model.layers.0': 0, 'model.layers.1': 0, 'model.layers.2': 0, 'model.layers.3': 0, 'model.layers.4': 1, 'model.layers.5': 1, 'model.layers.6': 1, 'model.layers.7': 1, 'model.layers.8': 1, 'model.layers.9': 2, 'model.layers.10': 2, 'model.layers.11': 2, 'model.layers.12': 2, 'model.layers.13': 2, 'model.layers.14': 3, 'model.layers.15': 3, 'model.layers.16': 3, 'model.layers.17': 3, 'model.layers.18': 3, 'model.layers.19': 4, 'model.layers.20': 4, 'model.layers.21': 4, 'model.layers.22': 4, 'model.layers.23': 4, 'model.layers.24': 5, 'model.layers.25': 5, 'model.layers.26': 5, 'model.layers.27': 5, 'model.layers.28': 5, 'model.layers.29': 6, 'model.layers.30': 6, 'model.layers.31': 6, 'model.norm': 6, 'model.rotary_emb': 6, 'lm_head': 6})
```


Changed files:
filename: src/transformers/integrations/accelerate.py
patch: @@ -43,7 +43,7 @@
 
 if is_accelerate_available():
     from accelerate import dispatch_model
-    from accelerate.utils import get_max_memory
+    from accelerate.utils import get_max_memory as accelerate_max_memory
     from accelerate.utils.modeling import clean_device_map, get_max_layer_size
 
 if TYPE_CHECKING:
@@ -196,6 +196,41 @@ def compute_module_total_buffer_size(model: nn.Module, hf_quantizer: "HfQuantize
     return module_sizes.get("", 0)
 
 
+def get_max_memory(max_memory: dict[int | str, int | str] | None = None):
+    """
+    Get the maximum memory available if nothing is passed, converts string to int otherwise.
+    Note: we need to overwrite this as accelerate does not take into account torch allocated but unused device memory...
+    """
+    # Get the max memory (it only uses free gpu memory, not torch allocated but free memory...)
+    final_max_memory = accelerate_max_memory(max_memory)
+
+    # Adjust for allocated but free memory
+    for device_name in final_max_memory:
+        if isinstance(device_name, int):  # it's a GPU device
+            # Only cuda and xpu use caching memory allocator
+            if is_torch_xpu_available():
+                unused_memory = torch.xpu.memory_reserved(device_name) - torch.xpu.memory_allocated(device_name)
+            elif torch.cuda.is_available():
+                unused_memory = torch.cuda.memory_reserved(device_name) - torch.cuda.memory_allocated(device_name)
+            else:
+                unused_memory = 0
+            # Add the pre-allocated but unused device memory
+            final_max_memory[device_name] += unused_memory
+        # Still respect the `max_memory` passed by the user if any
+        if max_memory is not None and device_name in max_memory:
+            final_max_memory[device_name] = min(max_memory[device_name], final_max_memory[device_name])
+
+    # If the user does not provide `max_memory`, accelerate sets the WHOLE cpu available memory as available.
+    # This is unwanted, as we don't want to set extremely tight bound and pressure for cpu if we are memory-constrained,
+    # especially if the model uses WeightConverter (because there will be some uncontrollable cpu memory spikes during
+    # the conversions before we resave the weights). In those cases, it's better to offload to disk a bit more
+    # if we were in-between, as otherwise we blow-up cpu memory
+    if max_memory is None and "cpu" in final_max_memory:
+        final_max_memory["cpu"] *= 0.90
+
+    return final_max_memory
+
+
 def get_balanced_memory(
     model: "PreTrainedModel",
     max_memory: dict[int | str, int | str] | None = None,
@@ -324,30 +359,9 @@ def _get_device_map(
         else:
             inferred_max_memory = get_max_memory(max_memory)
 
-        # If the user does not provide `max_memory`, accelerate sets the WHOLE cpu available memory as available.
-        # This is unwanted, as we don't want to set extremely tight bound and pressure for cpu if we are memory-constrained,
-        # especially if the model uses WeightConverter (because there will be some uncontrollable cpu memory spikes during
-        # the conversions before we resave the weights). In those cases, it's better to offload to disk a bit more
-        # if we were in-between, as otherwise we blow-up cpu memory
-        if max_memory is None and "cpu" in inferred_max_memory:
-            inferred_max_memory["cpu"] *= 0.90
-
         if hf_quantizer is not None:
             inferred_max_memory = hf_quantizer.adjust_max_memory(inferred_max_memory)
 
-        # `inferred_max_memory` contains non-reserved memory. There may be *unused* reserved memory in the GPU,
-        # which we can use to allocate parameters.
-        for device_name in inferred_max_memory:
-            if isinstance(device_name, int):  # it's a GPU device
-                if is_torch_xpu_available():
-                    unused_memory = torch.xpu.memory_reserved(device_name) - torch.xpu.memory_allocated(device_name)
-                else:
-                    unused_memory = torch.cuda.memory_reserved(device_name) - torch.cuda.memory_allocated(device_name)
-                inferred_max_memory[device_name] += unused_memory
-            # respect the `max_memory` passed by the user
-            if max_memory is not None and device_name in max_memory:
-                inferred_max_memory[device_name] = min(inferred_max_memory[device_name], max_memory[device_name])
-
         device_map = infer_auto_device_map(
             model,
             max_memory=inferred_max_memory,