Repository: huggingface/accelerate
Pull request: [`Accelerator`] Fix issue with 8bit models
# What does this PR do?

Fixes https://github.com/huggingface/accelerate/issues/1147

In theory it is not possible to fine-tune 8-bit models except if you use adapters that can be used only is a `PeftModel` is used for training (I also need to test the snippet below with a `PeftModel` to make sure this is relevant or not). EDIT: passing an 8bit `PeftModel` through `accelerator.prepare` seems to work fine.

But in some systems you can use the accelerator to load an 8bit model and use it out of the training scope (e.g. get the model's logits and use it in another model)

I am not sure if we should support 8-bit models using `Accelerator`, but if so, I propose the following changes in this PR
Happy also to revert the tests / `bnb` dependency 

To reproduce: 
```python
import torch
import torch.nn.functional as F

from transformers import AutoModelForCausalLM
from datasets import load_dataset
from accelerate import Accelerator

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neo-125m",
    load_in_8bit=True,
    device_map="balanced",
)

accelerator = Accelerator()


model = accelerator.prepare(model)

```

cc @sgugger 

Ran all the slow tests and got errors on DeepSpeed and FSDP tests but not sure if the failing is related to my PR

<img width="1166" alt="Screenshot 2023-03-06 at 18 24 12" src="https://user-images.githubusercontent.com/49240599/223184756-66e6fcbe-4179-444a-b26f-a6a70693cb17.png">

Problem statement:
Incompatible with 8bit models
### System Info

```Shell
If I load a model in 8bit, it has to be mapped automatically to gpus. However, when `accelerate` tries to load the model to several gpus, I get the following error.

`ValueError: `.to` is not supported for `8-bit` models. Please use the model as it is, since the model has already been set to the correct devices and casted to the correct `dtype`.`

How to fix this?
```


### 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

1.. Load model in 8bit

```
model = GPTJForCausalLM.from_pretrained(
    hyperparameters["model_checkpoint"],
    revision="float16",
    load_in_8bit=True,
    device_map="auto",
)
```

2. Run with pytorch accelerate
```
import torch
import torch.nn.functional as F
from datasets import load_dataset
from accelerate import Accelerator

accelerator = Accelerator()
optimizer = torch.optim.Adam(model.parameters())

dataset = load_dataset('my_dataset')
data = torch.utils.data.DataLoader(dataset, shuffle=True)

model, optimizer, data = accelerator.prepare(model, optimizer, data)

model.train()
for epoch in range(10):
    for source, targets in data:
         optimizer.zero_grad()
         output = model(source)
         loss = F.cross_entropy(output, targets)
         accelerator.backward(loss)
         optimizer.step()
```

### Expected behavior

```Shell
accelerate loads model to gpus.
```

Hints from issue discussion:
cc @younesbelkada Could you have the same test you did in the Trainer here in Accelerate?

Commit message: add bnb as dependency

Patch to review:
diff --git a/docker/accelerate-gpu/Dockerfile b/docker/accelerate-gpu/Dockerfile
index 44aa94b29fb..ae71607d299 100644
--- a/docker/accelerate-gpu/Dockerfile
+++ b/docker/accelerate-gpu/Dockerfile
@@ -25,6 +25,8 @@ RUN source activate accelerate && \
     git+https://github.com/huggingface/accelerate#egg=accelerate[testing,test_trackers] \
     --extra-index-url https://download.pytorch.org/whl/cu113
 
+RUN python3 -m pip install --no-cache-dir bitsandbytes
+
 # Stage 2
 FROM nvidia/cuda:11.2.2-cudnn8-devel-ubuntu20.04 AS build-image
 COPY --from=compile-image /opt/conda /opt/conda
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
index d5073bf6463..ec35ca4e85d 100644
--- a/src/accelerate/accelerator.py
+++ b/src/accelerate/accelerator.py
@@ -1119,8 +1119,27 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
         if device_placement is None:
             device_placement = self.device_placement and self.distributed_type != DistributedType.FSDP
         self._models.append(model)
-        if device_placement:
+        if getattr(model, "is_loaded_in_8bit", False):
+            model_devices = set(model.hf_device_map.values())
+            if len(model_devices) > 1:
+                raise ValueError(
+                    "You can't train a model that has been loaded in 8-bit precision on multiple devices."
+                )
+
+            current_device_index = list(model_devices)[0]
+            if torch.device(current_device_index) != self.device:
+                # if on the first device (GPU 0) we don't care
+                if (self.device.index is not None) or (current_device_index != 0):
+                    raise ValueError(
+                        "You can't train a model that has been loaded in 8-bit precision on a different device than the one "
+                        "you're training on. Make sure you loaded the model on the correct device using for example `device_map={'':torch.cuda.current_device()}"
+                    )
+
+            if "cpu" in model_devices or "disk" in model_devices:
+                raise ValueError("You can't train a model that has been loaded in 8-bit precision on CPU or disk.")
+        elif device_placement:
             model = model.to(self.device)
+
         if self.distributed_type == DistributedType.MULTI_GPU:
             if any(p.requires_grad for p in model.parameters()):
                 kwargs = self.ddp_handler.to_kwargs() if self.ddp_handler is not None else {}
diff --git a/tests/test_accelerator.py b/tests/test_accelerator.py
index 3af62c9b94d..9ca32e1c30e 100644
--- a/tests/test_accelerator.py
+++ b/tests/test_accelerator.py
@@ -6,8 +6,10 @@
 import torch
 from torch.utils.data import DataLoader, TensorDataset
 
+from accelerate import infer_auto_device_map, init_empty_weights
 from accelerate.accelerator import Accelerator
 from accelerate.state import PartialState
+from accelerate.test_utils import require_multi_gpu, slow
 from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
 from accelerate.utils import patch_environment
 
@@ -160,3 +162,57 @@ def load_config(models, input_dir):
 
             # mode.class_name is NOT loaded from config
             self.assertTrue(model.class_name != model.__class__.__name__)
+
+    @slow
+    def test_accelerator_bnb(self):
+        """Tests that the accelerator can be used with the BNB library."""
+        from transformers import AutoModelForCausalLM
+
+        model = AutoModelForCausalLM.from_pretrained(
+            "EleutherAI/gpt-neo-125m",
+            load_in_8bit=True,
+            device_map="auto",
+        )
+        accelerator = Accelerator()
+
+        # This should work
+        model = accelerator.prepare(model)
+
+        with init_empty_weights():
+            model = AutoModelForCausalLM.from_pretrained(
+                "EleutherAI/gpt-neo-125m",
+            )
+            device_map = infer_auto_device_map(model)
+            device_map["lm_head"] = "cpu"
+
+        model = AutoModelForCausalLM.from_pretrained(
+            "EleutherAI/gpt-neo-125m", device_map=device_map, load_in_8bit=True, llm_int8_enable_fp32_cpu_offload=True
+        )
+
+        # This should not work and get value error
+        with self.assertRaises(ValueError):
+            model = accelerator.prepare(model)
+
+    @slow
+    @require_multi_gpu
+    def test_accelerator_bnb_multi_gpu(self):
+        """Tests that the accelerator can be used with the BNB library."""
+        from transformers import AutoModelForCausalLM
+
+        with init_empty_weights():
+            model = AutoModelForCausalLM.from_pretrained(
+                "EleutherAI/gpt-neo-125m",
+            )
+            device_map = infer_auto_device_map(model)
+            device_map["lm_head"] = 1
+
+        model = AutoModelForCausalLM.from_pretrained(
+            "EleutherAI/gpt-neo-125m",
+            load_in_8bit=True,
+            device_map=device_map,
+        )
+        accelerator = Accelerator()
+
+        # This should not work and get value error
+        with self.assertRaises(ValueError):
+            _ = accelerator.prepare(model)


Code review comments:
diff_hunk: @@ -1119,8 +1119,27 @@ def prepare_model(self, model: torch.nn.Module, device_placement=None):
         if device_placement is None:
             device_placement = self.device_placement and self.distributed_type != DistributedType.FSDP
         self._models.append(model)
-        if device_placement:
+        if getattr(model, "is_loaded_in_8bit", False):
+            model_devices = set(model.hf_device_map.values())
+            if len(model_devices) > 1:
+                raise ValueError(
+                    "You can't train a model that has been loaded in 8-bit precision on multiple devices."
+                )
+
+            current_device_index = list(model_devices)[0]
+            if torch.device(current_device_index) != self.device:
+                # if on the first device (GPU 0) we don't care
+                if (self.device.index is not None) or (current_device_index != 0):
+                    raise ValueError(
+                        "You can't train a model that has been loaded in 8-bit precision on a different device than the one "
+                        "you're training on. Make sure you loaded the model on the correct device using for example `device_map={'':torch.cuda.current_device()}"
+                    )
+
+            if "cpu" in model_devices or "disk" in model_devices:
+                raise ValueError("You can't train a model that has been loaded in 8-bit precision on CPU or disk.")
path: src/accelerate/accelerator.py
text: @user1:
```suggestion
                raise ValueError("You can't train a model that has been loaded in 8-bit precision with CPU or disk offload.")
```