GPTNeoXAttention takes extra GPU memory footprint in torch.float16 precision mode.
Author: kikutakou
Pull request: no

### System Info

- `transformers` version: 4.30.1
- Platform: Linux-5.15.0-1034-gcp-x86_64-with-glibc2.2.5
- Python version: 3.8.10
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu117 (True)
- Tensorflow version (GPU?): 2.11.1 (True)
- Flax version (CPU?/GPU?/TPU?): 0.5.3 (cpu)
- Jax version: 0.3.6
- JaxLib version: 0.3.5
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No


@ArthurZucker and @younesbelkada

Hi.
I'm using a model, `GPTNeoXForCausalLM` (defined in `src/transformers/models/gpt_neox/modeling_gpt_neox.py`) with torch.float16 precision by calling `.from_pretrained(torch_dtype=torch.float16)`.
In this mode, the model is expected to calculate in float16 precision to save GPU memory usage.
However, [some of variables](https://github.com/huggingface/transformers/blob/b89fcccd44508110fd11579a554c1876bc10c0ad/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L265) in this model remain float32 and don't turn to float16, and they affects the subsequent calculation.
Eventually, the weight attention, which can be dominant memory consumer, is calculated in float32. 
GPU memory won't be saved as we expected.



The following is the problem detail:

1. setup model [`GPTNeoXForCausalLM`](https://github.com/huggingface/transformers/blob/b89fcccd44508110fd11579a554c1876bc10c0ad/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L598) with `torch_dtype=torch.float16`
2. [`self.cos_cached`](https://github.com/huggingface/transformers/blob/b89fcccd44508110fd11579a554c1876bc10c0ad/src/transformers/models/gpt_neox/modeling_gpt_neox.py#L264) and [`self.sin_cached`](https://github.com/huggingface/transformers/blob/b89fcccd44508110fd11579a554c1876bc10c0ad/src/transformers/models/gpt_neox/modeling_gpt_neox.py#LL265C14-L265C14) in `RotaryEmbedding` class held by `GPTNeoXAttention` are calcurated as float32 in __init__().
3. `GPTNeoXAttention.forward()` calls `RotaryEmbedding.forward()`.
4. `RotaryEmbedding.forward()` prepare the return values in float32.
5. `GPTNeoXAttention.forward()` receives the return values in float32.
6. Hereafter, all variables including `attn_weights` are calculated in float32. 
7. `attn_weights = attn_weights.to(value.dtype)` is called and `attn_weights` is returned to float16.

Because of step 7, the model forward() returns the float16, but it consumes float32 GPU footprint internally.




 









### 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

1. Checkout to [ko_gptneox_fp16_debug branch in https://github.com/kikutakou/transformers](https://github.com/kikutakou/transformers/tree/ko_gptneox_fp16_debug) (this branch only has additional debug print code on origin/main)
2. setup model by GPTNeoXForCausalLM.from_pretrained with torch_dtype=torch.float16
3. model.forward()


Here is a sample code.

```
import torch
from transformers import GPTNeoXForCausalLM, GPTNeoXTokenizerFast

torch.manual_seed(0)

MODEL_NAME = 'cyberagent/gpt-neox-1b-japanese'


# load text
input_text = 'this is test'

# tokenize text
tokenizer = GPTNeoXTokenizerFast.from_pretrained(MODEL_NAME, use_auth_token=True)
t = tokenizer(input_text, return_tensors='pt', truncation=True, padding='longest', add_special_tokens=False)
input_ids = t['input_ids'].cuda()
attention_mask = t['attention_mask'].cuda()
input_len = len(input_ids[0])

model = GPTNeoXForCausalLM.from_pretrained(MODEL_NAME, low_cpu_mem_usage=True,
                                           use_auth_token=True, torch_dtype=torch.float16)

model.eval()
model.cuda()

# generate
generation_len = (input_len + 50)
batch_params = dict(input_ids=input_ids,
                    attention_mask=attention_mask,
                    repetition_penalty=None, num_return_sequences=3, num_beams=1, do_sample=True,
                    temperature=None, top_p=0.95, pad_token_id=1, max_length=generation_len)
output_ids = model.generate(**batch_params).cpu()[0]


# decode
output_ids = output_ids[input_len:]
decoded = tokenizer.decode(output_ids, skip_special_tokens=False)
print(decoded)
```

### Expected behavior

It prints all dtypes if you execute on ko_gptneox_fp16_debug branch.
All dtypes are expected to be float16, but actually float 32.


Comments:
PR is prepared as #24262

cc @younesbelkada @ArthurZucker 

@younesbelkada @ArthurZucker
Hi. This is just a friendly reminder. 

Hi @kikutakou 
For fp16 models it is important to calculate the attention scores in full precision, mainly for numerical stability reasons. Check out for instance: https://github.com/huggingface/transformers/issues/17433 or the thread in (that includes authors from OPT models) https://github.com/huggingface/transformers/pull/17437 to start with. So the computation inside attention module to calculate `attn_weights` should always stay in full precision.
Regarding the positional embeddings, looking at the official implementation, it seems that indeed the positional embeddings are returned in half-precision: https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/positional_embeddings.py#L48 . Maybe @StellaAthena can help us confirm if the rotary embeddings should return fp16 values in half-precision modes

For rope, there was an attempt to fix this here: #23837, as it seems that in the original code they are re-computed each forward, with the correct dtype. It's very detailed! 

> Hi @kikutakou
> For fp16 models it is important to calculate the attention scores in full precision, mainly for numerical stability reasons. Check out for instance: #17433 or the thread in (that includes authors from OPT models) #17437 to start with. So the computation inside attention module to calculate `attn_weights` should always stay in full precision.
> Regarding the positional embeddings, looking at the official implementation, it seems that indeed the positional embeddings are returned in half-precision: https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/positional_embeddings.py#L48 . Maybe @StellaAthena can help us confirm if the rotary embeddings should return fp16 values in half-precision modes

I have no reason to think that you can’t compute rotary embed signs in half-precision.

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.