get warning: use Flash Attention 2.0 without specifying a torch dtype
Author: pspdada
Pull request: no
Labels: name: bug

name: Vision

name: Flash Attention

name: Multimodal

### System Info

- `transformers` version: 4.45.2
- Platform: Linux-5.15.0-120-generic-x86_64-with-glibc2.35
- Python version: 3.10.15
- Huggingface_hub version: 0.25.2
- Safetensors version: 0.4.5
- Accelerate version: 1.0.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0+cu124 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: None

### Who can help?

@amyeroberts, @qubvel

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

I create a llava 1.5 model and use flashattention 2 when specifying the torch_dtype=torch.float16 with the code below, then I get a warning: You are attempting to use Flash Attention 2.0 without specifying a torch dtype. This might lead to unexpected behaviour.
```    
def _create_v_1_5(self) -> tuple[LlavaForConditionalGeneration, LlavaProcessor]:
        model_name = f"llava-hf/llava-1.5-{self.model_size}-hf"
        processor: LlavaProcessor = LlavaProcessor.from_pretrained(
            model_name,
            cache_dir=self.model_dir,
        )
        model: LlavaForConditionalGeneration = LlavaForConditionalGeneration.from_pretrained(
            model_name,
            cache_dir=self.model_dir,
            torch_dtype=torch.float16,
            device_map="auto",
            low_cpu_mem_usage=True,
            attn_implementation=attn_implementation,
        ).to(self.device).eval()

        return model, processor
 ``` 
Looking into the code, I found that it will go into `_autoset_attn_implementation` three times and check `_check_and_enable_flash_attn_2` three times. It's becouse there are several components of the model llava, the code is shown below. When initing `vision_tower` of llava, it will return a vision_tower with attribute torch_dtype=None, which I suppose lead to the warning, what should I do?
```

class LlavaForConditionalGeneration(LlavaPreTrainedModel, GenerationMixin):
    def __init__(self, config: LlavaConfig):
        super().__init__(config)
        self.vision_tower = AutoModel.from_config(config.vision_config)

        self.multi_modal_projector = LlavaMultiModalProjector(config)
        self.vocab_size = config.text_config.vocab_size
        self.language_model = AutoModelForCausalLM.from_config(
            config.text_config, attn_implementation=config._attn_implementation
        )
        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1
        self.post_init()

```

### Expected behavior

After I specify the torch_dtype in `LlavaForConditionalGeneration.from_pretrained`, it shouldn't warn me You are attempting to use Flash Attention 2.0 without specifying a torch dtype.

Comments:
Should have been fixed by https://github.com/huggingface/transformers/pull/33953, can you try updating transformers version pls?

`!pip install --upgrade git+https://github.com/huggingface/transformers.git`

The general attn implementation setting is going to be modified to support VLMs in https://github.com/huggingface/transformers/pull/32238 and I'll merge it soon

This method seems to be working. Thank you for your help, and I apologize for not seeing the related issues clearly.

Closing as resolved then :)