There is bug in trainer of Llama:indices should be either on cpu or on the same device as the indexed tensor (cpu)
### System Info

Setting ds_accelerator to cuda (auto detect)

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.30.0.dev0
- Platform: Linux-5.4.0-132-generic-x86_64-with-glibc2.27
- Python version: 3.10.11
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 1.13.1+cu116 (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?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@ArthurZucker  @sgugger @pacman100

### Information

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

### Tasks

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

### Reproduction

```
from dataclasses import dataclass, field
from typing import Optional, Tuple
import transformers
from transformers import Trainer
from transformers.models.llama import LlamaForCausalLM  # noqa
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    AutoModel
)
from nl2sql_dataset import Nl2SqlJsonlDataset
import env
import time


@dataclass
class ModelArguments:
    pretrain_path: str = field(
        default=f"{env.MODEL_ROOT.joinpath('llama-7b-hf')}" # llama-13b-hf
    )


@dataclass
class DataArguments:
    train_file: str = field(
        default=f"{env.INPUT_ROOT.joinpath('trainset/config.json')}",
        metadata={"help": "A josnl file containing the training corpus"},
    )
    validation_file: str = field(
        default=f"{env.INPUT_ROOT.joinpath('devset/config.json')}",
        metadata={"help": "A jsonl file containing the validation corpus"},
    )
    max_seq_length: int = field(
        default=512, metadata={"help": "Max sequence length for training"}
    )
    pad_to_max_length: bool = field(default=False)


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    cache_dir: Optional[str] = field(default=None)
    optim: str = field(default="adamw_torch")
    model_max_length: int = field(
        default=512,
        metadata={
            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    num_train_epochs: int = 5
    evaluation_strategy: str = field(default="epoch")
    save_strategy: str = "epoch"
    fp16: bool = True
    save_total_limit: int = 5
    load_best_model_at_end: bool = False
    warmup_steps: int = 0
    logging_steps: int = 1
    gradient_checkpointing: bool = True
    ddp_timeout: int = 3600
    output_dir: str = field(
        default=f"{env.OUTPUT_ROOT.joinpath(time.strftime('%Y年%m月%d日%H时%M分%S秒'))}",
    )


def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
    """Collects the state dict and dump to disk."""
    state_dict = trainer.model.state_dict()
    if trainer.args.should_save:
        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
        del state_dict
        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa


def parse_args() -> Tuple[ModelArguments, DataArguments, TrainingArguments]:
    parser = transformers.HfArgumentParser(
        (ModelArguments, DataArguments, TrainingArguments)
    )
    return parser.parse_args_into_dataclasses()


def train():
    model_args, data_args, training_args = parse_args()
    if "chatglm" in model_args.pretrain_path:
        print(model_args.pretrain_path)
        model = AutoModel.from_pretrained(model_args.pretrain_path, trust_remote_code=True, empty_init=False)
    else:
        model = AutoModelForCausalLM.from_pretrained(model_args.pretrain_path)
    print(model_args, data_args, training_args)
    dataset = Nl2SqlJsonlDataset(
        pretrain_path=model_args.pretrain_path,
        train_file_path=data_args.train_file,
        validation_file_path=data_args.validation_file,
        max_seg_length=data_args.max_seq_length,
        pad_to_max_length=data_args.pad_to_max_length,
    )
    dataset.setup()

    # Tell Trainer not to attempt DataParallel
    model.is_parallelizable = True
    model.model_parallel = True

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset.train_dataset,
        eval_dataset=dataset.val_dataset,
        data_collator=dataset.collate_fn,
    )
    model.config.use_cache = False

    trainer.train()
    trainer.save_state()
    safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)


if __name__ == "__main__":
    train()



```

**error:** To simplify the output information, I ran only on one card
Setting ds_accelerator to cuda (auto detect)
[2023-06-06 12:16:17,633] [WARNING] [comm.py:152:init_deepspeed_backend] NCCL backend in DeepSpeed not yet implemented
[2023-06-06 12:16:17,633] [INFO] [comm.py:594:init_distributed] cdb=None
[2023-06-06 12:16:17,633] [INFO] [comm.py:625:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
[2023-06-06 12:17:21,017] [INFO] [partition_parameters.py:454:__exit__] finished initializing model with 6.74B parameters
Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████| 3/3 [00:41<00:00, 13.86s/it]ModelArguments(pretrain_path='/mnt/data/wxc/workspace/pretrained_models/2/llama-7b') DataArguments(train_file='/mnt/data/wxc/workspace/release/data/trainset/config.json', validation_file='/mnt/data/wxc/workspace/release/data/devset/config.json', max_seq_length=512, pad_to_max_length=False) TrainingArguments(

Parameter Offload: Total persistent parameters: 266240 in 65 params
0%| | 0/87850 [00:00<?, ?it/s]Traceback (most recent call last):
File "/mnt/data/wxc/workspace/release/start_ds_finetune.py", line 118, in
train()
File "/mnt/data/wxc/workspace/release/start_ds_finetune.py", line 112, in train
trainer.train()
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/trainer.py", line 1661, in train
return inner_training_loop(
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/trainer.py", line 1946, in _inner_training_loop
tr_loss_step = self.training_step(model, inputs)
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/trainer.py", line 2756, in training_step
loss = self.compute_loss(model, inputs)
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/trainer.py", line 2781, in compute_loss
outputs = model(**inputs)
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
return forward_call(*input, **kwargs)
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
ret_val = func(*args, **kwargs)
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1733, in forward
loss = self.module(*inputs, **kwargs)
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
result = forward_call(*input, **kwargs)
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/models/llama/modeling_llama.py", line 688, in forward
outputs = self.model(
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
result = forward_call(*input, **kwargs)
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/models/llama/modeling_llama.py", line 570, in forward
layer_outputs = torch.utils.checkpoint.checkpoint(
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 249, in checkpoint
return CheckpointFunction.apply(function, preserve, *args)
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/utils/checkpoint.py", line 107, in forward
outputs = run_function(*args)
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/models/llama/modeling_llama.py", line 566, in custom_forward
return module(*inputs, output_attentions, None)
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
result = forward_call(*input, **kwargs)
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/models/llama/modeling_llama.py", line 292, in forward
hidden_states, self_attn_weights, present_key_value = self.self_attn(
File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1212, in _call_impl
result = forward_call(*input, **kwargs)
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/models/llama/modeling_llama.py", line 202, in forward
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
File "/mnt/data/wxc/workspace/Llama-X/src/transformers/src/transformers/models/llama/modeling_llama.py", line 134, 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)


ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 12104) of binary: /home/wxc/miniconda3/envs/llamax/bin/python3.1
Traceback (most recent call last):
  File "/home/wxc/miniconda3/envs/llamax/bin/torchrun", line 8, in <module>
    sys.exit(main())
  File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 346, in wrapper
    return f(*args, **kwargs)
  File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/distributed/run.py", line 762, in main
    run(args)
  File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/distributed/run.py", line 753, in run
    elastic_launch(
  File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 132, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/home/wxc/miniconda3/envs/llamax/lib/python3.10/site-packages/torch/distributed/launcher/api.py", line 246, in launch_agent    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
start_ds_finetune.py FAILED
------------------------------------------------------------
Failures:
  <NO_OTHER_FAILURES>
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2023-06-06_12:20:18
  host      : omnisky
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 12104)
  error_file: <N/A>
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================


By addition,my deepspeed is the latest version,installed by git and compile.


### Expected behavior

This is a finetune with nl2sql dataset.

And my data format is 
{"source": "dusql_中国历史名城", "table": [{"table_name": "城市", "header": ["词条id", "名称", "所属省份", "常住人口", "城区面积", "建城年数"], "rows": []}, {"table_name": "都城", "header": ["朝代", "古称", "城市id", "建都起始时间", "建都结束时间", "建都年数"], "rows": []}], "sql": "select 名称 , 所属省份 from 城市 where 词条id not in ( select 城市id from 都城 )", "question": "哪些城市没有做过都城，给出这些城市名和其省份。"}
Hello, what is the launch command?
What version of accelerate is being used?
Also see this https://github.com/microsoft/DeepSpeed/issues/3678<|||||>hi~ thanks for your help!
CUDA_VISIBLE_DEVICES="1,2,4,5" torchrun  --nnodes=1 --nproc_per_node=4 start_ds_finetune.py --deepspeed deepspeed_config.json --learning_rate=2e-5 --per_device_train_batch_size=4 --gradient_accumulation_steps=1

accelerate == 0.19.0

> Hello, what is the launch command? What version of accelerate is being used? Also see this [microsoft/DeepSpeed#3678](https://github.com/microsoft/DeepSpeed/issues/3678)

<|||||>Can you update the accelerate from the main branch and use the PR of DeepSpeed LinkedIn the above DeepSpeed issue<|||||>Sorry I'm a beginner, do you mean this?  
pip install --upgrade accelerate    
git clone -b olruwase/ds_3678 https://github.com/microsoft/DeepSpeed.git
cd DeepSpeed
DS_BUILD_CPU_ADAM=1 DS_BUILD_AIO=1 DS_BUILD_UTILS=1 pip install -e .<|||||>My problem is solved！！！Thanks！！！How excellent you are!!

> Can you update the accelerate from the main branch and use the PR of DeepSpeed LinkedIn the above DeepSpeed issue

