Repository: huggingface/transformers (Python)
Date: 2024-03-07 15:56:14+00:00
Task type: Bug Fix

Issue:
Sliding window inconsistency between PyTorch and Flax
### System Info

transformers main (ae49b218c), Python 3.10.8

### Who can help?

@ArthurZucker, @sanchit-gandhi

### Reproduction

The attention `sliding_window` has different interpretation for PyTorch and Flax. Here's are matching examples:

**PyTorch**

```python
from transformers import MistralModel
import torch

model = MistralModel.from_pretrained("hf-internal-testing/tiny-random-MistralModel", sliding_window=2)

inputs = {
  "input_ids": torch.tensor([[10, 20, 30, 40, 50, 60, 70, 80, 0, 0]]),
  "attention_mask": torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
}

outputs = model(**inputs)

print(outputs.last_hidden_state[:, 1:4, 1:4])
```

**Flax**

```python
from transformers import FlaxMistralModel
import jax.numpy as jnp

model = FlaxMistralModel.from_pretrained("hf-internal-testing/tiny-random-MistralModel", sliding_window=2, from_pt=True)

inputs = {
  "input_ids": jnp.array([[10, 20, 30, 40, 50, 60, 70, 80, 0, 0]]),
  "attention_mask": jnp.array([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0]])
}

outputs = model(**inputs)

print(outputs.last_hidden_state[:, 1:4, 1:4])
```

Both snippets return different results, however, if we use `sliding_window=3` in the PyTorch version, the results are the same.

In the Flax implementation, `sliding_window=2` means that a position will attend to self, and two previous position inclusive (which intuitively seems correct to me). It looks like in the PyTorch version it is not inclusive. Which behaviour is expected?

### Expected behavior

The `sliding_window` meaning to be consistent.


Issue discussion:
Hey!  Pretty sure `MistralSdpaAttention` does not support sliding window yet!  Are you using `attn_implementation="flash_attention_2"`? 
@ArthurZucker I'm using the default implementation on the CPU, I've just checked to make sure and it's "eager". Initially I thought the issues may be in flash_attn, but you made me realise it's obviously not used on the CPU, so I tracked it down and I think there is an off by one error here:

https://github.com/huggingface/transformers/blob/3f60d11a8750992287cd0d1f3dbc9df6ffc34288/src/transformers/modeling_attn_mask_utils.py#L169


I think one way to fix it would be this:

```python
if sliding_window is not None:
    diagonal = past_key_values_length - sliding_window - 1

    context_mask = torch.tril(torch.ones_like(mask, dtype=torch.bool), diagonal=diagonal)
    mask.masked_fill_(context_mask, torch.finfo(dtype).min)
```
By a quick glance, the initial implementation ([ref](https://github.com/huggingface/transformers/blame/f09a081d2765c6535256b0e2d65bf54fc03f7fee/src/transformers/models/mistral/modeling_mistral.py#L81-L88)) was actually inclusive, and it got lost during refactoring.

[flash_attn](https://github.com/Dao-AILab/flash-attention) also says "inclusive" in function docs, so I think that's the expected behaviour (and does make more sense).
would you like to open a PR? 🤗 

Solution patch:
diff --git a/src/transformers/modeling_attn_mask_utils.py b/src/transformers/modeling_attn_mask_utils.py
--- a/src/transformers/modeling_attn_mask_utils.py
+++ b/src/transformers/modeling_attn_mask_utils.py
@@ -164,10 +164,10 @@ def _make_causal_mask(
 
         # add lower triangular sliding window mask if necessary
         if sliding_window is not None:
-            diagonal = past_key_values_length - sliding_window + 1
+            diagonal = past_key_values_length - sliding_window - 1
 
-            context_mask = 1 - torch.triu(torch.ones_like(mask, dtype=torch.int), diagonal=diagonal)
-            mask.masked_fill_(context_mask.bool(), torch.finfo(dtype).min)
+            context_mask = torch.tril(torch.ones_like(mask, dtype=torch.bool), diagonal=diagonal)
+            mask.masked_fill_(context_mask, torch.finfo(dtype).min)
 
         return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
 


Test patch:
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -1673,7 +1673,7 @@ def check_to_causal(self, mask_converter, q_len, kv_len, bsz=3):
     def compute_num_context_mask(self, kv_len, context, q_len):
         # This function computes the # of attention tokens that are added for
         # the sliding window
-        c_mask_len = kv_len - context
+        c_mask_len = kv_len - context - 1
         num_mask_triangle = c_mask_len * (c_mask_len + 1) // 2
         cut_mask_len = max(c_mask_len - q_len, 0)
         num_cut_mask = cut_mask_len * (cut_mask_len + 1) // 2