Add warning in `generate` & `device_map=auto` & half precision models
Author: younesbelkada
Pull request: yes

# What does this PR do?

This PR fixes mainly 2 issues that are : 
- https://github.com/TimDettmers/bitsandbytes/issues/42 
- https://github.com/huggingface/transformers/issues/19445

This issue seems to be unrelated to `bitsandbytes` 8-bit models but slightly more tricky than that. When instantiating a model using `device_map=auto` (EDIT: regardless the `dtype`), the model returns the logits on the same device as the input. I think that this is expected since this is how `accelerate` builds its hooks for each module. So I don"t expect the fix to be done on `accelerate` but more on `transformers` side.

Therefore, if a user calls `generate` or just a simple forward function with a half-precision model that has been instantiated with `device_map=auto` and if the input is initially on the CPU, they may encounter unexpected behaviours such as `top-k-cpu not implemented for Half` since the sampling operations (`top_k`, `top_p`, etc) are done on the same device as the logits - that are in this specific case on CPU. 

This PR addresses this fix by forcing the logits to be on the same device as the **first** module of the model (in case the model is sharded across multiple devices, it makes sense to have the `input_ids` to be on the `device` of the first module). The PR also adds a warning message, suggesting the user to explicitly put the `input_ids` on the same device type as the model.

The PR also adds a slow `bnb` test to make sure this situation will not happen in the future!

# How to reproduce the issue? 

With this simple snippet you can reproduce the issue:
```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MAX_NEW_TOKENS = 128
model_name = 'gpt2'

text = """
Q: On average Joe throws 25 punches per minute. A fight lasts 5 rounds of 3 minutes. 
How many punches did he throw?\n
A: Let’s think step by step.\n"""
tokenizer = AutoTokenizer.from_pretrained(model_name)
input_ids = tokenizer(text, return_tensors="pt").input_ids

model = AutoModelForCausalLM.from_pretrained(
  model_name,
  device_map='auto',
  torch_dtype=torch.float16
)
generated_ids = model.generate(input_ids, max_length=len(input_ids[0])+25, do_sample=True, top_p=0.7)
print(tokenizer.decode(generated_ids[0]))
```

Thanks!
cc @sgugger @ydshieh @Narsil 

Comments:
Thanks for the feedback @Narsil ! 

Regarding `accelerate` I am unsure about the magic that happens there and if it is fixable without breaking anything. If there is something I can fix from `accelerate` happy to open a PR there! Gently pinging @sgugger and @muellerzr to see how we can fix that from `accelerate`

I think that `next(self.parameters())` should do the trick too, regarding the ordering in pytorch, I have observed a similar phenomenon in https://github.com/huggingface/transformers/pull/18312 / when printing a module the output is sensitive to the order of each submodule. I think that `module.parameters()` uses the same logic as `print` -> it uses the `self._modules` attribute from [`nn.Module`](https://pytorch.org/docs/stable/_modules/torch/nn/modules/module.html#Module).
So I believe that if a model has been defined correctly (ie, `Embedding` layer at the beginning of the module) `next(self.parameters())` should return the parameter of the `Embedding` layer. This is an example script I used to confirm my intuition but small details that I might be missing can change my hypothesis!
```
>>> model = nn.Sequential(
... nn.Embedding(1, 2),
... nn.Linear(2, 3),
... nn.Linear(3, 4),
... )
>>> list(model.parameters())[0]
Parameter containing:
tensor([[0.1374, 1.0764]], requires_grad=True)
```
But maybe there is a better way to do so! I'll check what is done on `accelerate`.

Also in the worst case we can just leave the warning message only, and force the user to pass an `input_ids` that is on the same device as the model!

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

This should be neither in Accelerate nor here in my opinion. There is a legitimate use case where you want to:
- make the forward passes on your GPU
- have the generation stuff happen on CPU
because you lack GPU memory (for instance). This is why Accelerate does not force the inputs to be on the same device as the model.

We can add a warning on the Transformer side, but we shouldn't force the inputs to change device. Just tell the user in case they did a mistake and that generation might be slow.

Perfect thanks! 
Fine with this change! I have kept the warning message and reverted the force assignment in aaf6ecb4d2f8872204f7ad80b38f21ac04b7b267 Let me know what do you think ;) !

We should also question ourseleves whether we put the warning inside `sample` or more generally in `generate` before calling `sample` 

Perfect thanks! Adapted the changes in 2a89fcd50e18d0da3bec56e878378f513bbed0c1 Will merge once it's green ;) 
Thanks a lot @Narsil & @sgugger !

I am just reading this discussion for learning purpose, as you all know better than me on this stuff. But @younesbelkada, yoc mentioned

> When instantiating a model using device_map=auto and for half-precision models, the model returns the logits on the same device as the input.

I feel a bit confused: for models running in float32, **does the model return the logits on the same device as the input too**?
I know for float32, we don't have issue on CPU. But I am just wondering the relationship between `dtype` and `the same device`.

Thanks for bringing up my comment @ydshieh !
I think my comment is slightly misleading here, it indeed returns the logits on the same device as the `input_ids` even if the model is in float32. The thing is, some pytorch operations such as `topk` are supported under CPU for float32 but are not supported for float16. So there is no relationship between `dtype` and `the same device` - if you load a model using `device_map=auto` the forward pass of the model will return the output on the same device as the input! 
Will update my comment for clarification! Let me know if anything else is unclear 

No problem @younesbelkada. I was probably focusing on the words too much - bad (sometimes) habits :-)

Ahah don't worry at all! 💪 Agree that my previous statement was super confusing