Repository: huggingface/transformers (Python) — Bug Fix

Issue:
`modeling_opt.py` if `previous_key_values` given and `attention_mask==None` the model throws an error.
### System Info

- `transformers` version: 4.26.1
- Platform: Linux-4.18.0-147.el8.x86_64-x86_64-with-glibc2.28
- Python version: 3.9.16
- Huggingface_hub version: 0.12.1
- PyTorch version (GPU?): 1.13.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: yes

### Who can help?

@ArthurZucker @younesbelkada 

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

## Code
1. Load opt/tokenizer
```py
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "facebook/opt-125m"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer  = AutoTokenizer.from_pretrained(model_name)
```
2. Precompute `past_key_values`
```py
text1 = "let's find a"
tokenized1 = tokenizer(text1, return_tensors='pt')
past_key_values = model(**tokenized1, use_cache=True)["past_key_values"]
```
4. Compute another set of values without `attention_mask`
```py
text2 = "bug"
tokenized2 = tokenizer(text2, return_tensors='pt')
model(input_ids=tokenized2["input_ids"], past_key_values=past_key_values)
# error! The mistakenly created an attention_mask that is too small.
```
(try `distilgpt2` and it will work)

## stack trace

```
Traceback (most recent call last):
  File "/home/gkressi1/opt/ldet/rate_in-context.py", line 334, in <module>
    main()
  File "/home/gkressi1/opt/ldet/rate_in-context.py", line 325, in main
    output_config = compute_surprisals(config=config, model_object=model_object)
  File "/home/gkressi1/opt/ldet/rate_in-context.py", line 219, in compute_surprisals
    output_rating = model_object.incontext(config, prompt_list)
  File "/home/gkressi1/opt/ldet/src/model_objects/model_hf_causal_lm_big.py", line 85, in incontext
    output = self.get_model_output(rest_prompt, use_cache=True)
  File "/home/gkressi1/opt/ldet/src/model_objects/model_hf_causal_lm_big.py", line 63, in get_model_output
    output = self.model(
  File "/home/gkressi1/.conda/envs/llm/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/gkressi1/.conda/envs/llm/lib/python3.9/site-packages/accelerate/hooks.py", line 158, in new_forward
    output = old_forward(*args, **kwargs)
  File "/home/gkressi1/.conda/envs/llm/lib/python3.9/site-packages/transformers/models/opt/modeling_opt.py", line 932, in forward
    outputs = self.model.decoder(
  File "/home/gkressi1/.conda/envs/llm/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/gkressi1/.conda/envs/llm/lib/python3.9/site-packages/transformers/models/opt/modeling_opt.py", line 639, in forward
    attention_mask = self._prepare_decoder_attention_mask(
  File "/home/gkressi1/.conda/envs/llm/lib/python3.9/site-packages/transformers/models/opt/modeling_opt.py", line 546, in _prepare_decoder_attention_mask
    expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
RuntimeError: The size of tensor a (93) must match the size of tensor b (1679) at non-singleton dimension 3
```

### Expected behavior

The model should create the attention mask by itself and not throw an error.

From the surface, this seems to be an easy fix:
1. Delete line [635](https://github.com/huggingface/transformers/blob/ae54e3c3b18bac0832ad62ea9b896dfd52a09850/src/transformers/models/opt/modeling_opt.py#L635) and [636](https://github.com/huggingface/transformers/blob/ae54e3c3b18bac0832ad62ea9b896dfd52a09850/src/transformers/models/opt/modeling_opt.py#L635)
2. Move line [639-642](https://github.com/huggingface/transformers/blob/ae54e3c3b18bac0832ad62ea9b896dfd52a09850/src/transformers/models/opt/modeling_opt.py#L639) of what is currently line [637](https://github.com/huggingface/transformers/blob/ae54e3c3b18bac0832ad62ea9b896dfd52a09850/src/transformers/models/opt/modeling_opt.py#L637)
3. Check TF/Flax models (?).

All the best!


Discussion:
Hey! Thanks for submitting this issue! 
Passing attention  maks solves the problem, and usually we expect to pass attention masks when you are using the `past_key_values`(for example in generate). It is debatable whether the default behaviour should rely on the past_key_values. 
Do you have a specific usage in mind? 

The following works as expected: 
```python 
attn = torch.cat((tokenized1["attention_mask"], tokenized2["attention_mask"]), -1)
text2 = "bug"
tokenized2 = tokenizer(text2, return_tensors='pt')
model(input_ids=tokenized2["input_ids"], past_key_values=past_key_values,attention_mask =attn)
```
This way is the expected usage. When training or doing an inference, you should probably be in a for loop where the attention mask is defined based on the entire input. 

I agree that manually adding the attention_mask is an easy fix.

I am using a shared context as `past_key_values` and then computing different model outputs given the context. In that case I save the contexts `past_key_values` and use them later on. It is easy to recompute/save the contexts attention_mask and concat it for every output - but
* OPT model behavior is inconsistent to other model's I have been using (gpt-neo, bloom)
* it is [not documented](https://huggingface.co/docs/transformers/v4.26.1/en/model_doc/opt#transformers.OPTForCausalLM.forward.past_key_values) that the expected usage is passing the `attention_mask` when using `past_key_values`
* the thrown error is not descriptive of the issue

I do not understand what you mean with "default behaviour should rely on the past_key_values" - it seems to me that default behavior is not affected by changing this: line [636](https://github.com/huggingface/transformers/blob/ae54e3c3b18bac0832ad62ea9b896dfd52a09850/src/transformers/models/opt/modeling_opt.py#L636) seems to have exactly the same job that [639 - 642](https://github.com/huggingface/transformers/blob/ae54e3c3b18bac0832ad62ea9b896dfd52a09850/src/transformers/models/opt/modeling_opt.py#L639) has, just that it does not take into account `past_key_values` introducing the deviation of model behavior to other models.

I can understand if you say that passing `attention_mask` is expected behavior for using `past_key_values`, but maybe that could be mentioned somewhere?
Totally agree with you, will open a PR to adress this. I think this was also blocking us from adding the ONNX config for this model! 
Thanks for this 😉 


Resolving patch:
diff --git a/src/transformers/models/opt/modeling_opt.py b/src/transformers/models/opt/modeling_opt.py
--- a/src/transformers/models/opt/modeling_opt.py
+++ b/src/transformers/models/opt/modeling_opt.py
@@ -631,19 +631,21 @@ def forward(
         else:
             raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
-        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
-
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
 
+        batch_size, seq_length = input_shape
+        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+        # required mask seq length can be calculated via length of past
+        mask_seq_length = past_key_values_length + seq_length
+
         # embed positions
         if attention_mask is None:
-            attention_mask = torch.ones(inputs_embeds.shape[:2], dtype=torch.bool, device=inputs_embeds.device)
-        pos_embeds = self.embed_positions(attention_mask, past_key_values_length)
-
-        attention_mask = self._prepare_decoder_attention_mask(
+            attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+        causal_attention_mask = self._prepare_decoder_attention_mask(
             attention_mask, input_shape, inputs_embeds, past_key_values_length
         )
+        pos_embeds = self.embed_positions(attention_mask, past_key_values_length)
 
         if self.project_in is not None:
             inputs_embeds = self.project_in(inputs_embeds)
@@ -694,14 +696,14 @@ def custom_forward(*inputs):
                 layer_outputs = torch.utils.checkpoint.checkpoint(
                     create_custom_forward(decoder_layer),
                     hidden_states,
-                    attention_mask,
+                    causal_attention_mask,
                     head_mask[idx] if head_mask is not None else None,
                     None,
                 )
             else:
                 layer_outputs = decoder_layer(
                     hidden_states,
-                    attention_mask=attention_mask,
+                    attention_mask=causal_attention_mask,
                     layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                     past_key_value=past_key_value,
                     output_attentions=output_attentions,


Test patch:
diff --git a/tests/models/opt/test_modeling_opt.py b/tests/models/opt/test_modeling_opt.py
--- a/tests/models/opt/test_modeling_opt.py
+++ b/tests/models/opt/test_modeling_opt.py
@@ -182,6 +182,19 @@ def create_and_check_decoder_model_past_large_inputs(self, config, inputs_dict):
         # test that outputs are equal for slice
         self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
 
+        # test no attention_mask works
+        outputs = model(input_ids, attention_mask=attention_mask, head_mask=head_mask, use_cache=True)
+        _, past_key_values = outputs.to_tuple()
+        output_from_no_past = model(next_input_ids)["last_hidden_state"]
+
+        output_from_past = model(next_tokens, past_key_values=past_key_values)["last_hidden_state"]
+
+        random_slice_idx = ids_tensor((1,), output_from_past.shape[-1]).item()
+        output_from_no_past_slice = output_from_no_past[:, -3:, random_slice_idx].detach()
+        output_from_past_slice = output_from_past[:, :, random_slice_idx].detach()
+        # test that outputs are equal for slice
+        self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))
+
 
 @require_torch
 class OPTModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):