Llama2 inference in bfloat16
Author: JeevanBhoot
Pull request: no

### System Info

- `transformers` version: 4.36.2
- Platform: Linux-5.15.0-89-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.20.1
- Safetensors version: 0.4.1
- Accelerate version: 0.25.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.2+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 GPU in script?: Yes - 1x RTX 4090
- Using distributed or parallel set-up in script?: No

### Who can help?

@ArthurZucker @younesbelkada @SunMarc 

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

I am trying to run both GPTQ and unquantized Llama2 models:

```python
gptq_config = GPTQConfig(bits=4, disable_exllama=True)

model_path = "TheBloke/Llama-2-7B-Chat-GPTQ"
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=False, revision="main", quantization_config=gptq_config, torch_dtype=torch.bfloat16).to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)

input_ids = tokenizer("Tell me an interesting fact!", return_tensors="pt").input_ids.to("cuda").to(torch.bfloat16)
output = model.generate(input_ids)
```

and 

```python
model_path = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=False, revision="main", torch_dtype=torch.bfloat16).to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)

input_ids = tokenizer("Tell me an interesting fact!", return_tensors="pt").input_ids.to("cuda").to(torch.bfloat16)
output = model.generate(input_ids)
```

### Expected behavior

I am trying to run Llama2 inference in bfloat16 - I want all weights and computations to be in bfloat16. When I run the two snippets provided, I encounter the following error:
```
  File "/home/jeevan/miniconda3/envs/llama2_env/lib/python3.10/site-packages/torch/nn/functional.py", line 2233, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: Long, Int; but got CUDABFloat16Type instead (while checking arguments for embedding)
```

If I keep the input in int64, then this works fine i.e. changing just the model weights to bfloat16 works fine. But I want all computation to be performed in bfloat16.

Comments:
Hi @JeevanBhoot 
This is not supported in PyTorch, when passing input_ids they need to stay in `torch.int64`, the only way to run inference in bf16 is the following: 
```diff
model_path = "meta-llama/Llama-2-7b-chat-hf"
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", trust_remote_code=False, revision="main", torch_dtype=torch.bfloat16).to("cuda")
tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)

- input_ids = tokenizer("Tell me an interesting fact!", return_tensors="pt").input_ids.to("cuda").to(torch.bfloat16)
+ input_ids = tokenizer("Tell me an interesting fact!", return_tensors="pt").input_ids.to("cuda")
output = model.generate(input_ids)
```

@younesbelkada When input_ids in `torch.int64`, the output from the model is also `torch.int64`. Is there no way to have the model activations and outputs be `torch.bfloat16`?

@JeevanBhoot if you pass `torch.LongTensor` to a `nn.Embedding` layer the output should be a floating point. If you load your model in bf16 it should be `int64` -> `bfloat16`, you can inspect that by doing `model.get_input_embeddings()(input_ids).dtype`

@younesbelkada Ok, I can confirm that `model.get_input_embeddings()(input_ids).dtype` returns `bfloat16` but `output.dtype` is `int64`. 

Is the final output of the model cast from `bfloat16` to `int64`, or is the dtype changing at some point in the model?

@JeevanBhoot `output` will always be in `int64` since it corresponds to the predicted tokens. There is no need to cast the output indices in `bfloat16`, what is your motivation behind getting `output` in bf16?

@younesbelkada I don't need the final output to be in bf16, but I want to confirm that EVERY operation and output within the model is bf16.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.