Repository: huggingface/accelerate

Issue:
infer_auto_device_map fails using a device_map from get_balanced_memory for certain model sizes
### System Info

```Shell
- `Accelerate` version: 0.23.0.dev0
- Platform: Linux-5.15.0-70-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Numpy version: 1.22.2
- PyTorch version (GPU?): 2.1.0a0+4136153 (True)
- PyTorch XPU available: False
- PyTorch NPU available: False
- GPU type: NVIDIA A100 80GB PCIe (system has 2 GPUs)
```


### 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

Run this script (`hf_auto_dispatch_issue.py`) using a system with 2 GPUs:
```python
import accelerate
import torch
import transformers
from accelerate.utils import get_balanced_memory, infer_auto_device_map
from torch import nn
from accelerate import dispatch_model

class SimpleNetwork(nn.Module):
    def __init__(self, model_size):
        super().__init__()
        self.fc1 = nn.Linear(10, 100)
        self.fc2 = nn.Linear(100, 5)
        self.fc3 = nn.Linear(5, model_size)

    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x

def dispatch(model):
    max_memory = get_balanced_memory(
                        model,
                        max_memory=None,
                        dtype=torch.float32,
                        low_zero=False
                    )

    print(f'balanced max_memory: {max_memory}')

    device_map = infer_auto_device_map(
        model, max_memory=max_memory, dtype=torch.float32, verbose=True
    )

    print(f'device map: {device_map}')
    return dispatch_model(model, device_map=device_map)

inp = torch.rand((1, 10))

# large model
print('---DISPATCHING LARGE MODEL---')
model_large = SimpleNetwork(model_size=200)
model_large_dispatched = dispatch(model_large)
model_large_dispatched(inp) # works

# small model
print('\n---DISPATCHING SMALL MODEL---')
model_small = SimpleNetwork(model_size=1)
model_small_dispatched = dispatch(model_small)
model_small_dispatched(inp) # crash
```

Program output:
```
---DISPATCHING LARGE MODEL---
balanced max_memory: {0: 10285, 1: 84551270400, 'cpu': 1066935697408}

Treating module fc1.
Putting fc1 (size=4400) on 0 (available=5485).

Treating module fc2.
Not enough space on 0 to put fc2 (space available 1085, module size 2020).
This module cannot be split, going to the next device.

Treating module fc2.
Putting fc2 (size=2020) on 1 (available=84551270400).

Treating module fc3.
Putting fc3 (size=4800) on 1 (available=84551268380).
device map: {'fc1': 0, 'fc2': 1, 'fc3': 1}

---DISPATCHING SMALL MODEL---
balanced max_memory: {0: 5907, 1: 84517715968, 'cpu': 1066919186432}

Treating module fc1.
Not enough space on 0 to put fc1 (space available 3887, module size 4400).
This module cannot be split, going to the next device.

Treating module fc1.
Putting fc1 (size=4400) on 1 (available=84517715968).

Treating module fc2.
Putting fc2 (size=2020) on 1 (available=84517711568).

Treating module fc3.
Putting fc3 (size=24) on 1 (available=84517709548).
device map: {'': 1}
Traceback (most recent call last):
  File "<redacted>/hf_auto_dispatch_issue.py", line 53, in <module>
    model_small_dispatched(inp) # crash
  File "/workspace/pytorch/torch/nn/modules/module.py", line 1505, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/workspace/pytorch/torch/nn/modules/module.py", line 1514, in _call_impl
    return forward_call(*args, **kwargs)
  File "<redacted>/hf_auto_dispatch_issue.py", line 19, in forward
    x = self.fc1(x)
  File "/workspace/pytorch/torch/nn/modules/module.py", line 1505, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/workspace/pytorch/torch/nn/modules/module.py", line 1514, in _call_impl
    return forward_call(*args, **kwargs)
  File "/workspace/pytorch/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:1 and cpu! (when checking argument for argument mat1 in method wrapper_CUDA_addmm)
```

### Expected behavior

In the above script, I create two models of the same architecture but different sizes and use `get_balanced_memory`, `infer_auto_device_map` and `dispatch_model` in order to distribute the model across two GPUs in a manner that balances memory usage.

For the larger model size, `infer_auto_device_map` is able to produce a device map that balances the three layers across the two GPU devices (`device map: {'fc1': 0, 'fc2': 1, 'fc3': 1}`), and I'm able to properly execute a forward pass of the model leveraging the HF hooks that are inserted into the model.

However, with the small model size, `infer_auto_device_map` puts all of the layers on GPU1 (`device map: {'': 1}`) rather than properly balancing the models across both GPUs (unexpected behavior). The result of this is that when `dispatch_model` is called with this device map, no HF hooks are inserted into the model due to the following check in `dispatch_model`:

https://github.com/huggingface/accelerate/blob/ba6f11ec3ea9bd385c18d5e4bcc1d8054c31f300/src/accelerate/big_modeling.py#L346-L349

Since the model is on one device only (GPU1), `dispatch_model` moves all layers of the model to GPU1 rather than inserting hooks. Since there are no hooks, when we attempt to call a forward pass of this model with a CPU tensor, we get the device mismatch error above.

There are a couple of issues raised from this experiment:
**Issue 1:**
For the small model, I would have expected that the combination of `get_balanced_memory` and `infer_auto_device_map`  would have produced a `device_map` that balanced the model across both GPUs as was successfully done for the larger version of the model. More specifically, for the small model, `get_balanced_memory` produces a device map that can fit `fc1` (size 4400) on GPU0 (`balanced max_memory: {0: 5907, 1: 84517715968, 'cpu': 1066919186432}`). However, `infer_auto_device_map` is unable to properly assign `fc1` to GPU0 even though there is room `max_memory` for it: `Not enough space on 0 to put fc1 (space available 3887, module size 4400)`. In this case the issue is that the reported available size is `3887`, whereas `max_memory` indicated `5907` available. From initial investigation, the reason for this available size discrepancy is due to the following code:
https://github.com/huggingface/accelerate/blob/ba6f11ec3ea9bd385c18d5e4bcc1d8054c31f300/src/accelerate/utils/modeling.py#L989-L995.
From this code snippet, when trying to fit `fc1` on GPU0, `current_max_size=5907` is reduced by `max_layer_size=2020` (`5907-2020=3887`). I currently don't understand why `current_max_size` is reduced by `2020` (which is the size of `fc1` not `fc0`) when trying to fit `fc0`. Since this happens, the `4400` size of `fc0` is now too big to fit on GPU0 and `infer_auto_device_map` moves to GPU1 which is the next device. This is what leads to all the layers of the network being assigned to GPU1.

**Issue 2**
From a user experience perspective, this script demonstrates how there is no guarantee that hooks will be assigned to the model. This is important because we see the forward pass either can succeed or fail based on the model architecture/size. Regardless of architecture/size, the user expects that they can rely on the hooks to move their input tensor to the right device.


Discussion:
cc @SunMarc 
Thanks for this detailed report ! Here are some answer to your questions: 
1) We decided to not add hook when the model fits inside one gpu from user feedback. Before, we always added hooks but it was confusing for the user that they could not move a model that was on one gpu. 
2) For the issue 1, this happens because the modules are very unbalanced in term of memory (the first two modules account for 90% of the memory) + not a lot of modules (only 3 in this case). We reduce the `current_max_size` by the `max_layer_size` of the remaining layers. In this case, we are trying to fit the first layer, so we don't take into account. 
It is done here:
```       
# Max size in the remaining layers may have changed since we took one, so we maybe update it.
max_layer_names = [n for n in max_layer_names if n != name and not n.startswith(name + ".")]
```
We do that because we need to keep space in case we offload the one of the remaining layers to the CPU. So, we need to have enough space in GPU to bring back that layer during inference. 

3) For the issue 2, this was initially designed to enable inference with big model. This approach has been battle tested with most models. For unbalanced models (first two modules takes most of the memory), I would recommend to just put the model on one device (device_map={"":0} ) or use device_map="sequential" as this mostly happens with small models that may have a big embedding layer compared to the other layers. Let me know if it makes sense. If you feel to do a PR if you have any idea on how to solve that :). 
@SunMarc Thanks for the response, this clarifies things! It would be nice to have the guarantee that the model will be partitioned across all devices when using this balanced approach, but I see that this likely isn't a big issue for larger models/models with more layers. For my workflow, I should be able to use `device_map=sequential` to work around this.

What are your thoughts on adding an additional option like `force_hooks=True/False` to `dispatch_model` and potentially `load_checkpoint_and_dispatch` that will attach hooks even if the `device_map` only contains one device? That way a user could ensure that they will have the benefits of hooks (being able to pass a tensor of any device into the model) and not have to worry about whether the partitioning system chose to put the model on 1 or >1 devices. In my example script above it would mean I wouldn't have to worry about the changing architecture size impacting whether I could rely on hooks. `force_hooks` could default to `False` to handle the case where the user expects they can move the model if it's just on one GPU.
Good idea ! Feel free to open a PR if you want to do it. Otherwise, I will do it when I have time

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
@@ -304,6 +304,7 @@ def dispatch_model(
     offload_buffers: bool = False,
     skip_keys: Optional[Union[str, List[str]]] = None,
     preload_module_classes: Optional[List[str]] = None,
+    force_hooks: bool = False,
 ):
     """
     Dispatches a model according to a given device map. Layers of the model might be spread across GPUs, offloaded on
@@ -334,6 +335,9 @@ def dispatch_model(
             of the forward. This should only be used for classes that have submodules which are registered but not
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        force_hooks (`bool`, *optional*, defaults to `False`):
+            Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a
+            single device.
     """
     # Error early if the device map is incomplete.
     check_device_map(model, device_map)
@@ -343,10 +347,11 @@ def dispatch_model(
         getattr(model, "is_quantized", False) or getattr(model, "is_loaded_in_8bit", False)
     ) and getattr(model, "quantization_method", "bitsandbytes") == "bitsandbytes"
 
-    # We attach hooks if the device_map have at least 2 different devices. Otherwise, the model in already loaded
+    # We attach hooks if the device_map has at least 2 different devices or if
+    # force_hooks is set to `True`. Otherwise, the model in already loaded
     # in the unique device and the user can decide where to dispatch the model.
     # If the model is quantized, we always force-dispatch the model
-    if (len(set(device_map.values())) > 1) or is_bnb_quantized:
+    if (len(set(device_map.values())) > 1) or is_bnb_quantized or force_hooks:
         if main_device is None:
             if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
                 main_device = "cpu"
@@ -439,6 +444,7 @@ def load_checkpoint_and_dispatch(
     offload_state_dict: Optional[bool] = None,
     skip_keys: Optional[Union[str, List[str]]] = None,
     preload_module_classes: Optional[List[str]] = None,
+    force_hooks: bool = False,
 ):
     """
     Loads a (potentially sharded) checkpoint inside a model, potentially sending weights to a given device as they are
@@ -481,6 +487,9 @@ def load_checkpoint_and_dispatch(
             of the forward. This should only be used for classes that have submodules which are registered but not
             called directly during the forward, for instance if a `dense` linear layer is registered, but at forward,
             `dense.weight` and `dense.bias` are used in some operations instead of calling `dense` directly.
+        force_hooks (`bool`, *optional*, defaults to `False`):
+            Whether or not to force device hooks to be attached to the model even if all layers are dispatched to a
+            single device.
 
     Example:
 
@@ -541,4 +550,5 @@ def load_checkpoint_and_dispatch(
         offload_buffers=offload_buffers,
         skip_keys=skip_keys,
         preload_module_classes=preload_module_classes,
+        force_hooks=force_hooks,
     )


Test patch:
diff --git a/tests/test_big_modeling.py b/tests/test_big_modeling.py
--- a/tests/test_big_modeling.py
+++ b/tests/test_big_modeling.py
@@ -444,6 +444,18 @@ def test_dispatch_model_with_unused_submodules_multi_gpu(self):
             output = model(x)
             self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
 
+    @require_cuda
+    def test_dispatch_model_force_hooks(self):
+        model = ModelForTest()
+        device_map = {"": 0}
+
+        x = torch.randn(2, 3)
+        expected = model(x)
+
+        dispatch_model(model, device_map, force_hooks=True)
+        output = model(x)
+        self.assertTrue(torch.allclose(expected, output.cpu(), atol=1e-5))
+
     @require_cuda
     def test_load_checkpoint_and_dispatch(self):
         model = ModelForTest()