IDEFICS cannot work with past_key_values when I'm using prefix tuning from peft library.
Author: Kamichanw
Pull request: no
Labels: name: bug

name: Vision

name: PEFT

### System Info

- `transformers` version: 4.44.2
- Platform: Linux-4.15.0-76-generic-x86_64-with-glibc2.27
- Python version: 3.12.4
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.3
- Accelerate version: 0.34.2
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - mixed_precision: fp16
        - use_cpu: False
        - debug: False
        - num_processes: 4
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - enable_cpu_affinity: False
        - deepspeed_config: {'gradient_accumulation_steps': 16, 'gradient_clipping': 1.0, 'offload_optimizer_device': 'cpu', 'offload_param_device': 'cpu', 'zero3_init_flag': False, 'zero3_save_16bit_model': False, 'zero_stage': 3}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.3.1+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?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA GeForce RTX 3090

### Who can help?

@ArthurZucker @amyeroberts, @qubvel

### 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

```python
import peft
from transformers import IdeficsForVisionText2Text, AutoProcessor
import sys
import torch

sys.path.insert(0, "..")
import config  # my own file

device = torch.device("cuda:2")
model = IdeficsForVisionText2Text.from_pretrained(
    config.idefics_9b_path, torch_dtype=torch.float16
).to(device)
processor = AutoProcessor.from_pretrained(
    config.idefics_9b_path, torch_dtype=torch.float16
)
model = peft.get_peft_model(
    model,
    peft.PrefixTuningConfig(
        peft_type="PREFIX_TUNING",
        task_type="CAUSAL_LM",
        num_virtual_tokens=1,
        token_dim=4096,
        num_transformer_submodules=1,
        num_attention_heads=32,
        num_layers=32,
        encoder_hidden_size=768,
    ),
)
inputs = processor(["hello"]).to(device)
model(**inputs)
```

### Expected behavior

When I ran above code snippet, I got:
```
{
	"name": "RuntimeError",
	"message": "The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 2",
	"stack": "---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[3], line 2
      1 inputs = inputs.to(device)
----> 2 model(**inputs)

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/peft/peft_model.py:1680, in PeftModelForCausalLM.forward(self, input_ids, attention_mask, inputs_embeds, labels, output_attentions, output_hidden_states, return_dict, task_ids, **kwargs)
   1677 if peft_config.peft_type == PeftType.PREFIX_TUNING:
   1678     # overwrite past_kv in kwargs
   1679     kwargs[\"past_key_values\"] = self.get_prompt(batch_size)
-> 1680     return self.base_model(input_ids=input_ids, inputs_embeds=inputs_embeds, **kwargs)
   1681 else:
   1682     if inputs_embeds is None:

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/transformers/models/idefics/modeling_idefics.py:1496, in IdeficsForVisionText2Text.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, pixel_values, image_encoder_embeddings, perceiver_embeddings, image_attention_mask, labels, use_cache, output_attentions, output_hidden_states, interpolate_pos_encoding, return_dict)
   1493 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   1495 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> 1496 outputs = self.model(
   1497     input_ids=input_ids,
   1498     attention_mask=attention_mask,
   1499     position_ids=position_ids,
   1500     past_key_values=past_key_values,
   1501     inputs_embeds=inputs_embeds,
   1502     pixel_values=pixel_values,
   1503     image_encoder_embeddings=image_encoder_embeddings,
   1504     perceiver_embeddings=perceiver_embeddings,
   1505     image_attention_mask=image_attention_mask,
   1506     use_cache=use_cache,
   1507     output_attentions=output_attentions,
   1508     output_hidden_states=output_hidden_states,
   1509     interpolate_pos_encoding=interpolate_pos_encoding,
   1510     return_dict=return_dict,
   1511 )
   1513 hidden_states = outputs[0]
   1514 logits = self.lm_head(hidden_states)

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/transformers/models/idefics/modeling_idefics.py:1329, in IdeficsModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, pixel_values, image_encoder_embeddings, perceiver_embeddings, image_attention_mask, use_cache, output_attentions, output_hidden_states, interpolate_pos_encoding, return_dict)
   1312     layer_outputs = self._gradient_checkpointing_func(
   1313         vblock,
   1314         decoder_layer,
   (...)
   1326         self.gated_cross_attn_layers,
   1327     )
   1328 else:
-> 1329     layer_outputs = vblock(
   1330         decoder_layer,
   1331         hidden_states,
   1332         attention_mask=attention_mask,
   1333         position_ids=position_ids,
   1334         past_key_value=past_key_value,
   1335         image_hidden_states=image_hidden_states,
   1336         image_attention_mask=image_attention_mask,
   1337         cross_attention_gate=cross_attention_gate,
   1338         output_attentions=output_attentions,
   1339         use_cache=use_cache,
   1340         layer_idx=idx,
   1341         cross_layer_interval=self.cross_layer_interval,
   1342         gated_cross_attn_layers=self.gated_cross_attn_layers,
   1343     )
   1345 hidden_states = layer_outputs[0]
   1347 if use_cache:

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/transformers/models/idefics/modeling_idefics.py:1293, in IdeficsModel.forward.<locals>.vblock(main_block, hidden_states, attention_mask, position_ids, past_key_value, image_hidden_states, image_attention_mask, cross_attention_gate, output_attentions, use_cache, layer_idx, cross_layer_interval, gated_cross_attn_layers)
   1281     outputs = xblock(
   1282         hidden_states,
   1283         attention_mask=attention_mask,
   (...)
   1289         past_key_value=None,  # not implemented
   1290     )
   1291     hidden_states = outputs[0]
-> 1293 layer_outputs = main_block(
   1294     hidden_states,
   1295     attention_mask=attention_mask,
   1296     position_ids=position_ids,
   1297     past_key_value=past_key_value,
   1298     output_attentions=output_attentions,
   1299     use_cache=use_cache,
   1300 )
   1302 return layer_outputs

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/transformers/models/idefics/modeling_idefics.py:749, in IdeficsDecoderLayer.forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache)
    746 hidden_states = self.input_layernorm(hidden_states)
    748 # Self Attention
--> 749 hidden_states, self_attn_weights, present_key_value = self.self_attn(
    750     hidden_states=hidden_states,
    751     attention_mask=attention_mask,
    752     position_ids=position_ids,
    753     past_key_value=past_key_value,
    754     output_attentions=output_attentions,
    755     use_cache=use_cache,
    756 )
    757 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    758 hidden_states = residual + hidden_states

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1532, in Module._wrapped_call_impl(self, *args, **kwargs)
   1530     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1531 else:
-> 1532     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/torch/nn/modules/module.py:1541, in Module._call_impl(self, *args, **kwargs)
   1536 # If we don't have any hooks, we want to skip the rest of the logic in
   1537 # this function, and just call forward.
   1538 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1539         or _global_backward_pre_hooks or _global_backward_hooks
   1540         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1541     return forward_call(*args, **kwargs)
   1543 try:
   1544     result = None

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/transformers/models/idefics/modeling_idefics.py:640, in IdeficsAttention.forward(self, hidden_states, key_value_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache)
    638 if not is_cross_attention:
    639     cos, sin = self.rotary_emb(value_states, seq_len=max(kv_seq_len, q_len))
--> 640     query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
    641 # [bsz, nh, t, hd]
    643 if past_key_value is not None:
    644     # reuse k, v, self_attention

File ~/miniconda3/envs/icl/lib/python3.12/site-packages/transformers/models/idefics/modeling_idefics.py:509, in apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim)
    507 cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    508 sin = sin[position_ids].unsqueeze(unsqueeze_dim)
--> 509 q_embed = (q * cos) + (rotate_half(q) * sin)
    510 k_embed = (k * cos) + (rotate_half(k) * sin)
    511 return q_embed, k_embed

RuntimeError: The size of tensor a (2) must match the size of tensor b (3) at non-singleton dimension 2"
}
```

As we can see, due to appearance of `past_key_values`, the key states and the value states will be extent to a larger length, while query states won't. The `cos` here is generated by `position_ids` whose length is same as `attention_mask`.

Comments:
Seems like a duplicate of https://github.com/huggingface/transformers/issues/33763. Can you try installing transformers from source with `!pip install --upgrade git+https://github.com/huggingface/transformers.git`

I still get the same bug. I also got the following warnings:
Before loading checkpoint shards:
```
IdeficsForVisionText2Text has generative capabilities, as `prepare_inputs_for_generation` is explicitly overwritten. However, it doesn't directly inherit from `GenerationMixin`. From 👉v4.50👈 onwards, `PreTrainedModel` will NOT inherit from `GenerationMixin`, and this model will lose the ability to call `generate` and other related functions.
  - If you're using `trust_remote_code=True`, you can get rid of this warning by loading the model with an auto class. See https://huggingface.co/docs/transformers/en/model_doc/auto#auto-classes
  - If you are the owner of the model architecture code, please modify your model class such that it inherits from `GenerationMixin` (after `PreTrainedModel`, otherwise you'll get an exception).
  - If you are not the owner of the model architecture class, please contact the model code owner to update it.
Instantiating IdeficsAttention without passing a `layer_idx` is not recommended and will lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` when creating this class.
```
After loading checkpoint shards:
```
We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class (https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)
```

This warning is possibly from [here](https://github.com/huggingface/transformers/blob/22266be9703c2e30e4734fe7a91f7abab9ff6b23/src/transformers/models/idefics/modeling_idefics.py#L801). Did you forget to add `layer_idx` when instantiates `IdeficsAttention`?
```
Instantiating IdeficsAttention without passing a `layer_idx` is not recommended and will lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` when creating this class.
```

Hmm, seems like the `virtual_tokens` are being added to the attention mask and thus to position ids (making it sequence_length=3), but the input ids are left untouched. Running the same code without prefix tuning works without errors

cc @BenjaminBossan for peft related questions

Thanks for the ping. I'm trying to work on a solution via https://github.com/huggingface/peft/pull/2096. It's not trivial to fix as I'm not a transformers expert, any help is welcome :)

Note that this bug is not related to `DynamicCache`. In transformers v4.44 (new cache facilities are not introduced at this version), my code snippet will not work as well. @zucchini-nlp has mentioned that 
> virtual_tokens are being added to the attention mask and thus to position ids (making it sequence_length=3), but the input ids are left untouched.

I don't know what the correct way is, but it seems like adding some dummy token ids to the beginning of `input_ids` can fix this problem?

Ah I see, thanks for clarifying. I'll check if this can be rolled into the same PR or warrants a separate fix.

Could @zucchini-nlp or @BenjaminBossan provide a temporary solution? I'm currently conducting a series of important experiments, which are pivot to my upcoming CVPR paper. I really need it fixed as soon as possible 😢

I'm sorry to hear that. With these things, it's hard to predict if a simple change or a bigger refactor is required. If it's possible for you, you could try downgrading transformers and PEFT until you find versions that works.

Fine, I will wait for any progress. Thank you all for your magnificent work.

@zucchini-nlp To be honest, I think this bug should be fixed by transformers. We can observe from [here](https://github.com/huggingface/transformers/blob/0c4c2d7e07c299eb8781c0d525b297684297ce06/src/transformers/models/idefics/modeling_idefics.py#L657)
```python
if not is_cross_attention:
    cos, sin = self.rotary_emb(value_states, seq_len=max(kv_seq_len, q_len))
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
        # [bsz, nh, t, hd]
```
The `max(kv_seq_len, q_len)` operation indicates that the length of kv and q can be different. However, in [apply_rotary_pos_emb](https://github.com/huggingface/transformers/blob/0c4c2d7e07c299eb8781c0d525b297684297ce06/src/transformers/models/idefics/modeling_idefics.py#L493-L518), it always adds to k and q the positional encoding with the maximum length between k and q, which leads to this problem. I suppose a possible solution could be:
```python
  cos = cos[position_ids].unsqueeze(unsqueeze_dim)
  sin = sin[position_ids].unsqueeze(unsqueeze_dim)
  min_len = min(sin.shape[2], q.shape[2])
  q_embed = (q * cos[:, :, :min_len]) + (rotate_half(q) * sin[:, :, :min_len])
  min_len = min(sin.shape[2], k.shape[2])
  k_embed = (k * cos[:, :, :min_len]) + (rotate_half(k) * sin[:, :, :min_len])
```
What's more, [here](https://github.com/huggingface/transformers/blob/0c4c2d7e07c299eb8781c0d525b297684297ce06/src/transformers/models/idefics/modeling_idefics.py#L1690-L1691) should be fixed as well. In prefix tuning, `num_virual_tokens` ones are added to the front of `attention_mask`, which cause the shape mismatch between `input_ids` and `attention_mask`. Maybe the best practice can be allowing 4d attention_mask instead of 2d,

It is because we want to compute sin and cos for the max length that is available. For ex, if you had `kv_len=3` and `q_len=1`, we could get sin/cos for 3 positions `(0, 1, 2)`. Then these are sliced using position ids in `apply_rotary_embedding` to apply the RoPE on hidden states

The reason why we can have different length is that kv length accounts for current hidden states at hand and the prev cached states. So we should take the prev length and get RoPE for the actual length, which is further sliced with position ids. So the problem here is that `position ids` are not same shape as hidden states are, as mentioned earlier. The proposed solution will mess up the RoPE and thus lead to poor tuning/inference.

I still think it is the PEFT issue, since it seems to happen only with prefix tuning and the shape mismatch difference is exactly `num_virtual_tokens`. Actually I am not sure if prefix tuning worked in VLMs (if the issue is VLM-specific), as we have been mostly experimenting with QLORA internally when shipping the model

EDIT: Here is the source of inconsistency , but I am not very comfortable with prefix tuning so I'll leave that for PEFT team

https://github.com/huggingface/peft/blob/2a807359bd34df64dd8dd29c1bd1024125d13312/src/peft/peft_model.py#L1669-L1672

Thanks both of you for the further investigations. Hopefully I'll have a bit of spare time at my hands tomorrow to figure out if I can find a fix on the PEFT side.

I've tried to prepare an `attention_mask` with shape of `[batch_size, 1, seq_len, seq_len + num_virtual_tokens]`, but it still failed in `apply_rotary_pos_emb` with the same error. I really can't find any other way to make prefix tuning work successfully. Btw, don't you think the following warning should be fixed?

> This warning is possibly from [here](https://github.com/huggingface/transformers/blob/22266be9703c2e30e4734fe7a91f7abab9ff6b23/src/transformers/models/idefics/modeling_idefics.py#L801). Did you forget to add `layer_idx` when instantiates `IdeficsAttention`?
> 
> ```
> Instantiating IdeficsAttention without passing a `layer_idx` is not recommended and will lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` when creating this class.
> ```



Oh yeah, I'll hanlde the warnings :)

But I still think it's weird even I've prepared my own `attention_mask` with proper shape. Logically, q and kv can have different length. IDEFICS should support it but I can't find a way to achieve my goal.

> EDIT: Here is the source of inconsistency , but I am not very comfortable with prefix tuning so I'll leave that for PEFT team
> 
> https://github.com/huggingface/peft/blob/2a807359bd34df64dd8dd29c1bd1024125d13312/src/peft/peft_model.py#L1669-L1672

I tried changing some PEFT code to avoid the size mismatch of the `attention_mask` but I don't think that's right. For 2 tokens + 20 virtual tokens, we get 22 entries for `cos` and `sin`, which looks right. However, position ids is now only `[0, 1]`, not `[0, ..., 21]`.

I'm really not an expert on this, so I might well be wrong. But I wonder if the `key_states`, `query_states`, and `value_states` should not have different shapes. As is, the seq len is 2, which of course is because the seq len of `hidden_states` is also 2. Should it not be 22, given that we add 20 virtual tokens? However, I can't find in the code where the past cache is being used -- I only see it being updated and its `get_seq_length()` being called.

I still went with the fix above resulting in a sequence of only 2, but then we get an error when we try to update `past_key_values`:

> *** RuntimeError: Sizes of tensors must match except in dimension 2. Expected size 32 but got size 4 for tensor number 1 in the list.

So I think there is an issue with `past_key_values` having the wrong dimensions here. Does it need different shapes for Idefics?

https://github.com/huggingface/transformers/blob/d7950bff82b18c823193d17d72188c5e46d06c83/src/transformers/models/idefics/modeling_idefics.py#L664




@BenjaminBossan is there a PR or branch I can debug with? I think the position ids are only 2 because the input ids are also length=2. So when we make position ids, we make them as long as attention mask and then crop to match the input ids length. That is done to account for past cache


> As is, the seq len is 2, which of course is because the seq len of hidden_states is also 2. Should it not be 22, given that we add 20 virtual tokens? 

Yes, same thought. I wonder how these virtual tokens are/should be passed from peft? Is it done as a simple concat with input ids or you add then in past kv cache? If the first case, then indeed seems to be weird bug that loses 20 tokens on the way. In the second case, we'll need to investigate how the cache is updated and what is causing shape mismatch errors

> is there a PR or branch I can debug with?

You mean on PEFT? No, I think the normal `main` branch should be good.

> I wonder how these virtual tokens are/should be passed from peft? Is it done as a simple concat with input ids or you add then in past kv cache? 

In PEFT, we just generate the `past_key_values` and pass them on to transformers, with the expectation that transformers will correctly concatenate them. For causal LM and prefix tuning, that would be these lines:

https://github.com/huggingface/peft/blob/d9d3059e94eab8e4700f9889555feb8babd7d7f5/src/peft/peft_model.py#L1701-L1704

I did a further experiment where I used Llama instead of Idefics, assuming that this should work correctly. There, I found that we also only have a sequence of 2, so my assumption above was incorrect. However, what's different is that `position_ids` are `[20, 21]` instead of `[0, 1]`, i.e. the indices of the actual tokens, not of the cache. That looks right to me. 

Where do these `position_ids` come from? They are created here:

https://github.com/huggingface/transformers/blob/d7950bff82b18c823193d17d72188c5e46d06c83/src/transformers/models/llama/modeling_llama.py#L909-L915

In Idefics, the logic for `position_ids` is different, as it checks the `attention_mask`:

https://github.com/huggingface/transformers/blob/d7950bff82b18c823193d17d72188c5e46d06c83/src/transformers/models/idefics/modeling_idefics.py#L1211-L1219

If I remove the check, i.e. go to the `elif` part, I get the (presumably) correct `[20, 21]`. Unfortunately, after that, I still run into the error with updating the `past_key_values` that I mentioned earlier.

For reference, this is the code I used for testing:

```python
import peft
from transformers import IdeficsForVisionText2Text, AutoProcessor
from transformers import AutoModelForCausalLM
import sys
import torch

use_peft = True
use_idefics = True
device = torch.device("cuda:0")

if use_idefics:
    model_id = "HuggingFaceM4/idefics2-8b"
    model_id = "hf-internal-testing/tiny-random-IdeficsModel"
    model = IdeficsForVisionText2Text.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
    processor = AutoProcessor.from_pretrained(model_id, torch_dtype=torch.float16)

    if use_peft:
        model = peft.get_peft_model(
            model,
            peft.PrefixTuningConfig(
                peft_type="PREFIX_TUNING",
                task_type="CAUSAL_LM",
                num_virtual_tokens=20,
                token_dim=4096,
                num_transformer_submodules=1,
                num_attention_heads=32,
                num_layers=32,
                encoder_hidden_size=768,
            ),
        )

    max_size = model.config.image_size
    inputs = processor(["hello"]).to(device)
    inputs['pixel_values'] = inputs['pixel_values'][:, :, :, :max_size, :max_size]
    out = model(**inputs)
    print(out.logits)
else:
    model_id = "meta-llama/Llama-3.1-8B-Instruct"

    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16).to(device)
    if use_peft:
        model = peft.get_peft_model(
            model,
            peft.PrefixTuningConfig(
                peft_type="PREFIX_TUNING",
                task_type="CAUSAL_LM",
                num_virtual_tokens=20,
                num_transformer_submodules=1,
            ),
        )
    out = model(torch.tensor([[0, 1]]).to(device))
    print(out.logits)
```

@BenjaminBossan you're right! I see, idefics wasn't cropping position ids to match the input ids. I used your code and it seems to work after the position ids are fixed, if we use the "Idefics-9b" model (the example has idefics2-8b but not the correct class). So I guess we just need to fix that on transformers side, thanks for investigation!


Thanks for working on this fix so quickly @zucchini-nlp. Maybe @Kami-chanw you can test that branch with your original use case.

Thank you for your efforts!

@zucchini-nlp ATTENTION!! It seems that you forgot to fix [here](https://github.com/huggingface/transformers/blob/0c4c2d7e07c299eb8781c0d525b297684297ce06/src/transformers/models/idefics/modeling_idefics.py#L1690-L1691) in  #33907 . When I pass `labels` to model forward, I got an `IndexError` due to shape mismatch between `attention_mask` and `logits`.
```
{
	"name": "IndexError",
	"message": "The shape of the mask [1, 2] at index 1 does not match the shape of the indexed tensor [1, 1, 32002] at index 1",
	"stack": "---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[1], line 30
     16 model = peft.get_peft_model(
     17     model,
     18     peft.PrefixTuningConfig(
   (...)
     27     ),
     28 )
     29 inputs = processor([\"hello\"]).to(device)
---> 30 model(**inputs, labels=inputs[\"input_ids\"])

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/peft/peft_model.py:1680, in PeftModelForCausalLM.forward(self, input_ids, attention_mask, inputs_embeds, labels, output_attentions, output_hidden_states, return_dict, task_ids, **kwargs)
   1677 if peft_config.peft_type == PeftType.PREFIX_TUNING:
   1678     # overwrite past_kv in kwargs
   1679     kwargs[\"past_key_values\"] = self.get_prompt(batch_size)
-> 1680     return self.base_model(input_ids=input_ids, inputs_embeds=inputs_embeds, **kwargs)
   1681 else:
   1682     if inputs_embeds is None:

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/models/idefics/modeling_idefics.py:1643, in IdeficsForVisionText2Text.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, pixel_values, image_encoder_embeddings, perceiver_embeddings, image_attention_mask, labels, use_cache, output_attentions, output_hidden_states, interpolate_pos_encoding, return_dict, cache_position)
   1641 if attention_mask is not None:
   1642     shift_attention_mask = attention_mask[..., 1:].to(logits.device)
-> 1643     shift_logits = logits[..., :-1, :][shift_attention_mask != 0].contiguous()
   1644     shift_labels = labels[..., 1:][shift_attention_mask != 0].contiguous()
   1645 else:

IndexError: The shape of the mask [1, 2] at index 1 does not match the shape of the indexed tensor [1, 1, 32002] at index 1"
}
```
My code:
```python
import peft
from transformers import IdeficsForVisionText2Text, AutoProcessor
import sys
import torch

sys.path.insert(0, "..")
import config  # my own file

device = torch.device("cpu")
model = IdeficsForVisionText2Text.from_pretrained(
    config.idefics_9b_path, torch_dtype=torch.float16
).to(device)
processor = AutoProcessor.from_pretrained(
    config.idefics_9b_path, torch_dtype=torch.float16
)
model = peft.get_peft_model(
    model,
    peft.PrefixTuningConfig(
        peft_type="PREFIX_TUNING",
        task_type="CAUSAL_LM",
        num_virtual_tokens=1,
        token_dim=4096,
        num_transformer_submodules=1,
        num_attention_heads=32,
        num_layers=32,
        encoder_hidden_size=768,
    ),
)
inputs = processor(["hello"]).to(device)
model(**inputs, labels=inputs["input_ids"]) # add labels here
```

It also failed when generating something. You just need to replace `model(**inputs, labels=inputs["input_ids"])` with `model.generate(**inputs)` if you want to reproduce this bug.
```
{
	"name": "IndexError",
	"message": "index 2 is out of bounds for dimension 0 with size 2",
	"stack": "---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[1], line 30
     16 model = peft.get_peft_model(
     17     model,
     18     peft.PrefixTuningConfig(
   (...)
     27     ),
     28 )
     29 inputs = processor([\"hello\"]).to(device)
---> 30 model.generate(**inputs) # add labels here

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/peft/peft_model.py:1706, in PeftModelForCausalLM.generate(self, *args, **kwargs)
   1704             outputs = self.base_model.generate(*args, **kwargs)
   1705     else:
-> 1706         outputs = self.base_model.generate(**kwargs)
   1707 except:
   1708     self.base_model.prepare_inputs_for_generation = self.base_model_prepare_inputs_for_generation

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/utils/_contextlib.py:116, in context_decorator.<locals>.decorate_context(*args, **kwargs)
    113 @functools.wraps(func)
    114 def decorate_context(*args, **kwargs):
    115     with ctx_factory():
--> 116         return func(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/generation/utils.py:2139, in GenerationMixin.generate(self, inputs, generation_config, logits_processor, stopping_criteria, prefix_allowed_tokens_fn, synced_gpus, assistant_model, streamer, negative_prompt_ids, negative_prompt_attention_mask, **kwargs)
   2131     input_ids, model_kwargs = self._expand_inputs_for_generation(
   2132         input_ids=input_ids,
   2133         expand_size=generation_config.num_return_sequences,
   2134         is_encoder_decoder=self.config.is_encoder_decoder,
   2135         **model_kwargs,
   2136     )
   2138     # 12. run sample (it degenerates to greedy search when `generation_config.do_sample=False`)
-> 2139     result = self._sample(
   2140         input_ids,
   2141         logits_processor=prepared_logits_processor,
   2142         stopping_criteria=prepared_stopping_criteria,
   2143         generation_config=generation_config,
   2144         synced_gpus=synced_gpus,
   2145         streamer=streamer,
   2146         **model_kwargs,
   2147     )
   2149 elif generation_mode in (GenerationMode.BEAM_SAMPLE, GenerationMode.BEAM_SEARCH):
   2150     # 11. prepare beam search scorer
   2151     beam_scorer = BeamSearchScorer(
   2152         batch_size=batch_size,
   2153         num_beams=generation_config.num_beams,
   (...)
   2158         max_length=generation_config.max_length,
   2159     )

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/generation/utils.py:3099, in GenerationMixin._sample(self, input_ids, logits_processor, stopping_criteria, generation_config, synced_gpus, streamer, **model_kwargs)
   3096 model_inputs.update({\"output_hidden_states\": output_hidden_states} if output_hidden_states else {})
   3098 # forward pass to get next token
-> 3099 outputs = self(**model_inputs, return_dict=True)
   3101 if synced_gpus and this_peer_finished:
   3102     continue  # don't waste resources running the code we don't need

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/models/idefics/modeling_idefics.py:1616, in IdeficsForVisionText2Text.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, pixel_values, image_encoder_embeddings, perceiver_embeddings, image_attention_mask, labels, use_cache, output_attentions, output_hidden_states, interpolate_pos_encoding, return_dict, cache_position)
   1613 return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   1615 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-> 1616 outputs = self.model(
   1617     input_ids=input_ids,
   1618     attention_mask=attention_mask,
   1619     position_ids=position_ids,
   1620     past_key_values=past_key_values,
   1621     inputs_embeds=inputs_embeds,
   1622     pixel_values=pixel_values,
   1623     image_encoder_embeddings=image_encoder_embeddings,
   1624     perceiver_embeddings=perceiver_embeddings,
   1625     image_attention_mask=image_attention_mask,
   1626     use_cache=use_cache,
   1627     output_attentions=output_attentions,
   1628     output_hidden_states=output_hidden_states,
   1629     interpolate_pos_encoding=interpolate_pos_encoding,
   1630     return_dict=return_dict,
   1631     cache_position=cache_position,
   1632 )
   1634 hidden_states = outputs[0]
   1635 logits = self.lm_head(hidden_states)

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/models/idefics/modeling_idefics.py:1323, in IdeficsModel.forward(self, input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, pixel_values, image_encoder_embeddings, perceiver_embeddings, image_attention_mask, use_cache, output_attentions, output_hidden_states, interpolate_pos_encoding, return_dict, cache_position)
   1305     layer_outputs = self._gradient_checkpointing_func(
   1306         vblock,
   1307         decoder_layer,
   (...)
   1320         cache_position,
   1321     )
   1322 else:
-> 1323     layer_outputs = vblock(
   1324         decoder_layer,
   1325         hidden_states,
   1326         attention_mask=attention_mask,
   1327         position_ids=position_ids,
   1328         past_key_value=past_key_values,
   1329         image_hidden_states=image_hidden_states,
   1330         image_attention_mask=image_attention_mask,
   1331         cross_attention_gate=cross_attention_gate,
   1332         output_attentions=output_attentions,
   1333         use_cache=use_cache,
   1334         layer_idx=idx,
   1335         cross_layer_interval=self.cross_layer_interval,
   1336         gated_cross_attn_layers=self.gated_cross_attn_layers,
   1337         cache_position=cache_position,
   1338     )
   1340 hidden_states = layer_outputs[0]
   1342 if use_cache:

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/models/idefics/modeling_idefics.py:1285, in IdeficsModel.forward.<locals>.vblock(main_block, hidden_states, attention_mask, position_ids, past_key_value, image_hidden_states, image_attention_mask, cross_attention_gate, output_attentions, use_cache, layer_idx, cross_layer_interval, gated_cross_attn_layers, cache_position)
   1273     outputs = xblock(
   1274         hidden_states,
   1275         attention_mask=attention_mask,
   (...)
   1281         past_key_value=None,  # not implemented
   1282     )
   1283     hidden_states = outputs[0]
-> 1285 layer_outputs = main_block(
   1286     hidden_states,
   1287     attention_mask=attention_mask,
   1288     position_ids=position_ids,
   1289     past_key_value=past_key_value,
   1290     output_attentions=output_attentions,
   1291     use_cache=use_cache,
   1292     cache_position=cache_position,
   1293 )
   1295 return layer_outputs

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/models/idefics/modeling_idefics.py:722, in IdeficsDecoderLayer.forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position)
    719 hidden_states = self.input_layernorm(hidden_states)
    721 # Self Attention
--> 722 hidden_states, self_attn_weights, present_key_value = self.self_attn(
    723     hidden_states=hidden_states,
    724     attention_mask=attention_mask,
    725     position_ids=position_ids,
    726     past_key_value=past_key_value,
    727     output_attentions=output_attentions,
    728     use_cache=use_cache,
    729     cache_position=cache_position,
    730 )
    731 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
    732 hidden_states = residual + hidden_states

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1553, in Module._wrapped_call_impl(self, *args, **kwargs)
   1551     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1552 else:
-> 1553     return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/torch/nn/modules/module.py:1562, in Module._call_impl(self, *args, **kwargs)
   1557 # If we don't have any hooks, we want to skip the rest of the logic in
   1558 # this function, and just call forward.
   1559 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1560         or _global_backward_pre_hooks or _global_backward_hooks
   1561         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1562     return forward_call(*args, **kwargs)
   1564 try:
   1565     result = None

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/models/idefics/modeling_idefics.py:613, in IdeficsAttention.forward(self, hidden_states, key_value_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache, cache_position)
    611 if not is_cross_attention:
    612     cos, sin = self.rotary_emb(value_states, seq_len=max(kv_seq_len, q_len))
--> 613     query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
    614 # [bsz, nh, t, hd]
    616 if past_key_value is not None:
    617     # sin and cos are specific to RoPE models; cache_position needed for the static cache

File ~/miniconda3/envs/icl/lib/python3.9/site-packages/transformers/models/idefics/modeling_idefics.py:469, in apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim)
    448 def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    449     \"\"\"Applies Rotary Position Embedding to the query and key tensors.
    450 
    451     Args:
   (...)
    467         `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    468     \"\"\"
--> 469     cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    470     sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    471     q_embed = (q * cos) + (rotate_half(q) * sin)

IndexError: index 2 is out of bounds for dimension 0 with size 2"
}
```

Got it, seems like Idefics is one of the very few models that mask out logits using attention mask. But since in prefix-tuning the virtual tokens are always part of the cache, we fail to match sizes with attn mask. Will add this fix to the linked PR

Note that the current PR still cannot resolve `IndexError` when generate

That error is related to Idefics using cross attentions, so we init the cache with empty lists as of the last release. I think PEFT should handle tuple cache and DynamicCache differently internally, as for the cache object it is more reliable to call `get_seq_length()`.

After fixing that another error is from PEFT not updating the cache object with virtual tokens, which causes IndexError. While preparing inputs for generation, the check is `if model_kwargs["past_key_values"] is None and peft_config.peft_type == PeftType.PREFIX_TUNING` but cache can be not None and still empty. So this check can also be modified to handle new and old cache formats

cc @BenjaminBossan and maybe @gante 

As the generate error in related to PEFT-side, I'm merging the PR that fixes position ids slicing :)

> While preparing inputs for generation, the check is `if model_kwargs["past_key_values"] is None and peft_config.peft_type == PeftType.PREFIX_TUNING` but cache can be not None and still empty. So this check can also be modified to handle new and old cache formats

Could you please tell me which check you're referring to, so that I can fix it?

Btw when I run [my snippet](https://github.com/huggingface/transformers/issues/33852#issuecomment-2389184690) with your fix in #33907, I still get the error of the cache not being able to update ([in this line](https://github.com/zucchini-nlp/transformers/blob/ae678954109610f22abf2cd080ae36d796269060/src/transformers/models/idefics/modeling_idefics.py#L619)):

```
  [...]
  File "/tmp/transformers/src/transformers/models/idefics/modeling_idefics.py", line 619, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/transformers/src/transformers/cache_utils.py", line 447, in update
    self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
```

When I check the shapes, this is what I get:

```python
>>> self.key_cache[layer_idx].shape
torch.Size([1, 32, 20, 128])
>>> key_states.shape
torch.Size([1, 4, 2, 8])
```

Do you know if this is something that has yet to be fixed on the PEFT side?

@Kami-chanw checks in this method (https://github.com/huggingface/peft/blob/d9d3059e94eab8e4700f9889555feb8babd7d7f5/src/peft/peft_model.py#L1738)

The error you have is because you prob init PrefixTuning with incorrect num-layers and num-heads, can you double check all params? The script had dummy tiny model and thus params are not correct. 

> The error you have is because you prob init PrefixTuning with incorrect num-layers and num-heads, can you double check all params? The script had dummy tiny model and thus params are not correct.

:facepalm: Indeed, with the right config it works :tada: thanks

For future reference:

```python
model_id = "hf-internal-testing/tiny-random-IdeficsModel"
peft_config = peft.PrefixTuningConfig(
    peft_type="PREFIX_TUNING",
    task_type="CAUSAL_LM",
    num_virtual_tokens=20,
    token_dim=32,
    num_transformer_submodules=1,
    num_attention_heads=4,
    num_layers=5,
    encoder_hidden_size=768,
)
```