[`Llava`] Fix llava index errors
Author: younesbelkada
Pull request: yes

# What does this PR do?

Fixes errors on the Hub such as https://huggingface.co/llava-hf/llava-1.5-7b-hf/discussions/6 and https://huggingface.co/llava-hf/bakLlava-v1-hf/discussions/4

I did not managed to repro as the issue seems to happen on some specific custom images for some reason, however @gullalc managed to find a fix https://huggingface.co/llava-hf/llava-1.5-7b-hf/discussions/6#657a2aa96cd623f45c3c499f which do not affect generation as I can confirm by the slow tests. 

The fix is simply to mask out the indices that are out of range of the `extended_attention_mask` - added also the same fix on VipLlava architecture

cc @amyeroberts 

Fixes https://github.com/huggingface/transformers/issues/28197, Fixes https://github.com/huggingface/transformers/pull/27901

Comments:
Hi @amyeroberts , I agree it is quite hacky, let me take some time to further investigate and provide a proper fix

@younesbelkada I saw no other issue except here onwards `batch_index, non_attended_tokens = torch.where(first_layer_past_key_value == 0)`

`first_layer_past_key_value.shape[1]` is bigger than `extended_attention_mask.shape[1]`, so it should be expected that an index could be present in `non_attended_tokens` that is larger than `extended_attention_mask.shape[1]`.

That is why filtering out made sense to me at-least. To avoid this hack, I see that either `first_layer_past_key_value` or `extended_attention_mask` is the root of the problem.



hey - facing a similar issue: it seems to appear when both the inputs and generated outputs are long enough, hence different behaviour for different images. One way to replicate is:

```
import requests
from PIL import Image

import torch
from transformers import AutoProcessor, LlavaForConditionalGeneration

model_id = "llava-hf/llava-1.5-7b-hf"

k = 200
user_prompt = "Describe the image:?\n" * k
prompt = f"USER: <image>\n{user_prompt}ASSISTANT:"
image_file = "http://images.cocodataset.org/val2017/000000039769.jpg"

model = LlavaForConditionalGeneration.from_pretrained(
    model_id, 
    torch_dtype=torch.float16, 
).to(0)

processor = AutoProcessor.from_pretrained(model_id)

raw_image = Image.open(requests.get(image_file, stream=True).raw)
inputs = processor(prompt, raw_image, return_tensors='pt').to(0, torch.float16)

print(k, inputs['input_ids'].size())
output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
print(k, output.size())
print(processor.decode(output[0][-100:], skip_special_tokens=True))
```

Running on A10G on current main

```
!CUDA_LAUNCH_BLOCKING=1 python test_llava.py

2023-12-18 14:20:27.237501: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-12-18 14:20:27.237560: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-12-18 14:20:27.237608: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-12-18 14:20:27.244438: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
Loading checkpoint shards: 100%|██████████████████| 3/3 [00:01<00:00,  1.59it/s]
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
200 torch.Size([1, 1412])
/databricks/python/lib/python3.10/site-packages/torch/nn/modules/conv.py:459: UserWarning: Applied workaround for CuDNN issue, install nvrtc.so (Triggered internally at ../aten/src/ATen/native/cudnn/Conv_v8.cpp:80.)
  return F.conv2d(input, weight, bias, self.stride,


--- added some prints
Extended attention mask: (1, 575); 
Attention mask: (1, 1413); 
first_layer_past_key_value (1, 1987); 
Target seqlen: 1988; 
Batch index: tensor([0], device='cuda:0'); 
Non attended tokens: tensor([1881], device='cuda:0')
---


../aten/src/ATen/native/cuda/IndexKernel.cu:92: operator(): block: [0,0,0], thread: [0,0,0] Assertion `index >= -sizes[i] && index < sizes[i] && "index out of bounds"` failed.
Traceback (most recent call last):
  File "", line 25, in <module>
    output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
  File "/databricks/python/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-e7088835-8718-43c0-b531-9c937824ca9c/lib/python3.10/site-packages/transformers/generation/utils.py", line 1731, in generate
    return self.greedy_search(
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-e7088835-8718-43c0-b531-9c937824ca9c/lib/python3.10/site-packages/transformers/generation/utils.py", line 2592, in greedy_search
    outputs = self(
  File "/databricks/python/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/local_disk0/.ephemeral_nfs/envs/pythonEnv-e7088835-8718-43c0-b531-9c937824ca9c/lib/python3.10/site-packages/transformers/models/llava/modeling_llava.py", line 439, in forward
    extended_attention_mask[batch_index, non_attended_tokens] = 0
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

So if I understand correctly we should be masking the resulting `attention_mask` instead of `extended_attention_mask`?

Also, how do these zeros appear in `first_layer_past_key_value` in the first place?



Thanks for the reproducer, I'll try to run some experiments on my end

> Also, how do these zeros appear in first_layer_past_key_value in the first place?

Because the extended hidden states are initialized with all zeros , hence on the first layer they should stay un-touched so the first past kv cache should remain all zeros in the places where you have padd tokens

@adilzhan-ismailov-depop Thanks for sharing the example. I would also agree that it has something to do with length of the generated output for a certain image. I am not sure if it has anything to do with input length, as the three prompts I tried, the different between prompt length was not much.

@younesbelkada To answer your questions:
"Do you use one image per prompt? Are the prompts you use long? Can you somehow reproduce it with an image that you can find on the internet?"

Yes, I used one image per prompt. Three different prompts were used in different runs. Smallest prompt was 7 words and the biggest one was 17 words. I can try to find more example images on the internet for which the same error is thrown, if needed.

Hi @gullalc 

>  I can try to find more example images on the internet for which the same error is thrown, if needed.

Yes that would be really great, thanks !

> Because the extended hidden states are initialized with all zeros , hence on the first layer they should stay un-touched so the first past kv cache should remain all zeros in the places where you have padd tokens

Thanks - but I think in the example with batch size of one we shouldn't have any pad tokens? We can reproduce the error by adding a padding token to any input manually though:

```
import requests
from PIL import Image

import torch
from transformers import AutoProcessor, LlavaForConditionalGeneration

model_id = "llava-hf/llava-1.5-7b-hf"

prompt = f"USER: <image>\nDescribe the image:?\nASSISTANT:"
image_file = "http://images.cocodataset.org/val2017/000000039769.jpg"

model = LlavaForConditionalGeneration.from_pretrained(
    model_id, 
    torch_dtype=torch.float16, 
    device_map='auto'
)

processor = AutoProcessor.from_pretrained(model_id)

device = 'cuda'

raw_image = Image.open(requests.get(image_file, stream=True).raw)
inputs = processor(prompt, raw_image, return_tensors='pt').to(device)

# add padding token manually
pad_token_id = processor.tokenizer.pad_token_id #  32001
inputs['input_ids'] = torch.hstack([torch.ones((1, 1), dtype=torch.int64, device=device) * pad_token_id, inputs['input_ids']])
inputs['attention_mask'] = torch.hstack([torch.zeros((1, 1), dtype=torch.int64, device=device), inputs['input_ids']])

output = model.generate(**inputs, max_new_tokens=200, do_sample=False)
```

This fails because since we have an image, padding token is not at the first position, so we fail the first time we create extended_attention_mask and try to index it

So why does this happen without padding tokens, and likelihood is higher with longer inputs? This is likely to do with half-precision. If you run this experiment you can see that this happens for float16 much more frequently than for float32:

```
import torch
import altair as alt
import pandas as pd
from tqdm.auto import tqdm

# Function to run the experiment
def run_experiment(dtype, num_runs=1000):
    lengths = []
    for _ in tqdm(range(num_runs), desc=f"Running with dtype={dtype}"):
        random_tensor = torch.rand((2000, 1), dtype=dtype) # num. of elements in the original example
        batch_index, non_attended_tokens = torch.where(random_tensor == 0)
        lengths.append(len(batch_index))
    return lengths

# Running the experiments
lengths_float16 = run_experiment(torch.float16)
lengths_float32 = run_experiment(torch.float32)

# Creating a DataFrame for visualization
df = pd.DataFrame({
    "Length": lengths_float16 + lengths_float32,
    "Dtype": ["float16"] * len(lengths_float16) + ["float32"] * len(lengths_float32)
})

# Plotting the results
chart = alt.Chart(df).mark_bar().encode(
    x=alt.X('Length:Q', title="Num. of zero entries in 2000-element array"),
    y='count(Length):Q',
    column='Dtype:N'
).properties(
    width=220,
    height=200
)

chart
```
![image](https://github.com/huggingface/transformers/assets/13088690/1aa12bc1-54f2-4b5b-82f5-708115c8bde0)

In practical terms I think it's ok, but maybe there is a more elegant way to identify non-attended tokens. The logic that handles the attention mask is still an issue though in case we have real padding tokens in the batch

Hi :)
Thanks for investigating this! Just to let you know I'm facing the same issue when using images from the German split of the XM3600 dataset with a batch size > 1.

Here is some log extract:
```python
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 134, in run
    self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
  File "/homeXXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/lightning/pytorch/loops/evaluation_loop.py", line 391, in _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_args)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/lightning/pytorch/trainer/call.py", line 309, in _call_strategy_hook
    output = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/lightning/pytorch/strategies/strategy.py", line 416, in test_step
    return self.lightning_module.test_step(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/gitrepos/lmmm/lmmm/model/mixins/evaluation.py", line 201, in test_step
    return self._in_text_image_out_text(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/gitrepos/lmmm/lmmm/model/mixins/evaluation.py", line 101, in _in_text_image_out_text
    _, pred_text = self.generate(
                   ^^^^^^^^^^^^^^
  File "/home/XXX/gitrepos/lmmm/lmmm/model/lit_llava.py", line 73, in generate
    generated_ids: torch.Tensor = self.model.generate(
                                  ^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/transformers/generation/utils.py", line 1718, in generate
    return self.greedy_search(
           ^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/transformers/generation/utils.py", line 2579, in greedy_search
    outputs = self(
              ^^^^^
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/XXX/miniforge3/envs/lmmm/lib/python3.11/site-packages/transformers/models/llava/modeling_llava.py", line 428, in forward
    extended_attention_mask[batch_index, non_attended_tokens] = 0
    ~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: CUDA error: device-side assert triggered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
```

The fix introduced in this PR fixes the issue though

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_28032). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Note I will try to fix the SDPA regression (for users that perform multi-image & multi-prompt such as https://github.com/huggingface/transformers/issues/28184) in a separate PR , meanwhile users can always use the model with `attn_implementation="eager"` to revert the previous behaviour. 

Thanks a lot for the review! Tests are passing on my VM which is a 2xT4 with the same pytorch & bnb version as the docker image we use! Merging ! 🚀 
Thanks to all contributors for the insightful discussion and the fix! 

Have the same problem and this is very helpful. Thanks!

For anyone that wants to use this fix before the next release:
```bash
pip install -U git+https://github.com/huggingface/transformers.git
```

I tried this pr but it still gives me an error.
I think the core issue is 
```python
first_layer_past_key_value.size(-1) > extended_attention_mask.size(-1)
# May induce index errors
```
, which hasn't been addressed.
But I really don't quite understand what this piece of code is doing, the fix I think should look like this, but I don't know if that's correct:
```python
extended_attention_mask = torch.ones(
    (attention_mask.shape[0], target_seqlen - attention_mask.shape[1]),
    dtype=attention_mask.dtype,
    device=attention_mask.device,
)

# Zero-out the places where we don't need to attend
attention_mask = torch.cat((attention_mask, extended_attention_mask), dim=1)
attention_mask[batch_index, non_attended_tokens] = 0
# attention_mask.size() = first_layer_past_key_value.size() + 1
```
or

```python
extended_attention_mask = torch.ones(
    (attention_mask.shape[0], target_seqlen - attention_mask.shape[1]),
    dtype=attention_mask.dtype,
    device=attention_mask.device,
)

valid_indices = non_attended_tokens >= attention_mask.size(-1)
new_batch_index = batch_index[valid_indices]
new_non_attended_tokens = non_attended_tokens[valid_indices]

# Zero-out the places where we don't need to attend
extended_attention_mask[new_batch_index, new_non_attended_tokens - attention_mask.size(-1)] = 0
```
might be better

Thanks @NicholasCao !
Can you file a separate issue for this and tag me? If you can also provide a reproducer it would be great. This PR that got merged fixes the issue explained here: https://github.com/huggingface/transformers/pull/28032#issuecomment-1860650043 which hopefully should cover most of the issues related with llava and index errors

@NicholasCao I managed to repro your issue that seems to happen in the case one passes a custom past key value, which is the case for AWQ. It should be fixed in https://github.com/huggingface/transformers/pull/28239

I'm not using awq, I'm having this problem when I'm batch inference images, it's harder to reproduce to find the specific image

@NicholasCao #28239 should solve it, let me know if the PR fixes your issue

thx, it works

Thanks @NicholasCao !