Fix Unnecessary move of tensors from CPU to GPU in LlamaRotaryEmbedding
Author: ma787639046
Pull request: yes

# What does this PR do?
The original implementation of LlamaRotaryEmbedding does not use `cos_cached` & `sin_cached` tensors as the PyTorch Parameter or Buffer, thus these tensors do not move to GPU when we use `model.to(gpu_id)` or `model.cuda()`. They will keep in the device CPU.

This PR adjusts the `cos_cached` & `sin_cached` tensors to the Buffer with persistent=False. This keeps these tensors moving from CPU to GPU together with the model, while keeping them out of the model's state_dict as original.

# Fixes:
Fix unnecessary moves of tensors from CPU to GPU in LlamaRotaryEmbedding, for saving a large amount of CPU usage especially when we do inference.

Code for Reproducing the issue:
```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"    # Single card Generation

from tqdm import tqdm

import torch
from transformers.models.llama.modeling_llama import LlamaForCausalLM
from transformers.models.llama.tokenization_llama import LlamaTokenizer

tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf")
model = LlamaForCausalLM.from_pretrained("decapoda-research/llama-7b-hf", torch_dtype=torch.float16)
model = model.cuda()
model.eval()

# Batch generation
inputs = [
    "LLaMa is a large language model developed by Meta AI, for",
] * 32

batch = tokenizer(inputs, return_tensors="pt", add_special_tokens=False)
batch = batch.to(model.device)

# Here we do some high computational batched generation
for i in tqdm(range(5000)): 
    generated = model.generate(batch["input_ids"],
                                temperature=0.7, top_p=0.9, do_sample=True,
                                num_beams=1, max_new_tokens=600,)

```
Use `top` command in bash to watch the CPU usage.

Here are the comparison before applying this PR and after this PR:

Before:
| Fix    | USER | PR | NI | VIRT   | RES  | SHR    | S | %CPU | %MEM | TIME+   | COMMAND  |
|--------|------|----|----|--------|------|--------|---|------|------|---------|----------|
| Before | root | 20 | 0  | 108.6g | 1.9g | 411620 | R | 6263 | 0.2  | 40:28.1 | python   |

After:
| Fix   | USER | PR | NI | VIRT   | RES  | SHR    | S | %CPU | %MEM | TIME+   | COMMAND  |
|-------|------|----|----|--------|------|--------|---|------|------|---------|----------|
| After | root | 20 | 0  | 108.6g | 1.8g | 414360 | R | 98.3 | 0.2  | 03:21.6 | python   |

Here the CPU usage drops to a normal level because the `cos_cached` & `sin_cached` tensors can move to GPU correctly with the model. This helps avoid unnecessary moves of tensors from CPU to GPU in LlamaRotaryEmbedding.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?

## Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

 @ArthurZucker and @younesbelkada


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

Did you accidentally break meta loading?
`with init_empty_weights():`
leaves `cos_cached` and `sin_cached` on meta device and they won't be initialized because they are not persistent. 

Since `inv_freq` is a persistent buffer, it should be ok to also make harmonics persistent

Hi @BlackSamorez, would you like to open a PR with these suggested changes including details about the issue they resolve? 

@BlackSamorez `init_empty_weights` ignores buffers by default, so this should not cause any problem. We have multiple instance of non-persistent buffers in the lib and this is not a problem. I've also run Llama without any issue after it being initialized on the meta device.

Hi, I test the following two codes on my device. It seems the meta device works correctly in this PR.

```python
import pickle
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import torch
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaForCausalLM

model_name_or_path = "decapoda-research/llama-7b-hf"

model1 = LlamaForCausalLM.from_pretrained(model_name_or_path, torch_dtype=torch.float16)
model1 = model1.to(torch.device("cuda:0"))

# Save a initialized cos_cached tensor to `cos1.pt`, for comparasion with meta device loading
cos1 = model1.model.layers[0].self_attn.rotary_emb.cos_cached.to(torch.device("cpu"))
pickle.dump(cos1, open("cos1.pt", 'wb'))
```

```python
import pickle
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import torch
from transformers.models.llama.configuration_llama import LlamaConfig
from transformers.models.llama.modeling_llama import LlamaForCausalLM

from accelerate import init_empty_weights, infer_auto_device_map, load_checkpoint_and_dispatch

model_name_or_path = "decapoda-research/llama-7b-hf"

config = LlamaConfig.from_pretrained(model_name_or_path, torch_dtype=torch.float16)
with init_empty_weights():
    model0 = LlamaForCausalLM(config)

model0 = load_checkpoint_and_dispatch(
    model0, model_name_or_path, device_map='auto', 
)
# Compare the `cos_cached` tensor
cos0 = model0.model.layers[0].self_attn.rotary_emb.cos_cached.to(torch.device("cpu"))
cos1 = pickle.load(open("cos1.pt", 'rb'))
all((cos0==cos1).tolist())      # True
```

@BlackSamorez Maybe you can check the results on your device.

Yes, you're right and I was wrong. It works and the problem was in entirely different part of my program.
Consider https://github.com/huggingface/transformers/pull/22234#discussion_r1143320748 and https://github.com/huggingface/transformers/pull/22234#discussion_r1143321732 invalid.
Thank you!

I'm still facing this issue with latest deepspeed (0.9.5+1491e14e) and transformers (4.31.0.dev0). I feel this issue is more likely related to the LLaMA implementation here (LlamaRotaryEmbedding).

```
RuntimeErrorcos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]: 
indices should be either on cpu or on the same device as the indexed tensor (cpu)
    RuntimeErrorcos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]: 
indices should be either on cpu or on the same device as the indexed tensor (cpu)    
RuntimeErrorcos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]: 
indices should be either on cpu or on the same device as the indexed tensor (cpu)
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)    
cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
RuntimeError    : cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]indices should be either on cpu or on the same device as the indexed tensor (cpu)

RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
```

> I'm still facing this issue with latest deepspeed (0.9.5+1491e14e) and transformers (4.31.0.dev0). I feel this issue is more likely related to the LLaMA implementation here (LlamaRotaryEmbedding).
> 
> ```
> RuntimeErrorcos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]: 
> indices should be either on cpu or on the same device as the indexed tensor (cpu)
>     RuntimeErrorcos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]: 
> indices should be either on cpu or on the same device as the indexed tensor (cpu)    
> RuntimeErrorcos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]: 
> indices should be either on cpu or on the same device as the indexed tensor (cpu)
> RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)    
> cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
> RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
>     cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
> RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
>     cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
> RuntimeError    : cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]indices should be either on cpu or on the same device as the indexed tensor (cpu)
> 
> RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
> ```

I encountered exactly the same issue，training failed when using zero3

Thanks for the input, will investigate! 

any updates? I also meet this issues. with ds==0.9.3, transformers==4.32.0dev

Did not have time to investigate,  we are going to need a reproducer if you want some help here. Pinging @pacman100 when we have a reproducer shared!