There is bug in trainer of Llama:indices should be either on cpu or on the same device as the indexed tensor (cpu)
Author: 2018211801
Pull request: no

### 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": "哪些城市没有做过都城，给出这些城市名和其省份。"}

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


could you make a brief about how to solve this bug? I also meet this. Should i change the version of deepspeed?

@2018211801 

@2018211801 


> 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

What is your solution to this problem? Could you please share it?

update deepspeed to latest version solve it