DoLa decoding fails on cuda
Author: kosstbarz
Pull request: no
Labels: name: Generation

DoLa decoding on Mixtral model with multi-GPU setup returns error:
```
Traceback (most recent call last):
  File "~/src/evaluation/test.py", line 8, in <module>
    generate_ids = model.generate(inputs.input_ids, max_length=200, dola_layers='low')
  File "~/src/evaluation/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "~/src/evaluation/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 1921, in generate
    result = self._dola_decoding(
  File "~/src/evaluation/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 2366, in _dola_decoding
    next_token_logits = _dola_select_contrast(
  File "~/src/evaluation/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 4440, in _dola_select_contrast
    kl1 = F.kl_div(log_softmax_mature_layer[None, :, :], avg_dist, reduction="none").mean(-1)
  File "~/src/evaluation/.venv/lib/python3.10/site-packages/torch/nn/functional.py", line 2988, in kl_div
    reduced = torch.kl_div(input, target, reduction_enum, log_target=log_target)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
```

### System Info

- `transformers` version: 4.43.0.dev0
- Platform: Linux-5.15.0-1042-oracle-x86_64-with-glibc2.31
- Python version: 3.10.14
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: 0.29.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: Yes
- Using GPU in script?: Yes
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

This script uses DoLa decoding new feature with multi GPU setup.
Run the file test.py:
```
import transformers
model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
model = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)

prompt = "Hey, are you conscious? Can you talk to me?\n\n"
inputs = tokenizer(prompt, return_tensors="pt")
generate_ids = model.generate(inputs.input_ids, max_length=200, dola_layers='low')
result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
```
with this command
`CUDA_VISIBLE_DEVICES=0,1,2,3 python test.py`

### Expected behavior

No errors.

Comments:
cc @gante 

The issue arises because tensors are being allocated on different devices (CUDA and CPU), causing a `RuntimeError`. This often happens when using multi-GPU setups and certain layers or operations are not properly configured to use the same device.

> **Possible Solution** : Ensure all tensors are on the same device (CUDA) before running the `generate` function. 

#### Please try below script may be it helps

```python
import transformers

# Load the model and tokenizer
model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
model = transformers.AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# Define the prompt and tokenize
prompt = "Hey, are you conscious? Can you talk to me?\n\n"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

# Generate the output
generate_ids = model.generate(inputs.input_ids, max_length=200, dola_layers='low')

# Decode and print the result
result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
print(result)
```

> _By moving the input tensors to CUDA using `.to("cuda")`, we ensure all operations are performed on the same device, thus resolving the error._

Thank you @ashishpatel26 for your response.
I tried `.to("cuda")` and get
`Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cuda:3!`
The same error I get with `.to("cuda:0")` and `.to("cuda:3")`
The same result when I use `model.generate(inputs.input_ids.to("cuda:3")`
Frankly I expected that model with device_map == "auto" would put the input tensor on the correct device. Because base Mixtral model (with no dola_layers) works fine with input tensor on cpu.

Gentle ping @gante 

hehe a `.to` was missing in the DoLa body, opening a PR to fix it

(TL;DR the logits were on the same device as the model inputs, but the auxiliary DoLa tensors were on the same device as the lm head. On multi-gpu settings, that's not the same device.)