DeepSpeed ZeRO stage3+huggyllama/llama: RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
### System Info

- deepspeed: 0.9.5+1491e14e
- `transformers` version: 4.31.0.dev0
- Platform: Linux-5.10.133+-x86_64-with-glibc2.17
- Python version: 3.8.15
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.6.2 (gpu)
- Jax version: 0.3.22
- JaxLib version: 0.3.22
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@pacman100

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Is this issue highly related to this [PR](https://github.com/huggingface/transformers/pull/22234 )?
Used the native deepspeed running example (no accelerate).
```
model_engine, optimizer, _, scheduler = deepspeed.initialize(config=args.deepspeed_config, model=model, model_parameters=model_parameters)
```

Tried stage-3 with 7B/13B/30B ckpt and they all errored out, but they worked well with stage2.



Error message:
```
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    outputs = run_function(*args)
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 566, in custom_forward
    result = forward_call(*args, **kwargs)
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 205, in forward
    return module(*inputs, output_attentions, None)
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 137, in apply_rotary_pos_emb
    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)
    result = forward_call(*args, **kwargs)
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 293, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 205, in forward
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
  File "/export/share/ruimeng/env/anaconda/envs/codegen/lib/python3.8/site-packages/transformers/models/llama/modeling_llama.py", line 137, in apply_rotary_pos_emb
    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)
```

### Expected behavior

normal training
Hello @memray, this issue seems unrelated to DeepSpeed integration. Could you provide a minimal reproducible example?<|||||>No, I also tested `mosaicml/mpt-7b` and I found it worked. I believe it's more related to the implementation of LLaMA because `openlm-research/open_llama_7b` also fails. Can you let me know who can help with this issue? <|||||>@ArthurZucker and @younesbelkada can you check if this is related to the model implementation (say this [PR](https://github.com/huggingface/transformers/pull/22234))?<|||||>Hey! thanks for reporting, yes will check this! Seems like this might be it. Also if so fix should be straight forward<|||||>I am also facing the same issue with [chaoyi-wu/PMC_LLAMA_7B](https://huggingface.co/chaoyi-wu/PMC_LLAMA_7B)<|||||>Okay, I forgot to ask, could you provide a full reproducing script? The ROPE was recently updated too, not sure if this was adressed yet.<|||||>Same problem here.<|||||>We cannot help you if we don't have a minimal reproducing script, just commenting same problem will not help 😅 