Trained Gemma 2 crashes with CUDA error during infrence
Author: effortprogrammer
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.44.0
- Platform: Linux-6.8.0-36-generic-x86_64-with-glibc2.39
- Python version: 3.10.14
- Huggingface_hub version: 0.23.5
- Safetensors version: 0.4.3
- Accelerate version: 0.30.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: distributed
- GPU type: NVIDIA A100 80GB PCIe


### Who can help?

@ArthurZucker 
@Narsil 
@SunMarc


### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

This model was Supervised Fine Tuned from gemma-2-9b(base model) with my own SFT Dataset: 
https://huggingface.co/kaki-paper/gemma-2-9b-test-model-for-debug/tree/main

The model has been trained with default chatML provided in this format: 

user
{question}


assistant
{answer}

This is the inference pipeline for this model: 
```python
torch.multiprocessing.set_start_method("spawn")

model = AutoModelForCausalLM.from_pretrained(
    model_path,
    low_cpu_mem_usage=True,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    # load_in_8bit=True,  
    use_cache=False,
    attn_implementation="eager",
  )
model = torch.compile(model)
model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_path, model_max_len=4096, truncation_side='left')
start = time.time()

terminators = [
tokenizer.eos_token_id,
# tokenizer.convert_tokens_to_ids("
")
]

generator = pipeline(
    task="text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=256,
    repetition_penalty=1.2,
    do_sample=False,
    pad_token_id=tokenizer.eos_token_id, 
    eos_token_id = terminators,
    
)

with torch.no_grad():
    for i in tqdm(range(len(df))):        
        print("start Eval")
        doc1 = df.iloc[i, 0]  
        description = doc1
        text = ""
        text += description
        format_finished = "
user" + text + "
" + "
assistant"
        print("===============================================================")
        print(format_finished)
        print("===============================================================")
        encoded_input = tokenizer(format_finished, truncation=True, return_tensors="pt").to("cuda")
        prompt = tokenizer.decode(encoded_input["input_ids"][0], skip_special_token=False)        
        output = generator(prompt, return_full_text=False)
        print(output)
        generated_answer = output[0]["generated_text"]
        df.loc[i, "generated_text"] = generated_answer
```

and got this error: 
```
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [26,0,0], thread: [1,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [26,0,0], thread: [2,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [26,0,0], thread: [3,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [26,0,0], thread: [4,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [26,0,0], thread: [5,0,0] Assertion `-sizes[i] <= index && index < sizes[i] && "index out of bounds"` failed.
{...}
Traceback (most recent call last):
  File "/home/howard/workspace/new_workspace/gemma2_9b_sft/trained_gemma2_9b_eval.py", line 177, in <module>
    output = generator(prompt, return_full_text=False)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/text_generation.py", line 262, in __call__
    return super().__call__(text_inputs, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1257, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1264, in run_single
    model_outputs = self.forward(model_inputs, **forward_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1164, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/text_generation.py", line 351, in _forward
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 2024, in generate
    result = self._sample(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 2982, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 999, in forward
    outputs = self.model(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 847, in forward
    layer_outputs = decoder_layer(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 599, in forward
    hidden_states = self.post_attention_layernorm(hidden_states)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 121, in forward
    output = self._norm(x.float())
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 118, in _norm
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```


### Expected behavior

I knew there were some similar issues before(#31848), but still got this issue on the latest transformers version. 


Comments:
Hey @effortprogrammer , thanks for reporting the issue! Can you give the traceback when you set the env var `CUDA_LAUNCH_BLOCKING=1` ? We'll get a clearer understanding of what's happening that way 

 Most probably something related to new tokens but not having resized the embedding!

> Hey @effortprogrammer , thanks for reporting the issue! Can you give the traceback when you set the env var `CUDA_LAUNCH_BLOCKING=1` ? We'll get a clearer understanding of what's happening that way

This is the full traceback after setting the env var `CUDA_LAUNCH_BLOCKING=1`.

```markdown
Traceback (most recent call last):
  File "/home/howard/workspace/new_workspace/gemma2_9b_sft/trained_gemma2_9b_eval.py", line 177, in <module>
    output = generator(prompt, return_full_text=False)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/text_generation.py", line 262, in __call__
    return super().__call__(text_inputs, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1257, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1264, in run_single
    model_outputs = self.forward(model_inputs, **forward_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1164, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/text_generation.py", line 351, in _forward
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 2024, in generate
    result = self._sample(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 2982, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 999, in forward
    outputs = self.model(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 847, in forward
    layer_outputs = decoder_layer(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 590, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 293, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/cache_utils.py", line 1532, in update
    return update_fn(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/cache_utils.py", line 1507, in _static_update
    k_out[:, :, cache_position] = key_states
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```



> Most probably something related to new tokens but not having resized the embedding!

I knew well-known issue comes when adding new tokens but not having the resized embedding. But, I didn't add new tokens during training, which makes me hard to find the actual reason of this problem..

Alright, what seems to happen is `k_out` and `key_states` have different shapes. That would cause the CUDA index error you're seeing. Might be a cache issue we can look into. 
```python
File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/cache_utils.py", line 1507, in _static_update
k_out[:, :, cache_position] = key_states
```
However I'm noticing you're passing `terminators` to `eos_token_id` in the pipeline, which is a list, and I don't know how the pipeline handles this - can you pass directly `eos_token_id`?


> Alright, what seems to happen is `k_out` and `key_states` have different shapes. That would cause the CUDA index error you're seeing. Might be a cache issue we can look into.
> 
> ```python
> File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/cache_utils.py", line 1507, in _static_update
> k_out[:, :, cache_position] = key_states
> ```
> 
> However I'm noticing you're passing `terminators` to `eos_token_id` in the pipeline, which is a list, and I don't know how the pipeline handles this - can you pass directly `eos_token_id`?

```markdown
Traceback (most recent call last):
  File "/home/howard/workspace/new_workspace/gemma2_9b_sft/trained_gemma2_9b_eval.py", line 177, in <module>
    output = generator(prompt, return_full_text=False)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/text_generation.py", line 262, in __call__
    return super().__call__(text_inputs, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1257, in __call__
    return self.run_single(inputs, preprocess_params, forward_params, postprocess_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1264, in run_single
    model_outputs = self.forward(model_inputs, **forward_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/base.py", line 1164, in forward
    model_outputs = self._forward(model_inputs, **forward_params)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/pipelines/text_generation.py", line 351, in _forward
    generated_sequence = self.model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 2024, in generate
    result = self._sample(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/generation/utils.py", line 2982, in _sample
    outputs = self(**model_inputs, return_dict=True)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 999, in forward
    outputs = self.model(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 847, in forward
    layer_outputs = decoder_layer(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 590, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/accelerate/hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/models/gemma2/modeling_gemma2.py", line 293, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/cache_utils.py", line 1532, in update
    return update_fn(
  File "/data/howard/new_workspace/.venv/lib/python3.10/site-packages/transformers/cache_utils.py", line 1507, in _static_update
    k_out[:, :, cache_position] = key_states
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

```

This is the traceback after directly passing eos_token_id to the pipeline, which seems like the same issue.

Do you need additional information for debugging this issue?

I solved this problem when updating the pytorch version from 2.3.0+cu121 to 2.4.0+cu124 version. So, this problem may be come from the pytorch related issue, I guess. Since I solved this problem, I am closing this issue. 

That is super weird 😅 thanks for updating us! 

Yes... it was pretty interesting experience. Thanks for the support!

It seems like the problem was not the torch version, so reopening this issue.

@ArthurZucker @molbap  

I guess it was some other issue, which is pretty weird because I found the environment version that doesn't show this kind of issue

Environment that has no issue with trained model:
```markdown
- `transformers` version: 4.44.0.dev0
- Platform: Linux-5.4.0-137-generic-x86_64-with-glibc2.35
- Python version: 3.10.6
- Huggingface_hub version: 0.23.3
- Safetensors version: 0.4.3
- Accelerate version: 0.30.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: distributed
- Using GPU in script?: yes
- GPU type: CUDA GPU

```

Environment that has same error occured:
```markdown
- `transformers` version: 4.43.4
- Platform: Linux-6.8.0-36-generic-x86_64-with-glibc2.39
- Python version: 3.10.14
- Huggingface_hub version: 0.23.5
- Safetensors version: 0.4.3
- Accelerate version: 0.30.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: distributed
- Using GPU in script?: yes
- GPU type: NVIDIA A100 80GB PCIe
```

I still cannot figure out what is the issue.. 

Edit: When I run this inference code on cpu(without cuda), the inference code works, and gives me the model result. 

Edit2: When I run this inference code in the single GPU, model output came successful, which assumes this problem alies in the multi-gpu environment.

This fixed with the change from this PR: https://github.com/huggingface/transformers/pull/32543

Check this PR if you encountered the same issue as mine. Closing this issue now. 

In my case, I solved this problem by setting `attn_implementation=“flash_attention_2”.

You now have `flex_attention` that's available!