Qwen2FlashAttention sliding windows are applied to wrong layers

[hzhua]

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

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

[ArthurZucker]

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

[hzhua]

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.

[ArthurZucker]

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

[bigximik]

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:

  transformers/src/transformers/models/qwen2/configuration_qwen2.py

  Line 112 in f2ab182

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

• And its implementation:

  transformers/src/transformers/models/qwen2/modular_qwen2.py

  Line 74 in f2ab182

  and self.layer_idx >= self.config.max_window_layers

Which one is correct?
Thanks!