RuntimeError: CUDA error: device-side assert triggered

I am getting the following error when trying to fine-tune the 7B models from a Llama 2 base:

../aten/src/ATen/native/cuda/Indexing.cu:1146: indexSelectLargeIndex: block: [640,0,0], thread: [127,0,0] Assertion `srcIndex < srcSelectDimSize` failed.                                                   
Traceback (most recent call last):                                                                                                                                                                          
  File "/home/ubuntu/project/project/fastchat/FastChat/fastchat/train/train_mem.py", line 13, in <module>                                                                                                   
    train()                                                                                                                                                                                                 
  File "/home/ubuntu/project/project/fastchat/FastChat/fastchat/train/train.py", line 270, in train                                                                                                         
    trainer.train(resume_from_checkpoint=True)                                                                                                                                                              
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/transformers/trainer.py", line 1539, in train                                                                                      
    return inner_training_loop(                                                                                                                                                                             
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/transformers/trainer.py", line 1809, in _inner_training_loop                                                                       
    tr_loss_step = self.training_step(model, inputs)                                                                                                                                                        
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/transformers/trainer.py", line 2654, in training_step                                                                              
    loss = self.compute_loss(model, inputs)                                                                                                                                                                 
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/transformers/trainer.py", line 2679, in compute_loss                                                                               
    outputs = model(**inputs)                                                                                                                                                                               
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl                                                                              
    return forward_call(*args, **kwargs)                                                                                                                                                                    
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn                                                                                   
    ret_val = func(*args, **kwargs)                                                                                                                                                                         
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 1769, in forward
    loss = self.module(*inputs, **kwargs)
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 806, in forward
    outputs = self.model(
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 646, in forward
    inputs_embeds = self.embed_tokens(input_ids)
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1538, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/torch/nn/modules/sparse.py", line 162, in forward
    return F.embedding(
  File "/home/ubuntu/miniconda3/envs/lmflow/lib/python3.10/site-packages/torch/nn/functional.py", line 2210, in embedding
    return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: device-side assert triggered 
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Environment

• System: Ubuntu 20.04
• GPU: 8xA100 40gb and 4xA100 80gb (tried both)
• PyTorch version: 2.0.1+cu117
• fastchat version: 0.2.20
• transformers version: 4.31.0
• accelerate version: 0.21.0
• flash_attn 2.0.1

Start with fastchat/train/train_mem.py and fastchat/train/train.py at the reported train() call around line 270, then trace the inputs into the Llama embedding path shown in the traceback. Reproduce the fine-tuning run using the listed environment and inspect the input indices and model configuration. Done means the cause is identified and the training run no longer triggers the device-side assertion.