idefics2-8b-AWQ failed when doing multiple calls
Author: larme
Pull request: no
Labels: name: Quantization

### System Info

os: ubuntu 22.04

python version: 3.10.12

requirements.txt:
```
pillow==10.3.0
torch==2.2.2
transformers== 4.40.0
```

### Who can help?

_No response_

### Information

- [X] The official example scripts
- [ ] 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

```python
from transformers import AutoProcessor, AutoModelForVision2Seq, AwqConfig
from transformers.image_utils import load_image

DEVICE = "cuda:0"
PROCESSOR_ID = "HuggingFaceM4/idefics2-8b"
MODEL_ID = "HuggingFaceM4/idefics2-8b-AWQ"

prompt = "What do we see in this image?"
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image"},
            {
                "type": "text",
                "text": prompt,
            }
        ],
    },
]

quantization_config = AwqConfig(
    bits=4,
    fuse_max_seq_len=4096,
    modules_to_fuse={
        "attention": ["q_proj", "k_proj", "v_proj", "o_proj"],
        "mlp": ["gate_proj", "up_proj", "down_proj"],
        "layernorm": ["input_layernorm", "post_attention_layernorm", "norm"],
        "use_alibi": False,
        "num_attention_heads": 32,
        "num_key_value_heads": 8,
        "hidden_size": 4096,
    }
)

processor = AutoProcessor.from_pretrained(PROCESSOR_ID)
model = AutoModelForVision2Seq.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map=DEVICE,
)

image = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg")

prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
inputs = processor(text=prompt, images=[image], return_tensors="pt")
inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

# first call, everything is OK
generated_ids = model.generate(**inputs, max_new_tokens=500)
generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)
print("first call:", generated_texts)

# second call, failed
generated_ids = model.generate(**inputs, max_new_tokens=500)
generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)
print("second call:", generated_texts)
```

### Expected behavior

the script should output:

```
first call: ['User: What do we see in this image? \nAssistant: Liberty.']
second call: ['User: What do we see in this image? \nAssistant: Liberty.']
```

However, the second call will fail (error message [here](https://gist.github.com/larme/81a3d1ffad497fc0a86824ceba9ec072)). This script will work if we are not using the AWQ quantization.

I did some investigation on my own and I believe `MistralModel.get_input_embeddings` is the cause of the problem. Without quantization, [this line](https://github.com/huggingface/transformers/blob/main/src/transformers/models/idefics2/modeling_idefics2.py#L1608) will always return a tensor of shape in `[1, 1, 4096]`. However, the quantized version will return a tensor of shape like `[1, 343, 4096]`. The second dimension will be `query_length` in [this part](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_attn_mask_utils.py#L283-L285), and makes `ignore_causal_mask` equal to `False`, which eventually will result in a tensor size not matching error.

For reference, `MistralModel.get_input_embeddings` after quantization is in the form of:

```python
<bound method MistralModel.get_input_embeddings of MistralModel(
  (embed_tokens): Embedding(32003, 4096, padding_idx=0)
  (layers): ModuleList(
    (0-31): 32 x MistralDecoderLayer(
      (self_attn): QuantAttentionFused(
        (qkv_proj): WQLinear_GEMM(in_features=4096, out_features=6144, bias=False, w_bit=4, group_size=128)
        (o_proj): WQLinear_GEMM(in_features=4096, out_features=4096, bias=False, w_bit=4, group_size=128)
        (rope): RoPE()
      )
      (mlp): QuantFusedMLP(
        (down_proj): WQLinear_GEMM(in_features=14336, out_features=4096, bias=False, w_bit=4, group_size=128)
        (activation): SiLU()
      )
      (input_layernorm): FasterTransformerRMSNorm()
      (post_attention_layernorm): FasterTransformerRMSNorm()
    )
  )
  (norm): FasterTransformerRMSNorm()
)>
```

It is in the following form before the quantization:

```python
<bound method MistralModel.get_input_embeddings of MistralModel(
  (embed_tokens): Embedding(32003, 4096, padding_idx=0)
  (layers): ModuleList(
    (0-31): 32 x MistralDecoderLayer(
      (self_attn): MistralSdpaAttention(
        (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
        (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        (rotary_emb): MistralRotaryEmbedding()
      )
      (mlp): MistralMLP(
        (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (up_proj): Linear(in_features=4096, out_features=14336, bias=False)
        (down_proj): Linear(in_features=14336, out_features=4096, bias=False)
        (act_fn): SiLU()
      )
      (input_layernorm): MistralRMSNorm()
      (post_attention_layernorm): MistralRMSNorm()
    )
  )
  (norm): MistralRMSNorm()
)>
```

@SunMarc @younesbelkada @amyeroberts 

Comments:
Hi @larme 
Thanks for the issue !
This is expected and be fixed in https://github.com/casper-hansen/AutoAWQ/pull/444 - cc @casper-hansen
Meanwhile, you can use this method:
```python
def reset_awq_cache(model):
    """
    Simple method to reset the AWQ fused modules cache
    """
    from awq.modules.fused.attn import QuantAttentionFused

    for name, module in model.named_modules():
        if isinstance(module, QuantAttentionFused):
            module.start_pos = 0
```
and call it after each `generate()`

@younesbelkada Thanks! I will close this issue for now.