[`pipeline`] A simple fix for half-precision & 8bit models 
Author: younesbelkada
Pull request: yes

# What does this PR do?

Currently on the `main` branch of `transformers` if a user wants to run a `pipeline` using large models (thus, ideally loaded with `device_map=...`) and in half precision (or in int8), they may encounter some issues when calling `pipeline` with `top_p` & `top_k` sampling:
```bash
RuntimeError: "topk_cpu" not implemented for 'Half'
```

## Snippet to reproduce & explanations:

```python
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

model_id = "EleutherAI/gpt-neo-125M"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=torch.float16)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=20, temperature=1, do_sample=True, top_p=0.95, top_k=60, num_return_sequences=3)
text = "What can you tell me about the LHC?"
response = pipe(text)
print(response[0]["generated_text"])
```

This is because the `input_ids` are automatically set on `cpu` since the argument `device` is not passed when initializing the `pipeline`. A model that is loaded with `device_map=...` (i.e. with `accelerate`) always sets the output tensor of the model to the `device` of the input tensor thanks to the forward hooks. Therefore when calling the top_k method, the output tensor is in fp16 (because the model has been loaded in fp16) & on `cpu` hence the torch error above.

Currently a hack to fix this is to add `device=0` when initializing the `pipeline` but this leads to inconsistent and undesirable behaviours for some cases, for example when loading large models in several GPUs, since the call `model.to(self.device)` will break some internals (the hooks will be still there but the weights will be set on the wrong devices). A snippet to reproduce below:

```python
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

model_id = "EleutherAI/gpt-neo-125M"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="balanced", torch_dtype=torch.float16)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=20, temperature=1, do_sample=True, top_p=0.95, top_k=60, num_return_sequences=3, device=0)
text = "What can you tell me about the LHC?"
response = pipe(text)
print(response[0]["generated_text"])
```

adding this hack also breaks the usage of `pipeline` with int8 models, since the `to` method is blocked for these models:

```bash
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`.
```

Thus, I propose to fix this by simply checking whether a model has been loaded with `accelerate` by looking at the attribute `hf_device_map` , and set the model on the correct device only if it has not been loaded with accelerate as backend. This fixes 3 bugs: using `pipeline` with a fp16 model that has been loaded with `accelerate` without having any error in case of multi-gpu usage, using `pipeline` with a fp16 model w `accelerate` & sampling strategies, and using `pipeline` with int8 models & sampling strategies.

cc @sgugger @Narsil

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Thanks for the well thought out issue and proposed fix. 

I don't particularly like the fix because it depends on some weird internal and still forces users to use `device_map` and `device` iiuc.

Couldn't we just use `device_map` and use `accelerate` api to figure out where to put the inputs? (most likely cuda:0 but still cpu if no gpus are available I think.)
That or just do something special for `device_map` without asking where the model is (if the API doesn't exist or is tricky).

Imo using `device_map` and `device` should be an error (ambiguous intent) 

Thanks for the feedback @Narsil !
I think 
> Imo using device_map and device should be an error (ambiguous intent)

Makes sense !
Another fix would be to force-upcast the logits in fp32 when doing top_k & top_p sampling on the generation side only if the logits are on `cpu`, is this solution a reasonable fix @gante ? Happy to open a PR to fix it!

> force-upcast 

I would highly advise against it too. There's a limit to magic. Doing half precision on cpu should crash in a lot of places. We shouldn't upcast on behalf of a user that explicitely asked for half precision imo. That's breaking user intent.
But the user also asked for GPU, that's where we're breaking his intent and that's what should be fixed IMO.

Does `accelerate` allow to know on which device is the start of the model ?

I see, makes sense! 

> Does accelerate allow to know on which device is the start of the model ?

I am not sure here, maybe @sgugger & @muellerzr knows better

> I am not sure here, maybe @sgugger & @muellerzr knows better

if not the pipeline could have the simplest heuristic `'cuda:0' if torch.cuda.is_avalaible() else 'cpu'` which should work 99% of the time. 
But it wouldn't if a user specified an odd map (which is why having direct access would be better).

What do you think:

```diff
diff --git a/src/transformers/pipelines/base.py b/src/transformers/pipelines/base.py
index 30402b36e..e698f1aa3 100644
--- a/src/transformers/pipelines/base.py
+++ b/src/transformers/pipelines/base.py
@@ -749,7 +749,7 @@ class Pipeline(_ScikitCompat):
         framework: Optional[str] = None,
         task: str = "",
         args_parser: ArgumentHandler = None,
-        device: Union[int, str, "torch.device"] = -1,
+        device: Union[int, str, "torch.device"] = None,
         torch_dtype: Optional[Union[str, "torch.dtype"]] = None,
         binary_output: bool = False,
         **kwargs,
@@ -764,6 +764,20 @@ class Pipeline(_ScikitCompat):
         self.image_processor = image_processor
         self.modelcard = modelcard
         self.framework = framework
+
+        # Special handling
+        if self.framework == "pt" and device is not None:
+            self.model = self.model.to(device=device)
+
+        if device is None:
+            # `accelerate` device map
+            hf_device_map = getattr(self.model, "hf_device_map", None)
+            if hf_device_map is not None:
+                # Take the first device used by `accelerate`.
+                device = next(iter(hf_device_map.values()))
+            else:
+                device = -1
+
         if is_torch_available() and self.framework == "pt":
             if isinstance(device, torch.device):
                 self.device = device
@@ -775,13 +789,10 @@ class Pipeline(_ScikitCompat):
                 self.device = torch.device(f"cuda:{device}")
         else:
             self.device = device
+
         self.torch_dtype = torch_dtype
         self.binary_output = binary_output

-        # Special handling
-        if self.framework == "pt" and self.device.type != "cpu":
-            self.model = self.model.to(self.device)
-
         # Update config with task specific parameters
         task_specific_params = self.model.config.task_specific_params
         if task_specific_params is not None and task in task_specific_params:
```

Here we just modify the default device **when** the model uses `accelerate` 's `device_map`.
We still depend on something magic, but it only modifies the `default` device, and doesn't modify `model` unless `device` was specified by user (which is correct in terms of intent IMO)

I think that would totally work @Narsil ! Happy to change the PR with your proposed changes, let me know!

Sure. Let's update the doc too.

Side notes:
- you should probably do something if the user passes a device and the model has a `hf_device_map` (at least a warning) as the line `model.to(device)` will probably screw things up (it will at least error if there are some weights offloaded)
- the device on which the model is executed is determined by this rule in Accelerate, maybe you should use the same code? I can also store it on the Accelerate side in a special attribute but then you'd have to wait for a release.
```py
    if set(device_map.values()) == {"cpu"} or set(device_map.values()) == {"cpu", "disk"}:
        main_device = "cpu"
    else:
        main_device = [d for d in device_map.values() if d not in ["cpu", "disk"]][0]
```

Thanks a lot for the valuable feedback @Narsil @sgugger ! 
I updated the PR and added more clarification (and also a new section) on the docs

the failing test is indenpendent to our PR! Merging!
Thanks for all your comments!