Qwen2LM performs well when using flash_attention_2 or SDPA, but its performance drops when using the original attention implementation (i.e., attn_implementation="eager").
Author: WellDonePF
Pull request: no
Labels: name: bug

name: Generation

name: Flash Attention

### System Info

- transformers version: 4.43.0.dev0
- Platform: Linux-4.19.91-014.15-kangaroo.alios7.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.14
- Huggingface_hub version: 0.25.0
- Safetensors version: 0.4.5
- Accelerate version: 0.34.2
- PyTorch version (GPU): 2.1.2+cu121
- Using distributed or parallel set-up in script: false
- GPU type: NVIDIA H800

### Who can help?
@ArthurZucker @zucchini-nlp 


#### Problem
Some models perform well when using flash_attention_2 or SDPA, but their performance drops significantly when using the original attention (i.e., attn_implementation="eager"). For example, Qwen2LM. 

I found this issue when working with the lmms-lab/llava-onevision-qwen2-7b-ov model and qwen2vl.（the transformers library is the latest version.）

by the way, is there a way to extract the attention score matrix (the matrix after applying softmax) when using flash attention 2 or SDPA?


### 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

```python
import json
import argparse
from PIL import Image
import copy
import torch
import accelerate

from llava.model.builder import load_pretrained_model
from llava.mm_utils import process_images, tokenizer_image_token
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN
from llava.conversation import conv_templates

def main(args):
    accelerator = accelerate.Accelerator()
    llava_model_args = {
            "multimodal": True,
        }
    overwrite_config = {}
    overwrite_config["image_aspect_ratio"] = "pad"
    llava_model_args["overwrite_config"] = overwrite_config

    tokenizer, model, image_processor, max_length = load_pretrained_model(
        model_path=args.pretrained_model_name_or_path,
        model_base=None,
        model_name=args.model_name,
        device_map=accelerator.device,
        attn_implementation=args.attn_implementation,
        **llava_model_args
    )
    model.eval()

    info = json.loads(open(args.data_path,'r',encoding="utf-8").read())
    model = accelerator.prepare(model)
    if accelerator.num_processes > 1:
        model = model.module

    cnt = -1
    for item in info:
        cnt += 1
        context = item["context"]
        context = context.replace("<ImageHere>", DEFAULT_IMAGE_TOKEN)
        images = item["raw_img_list"]
        images = [Image.open(img) for img in images]
        images = process_images(images, image_processor, model.config)
        images = [img.to(dtype=torch.float16, device=accelerator.device) for img in images]

        conv = copy.deepcopy(conv_templates[args.conv_template])
        conv.append_message(conv.roles[0], context)
        conv.append_message(conv.roles[1], None)
        prompt_question = conv.get_prompt()

        input_ids = tokenizer_image_token(prompt_question, tokenizer, IMAGE_TOKEN_INDEX, return_tensors="pt").unsqueeze(0).to(accelerator.device)
        image_sizes = [img.size for img in images]

        try:
            cont = model.generate(
                input_ids,
                images=images,
                image_sizes=image_sizes,
                do_sample=False,
                temperature=0,
                max_new_tokens=4096,
            )
        except Exception as e:
            print(e)
            continue
        result = tokenizer.batch_decode(cont, skip_special_tokens=True)
        result = result[0]

        print(result)


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--pretrained_model_name_or_path", type=str, default="lmms-lab/llava-onevision-qwen2-7b-ov")
    parser.add_argument("--model_name", type=str, default="llava_qwen")
    parser.add_argument("--root_path", type=str, default="******")
    parser.add_argument("--data_path", type=str, default="******")
    parser.add_argument("--attn_implementation", type=str, default="eager", choices=["eager", "flash_attention_2", "sdpa"])
    parser.add_argument("--conv_template", type=str, default="qwen_1_5")

    return parser.parse_args()


if __name__ == "__main__":
    main(parse_args())

```

### Expected behavior

I am testing the llava-onevision model with multiple images and image-text interleaved. Expecting it to output the answer to a multiple-choice question, but instead, it outputs '!!!!!!!!!!!!!!!!!!!!!!'.

Comments:
Hey! Could you give a bit more details about what you mean when you say `performance drops` ? 
- generation is worse
- generation throughput is worse
- latency is worse 
(performance alone can refer to a lot of stuff). 
We expect generations to be the same, but performance (memory consumption / token/s ) to be worse

> * generation is worse

Thanks for your reply. I mean the generated is worse.

For the "lmms-lab/llava-onevision-qwen2-7b-ov" model, if the original attention is used for multi-image and interleaving text and images testing, all several hundred test cases output "!!!" (transformers==4.43.0.dev0, transformers==latest ver).

For Qwen2VL, the test result acc will drop from 90%+ to 80% (transformers==latest ver) if using original attention.

However, for the 'lmms-lab/llava-onevision-qwen2-0.5b-ov' model, there is no such issue when using the original attention.

They both use Qwen2LM, so I guess if this could be some bug in the Qwen2 code. 


> Hey! Could you give a bit more details about what you mean when you say `performance drops` ?
> 
> * generation is worse
> * generation throughput is worse
> * latency is worse
>   (performance alone can refer to a lot of stuff).
>   We expect generations to be the same, but performance (memory consumption / token/s ) to be worse

https://github.com/LLaVA-VL/LLaVA-NeXT/issues/254#issuecomment-2365026841

@WellDonePF , hmm that is weird! I could reproduce it for llava-ov and also using the same code with "Qwen/Qwen2-7B-instruct" fails to generate normal text. But the 0.5B checkpoint on Qwen2 is okay. 

So I believe it is from the model weights for 7B qwen, especially because generating in bf16 precision helps to boost up the quality back. In terms of code,  it is same code as in all LLMs, and we haven't seen any failures with sdpa-eager mismatch with gibberish generation before. I'd recommend to try bf16 for now

> @WellDonePF , hmm that is weird! I could reproduce it for llava-ov and also using the same code with "Qwen/Qwen2-7B-instruct" fails to generate normal text. But the 0.5B checkpoint on Qwen2 is okay.
> 
> So I believe it is from the model weights for 7B qwen, especially because generating in bf16 precision helps to boost up the quality back. In terms of code, it is same code as in all LLMs, and we haven't seen any failures with sdpa-eager mismatch with gibberish generation before. I'd recommend to try bf16 for now

I've tried bf16. It works for easy task, but the problem still exists when task become complex.

## EASY TASK
[QUESTION]
'Deliver concise, accurate answers to questions prompted by the images\n\nContext:  \n\nQuestion: What is the special magic Thimphu number in the images?\nRespond with the answer and your thought process for reasoning.'

[MODEL_OUTPUT]
'!425900'

[EXPECTED ANSWER]
'425900'

## COMPLEX TASK
[QUESTION]
I will give you some pictures, and each group of pictures will correspond to a question. Please answer it briefly. You must choose your answer from the Choice List. \nQuestion: what is the title for the name robert kauk ?\nChoice list: \nA. VICE PRESIDENT, BRUSH BEET GROWERS ASS'N\nB. TREASURER, BRUSH BEET GROWERS ASS'N\nC. PRESIDENT, BRUSH BEET GROWERS ASS'N\nD. SECRETARY, BRUSH BEET GROWERS ASS'N\n\nRespond with the answer and your thought process for reasoning."

[MODEL_OUTPUT]
'!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!'

Forgot that we already had several things going on with Qwen2, so might be related to https://github.com/huggingface/transformers/pull/33317

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.