DeepSpeed ZeRO stage3+huggyllama/llama: RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)
Author: memray
Pull request: no

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

Comments:
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 😅 

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.