Qwen2FlashAttention sliding windows are applied to wrong layers
Author: hzhua
Pull request: no
Labels: name: bug

According to the definition of max_window_layers

```max_window_layers  — The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.````

When `max_window_layers == num_hidden_layers` all layers should be applied with SWA. But current implementation only uses SWA when `self.layer_idx >= self.config.max_window_layers`, making all all layers using full attention instead.

https://github.com/huggingface/transformers/blob/fc269f77da72d4c65b2e71e6d4896cd16c6f1e76/src/transformers/models/qwen2/modular_qwen2.py#L71C1-L75C11

Changing `self.layer_idx >= self.config.max_window_layers` to `self.layer_idx < self.config.max_window_layers` may solve issue.

Please correct me if my understanding of max_window_layers is wrong.



### Who can help?

@ArthurZucker 

### Reproduction

1. Change the `config.json` of a Qwen2 model to use sliding window

```
{
  "architectures": [
    "Qwen2ForCausalLM"
  ],
  "attention_dropout": 0.0,
  "bos_token_id": 151643,
  "eos_token_id": 151643,
  "hidden_act": "silu",
  "hidden_size": 2048,
  "initializer_range": 0.02,
  "intermediate_size": 11008,
  "max_position_embeddings": 32768,
  "max_window_layers": 36,
  "model_type": "qwen2",
  "num_attention_heads": 16,
  "num_hidden_layers": 36,
  "num_key_value_heads": 2,
  "rms_norm_eps": 1e-06,
  "rope_theta": 1000000.0,
  "sliding_window": 500,
  "tie_word_embeddings": true,
  "torch_dtype": "bfloat16",
  "transformers_version": "4.40.1",
  "use_cache": true,
  "use_mrope": false,
  "use_sliding_window": true,
  "vocab_size": 151936
}
````

2. Add `print(sliding_window, self.config.max_window_layers, self.layer_idx)` to the forward function of `Qwen2FlashAttention2`.
3. run generation with a Qwen2 model

```python
from transformers import AutoTokenizer, Qwen2ForCausalLM
import torch

# Load the model and tokenizer
model = Qwen2ForCausalLM.from_pretrained("/home/hanzhenhua/Qwen2.5-3B", attn_implementation="flash_attention_2", torch_dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained("/home/hanzhenhua/Qwen2.5-3B")

# Move the model to GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Prepare the input prompt
prompt = '\n'.join(["Hey, are you conscious? Can you talk to me? " for _ in range(300)])
print(prompt)
# Tokenize the input
input_ids = tokenizer.encode(prompt, return_tensors="pt")
input_ids = input_ids.to(device)

# Generate text
# Generate output with a maximum of 50 tokens
output = model.generate(input_ids, max_new_tokens=50)

# Decode the output tokens to text
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

print(generated_text)
```


### Expected behavior

The debug printer shows `sliding_window = None`, which means sliding_window is not taking effect in flash attention.
```None 36 0
None 36 1
None 36 2
None 36 3
None 36 4
None 36 5
None 36 6
None 36 7
None 36 8
None 36 9
None 36 10
None 36 11
None 36 12
None 36 13
None 36 14
None 36 15
None 36 16
None 36 17
None 36 18
None 36 19
None 36 20
None 36 21
None 36 22
None 36 23
None 36 24
None 36 25
None 36 26
None 36 27
None 36 28
None 36 29
None 36 30
None 36 31
None 36 32
None 36 33
None 36 34
None 36 35
```

Comments:
Hey! I think it depends on the size of the model no? 

Vanilla Qwen2 model weight does not use sliding window so this bug does not take effect. I was trying to train a new version based on Qwen2-3B architecture with a sliding window, thus met this issue.

Yep sorry #36162 should fix it! Thanks for reporting!

Hello, @ArthurZucker,  

It seems that #36162 does not fix the issue described above by @hzhua.  

There is still a contradiction between:  
- The description of `max_window_layers`: https://github.com/huggingface/transformers/blob/f2ab182dca2a9aa4a053a1903bd4b46b90ff9f78/src/transformers/models/qwen2/configuration_qwen2.py#L112
- And its implementation: https://github.com/huggingface/transformers/blob/f2ab182dca2a9aa4a053a1903bd4b46b90ff9f78/src/transformers/models/qwen2/modular_qwen2.py#L74

Which one is correct?  
Thanks!