"RuntimeError: FlashAttention only supports fp16 and bf16 data type" when training llama-2-7b-hf on databricks-dolly-15k dataset
Author: dshi2
Pull request: no

### System Info

transformers version: 4.34.0.dev0
Platform: Linux-5.15.109+-x86_64-with-glibc2.35
Python version: 3.10.12
Huggingface_hub version: 0.16.4
Safetensors version: 0.3.3
Accelerate version: 0.23.0.dev0
Accelerate config: Distributed environment: NO
Num processes: 1
Process index: 0
Local process index: 0
Device: cuda

Mixed precision type: no

PyTorch version (GPU?): 2.0.1+cu118 (True)
Tensorflow version (GPU?): 2.13.0 (NA)
Flax version (CPU?/GPU?/TPU?): 0.7.2 (NA)
Jax version: 0.4.14
JaxLib version: 0.4.14
Using GPU in script?: Yes, A100
Using distributed or parallel set-up in script?:

### Who can help?

@philschmid @ArthurZucker @sgugger 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

from datasets import load_dataset
from random import randrange

#### Load dataset from the hub
dataset = load_dataset("databricks/databricks-dolly-15k", split="train")

print(f"dataset size: {len(dataset)}")
print(dataset[randrange(len(dataset))])



def format_instruction(sample):
	return f"""### Instruction:
Use the Input below to create an instruction, which could have been used to generate the input using an LLM.

### Input:
{sample['response']}

### Response:
{sample['instruction']}
"""


from random import randrange

print(format_instruction(dataset[randrange(len(dataset))]))

!python -c "import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not supported for Flash Attention'"
!pip install ninja packaging
!MAX_JOBS=4 pip install flash-attn --no-build-isolation


import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

use_flash_attention = True

if torch.cuda.get_device_capability()[0] >= 8:
  from utils.llama_patch import replace_attn_with_flash_attn
  print("Using flash attention")
  replace_attn_with_flash_attn()
  use_flash_attention = True


### Hugging Face model id
model_id = "NousResearch/Llama-2-7b-hf" # non-gated
### model_id = "meta-llama/Llama-2-7b-hf" # gated


### BitsAndBytesConfig int-4 config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

### Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=bnb_config, use_cache=False, device_map="auto")
model.config.pretraining_tp = 1

### Validate that the model is using flash attention, by comparing doc strings
if use_flash_attention:
    from utils.llama_patch import forward
    assert model.model.layers[0].self_attn.forward.__doc__ == forward.__doc__, "Model is not using flash attention"


tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"


from peft import LoraConfig, prepare_model_for_kbit_training, get_peft_model

### LoRA config based on QLoRA paper
peft_config = LoraConfig(
        lora_alpha=16,
        lora_dropout=0.1,
        r=64,
        bias="none",
        task_type="CAUSAL_LM",
)


### prepare model for training
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, peft_config)


from transformers import TrainingArguments

args = TrainingArguments(
    output_dir="llama-7-int4-dolly",
    num_train_epochs=3,
    per_device_train_batch_size=6 if use_flash_attention else 4,
    gradient_accumulation_steps=2,
    gradient_checkpointing=True,
    optim="paged_adamw_32bit",
    logging_steps=10,
    save_strategy="epoch",
    learning_rate=2e-4,
    bf16=True,
    tf32=True,
    max_grad_norm=0.3,
    warmup_ratio=0.03,
    lr_scheduler_type="constant",
    disable_tqdm=True # disable tqdm since with packing values are in correct
)


from trl import SFTTrainer

max_seq_length = 2048 # max sequence length for model and packing of the dataset

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    max_seq_length=max_seq_length,
    tokenizer=tokenizer,
    packing=True,
    formatting_func=format_instruction,
    args=args,
)


from trl import SFTTrainer

max_seq_length = 2048 # max sequence length for model and packing of the dataset

trainer = SFTTrainer(
    model=model,
    train_dataset=dataset,
    peft_config=peft_config,
    max_seq_length=max_seq_length,
    tokenizer=tokenizer,
    packing=True,
    formatting_func=format_instruction,
    args=args,
)


### train
trainer.train() # there will not be a progress bar since tqdm is disabled

### save model
trainer.save_model()


### Expected behavior

Hi,
I was trying to train llama-2-7b-hf (quantized to 4-bit) on databricks/databricks-dolly-15k dataset. Essentially, I was trying to follow this amazing tutorial https://www.philschmid.de/instruction-tune-llama-2. But for some reason, I'm running into an error that states that **RuntimeError: FlashAttention only supports fp16 and bf16 data type**. 

Colab Notebook: https://colab.research.google.com/drive/13UUTyjXmSSkSkBjnLz_2gy4yQHBVCEcg?usp=sharing

This error seems to occur when the trainer.train() method is called. The error backtrace seems to go all the way down to the flash_attn module.

Things I have tried:
1. I updated the package versions --installed from the branch itself.
2. I tried setting tf32=False, but it didn't help.
3. I also tried to run the script without bf16 and tf32 + enabled fp16=True 
4. Lastly, I also changed the optimizer from **paged_adamw_32bit** (which is a 32-bit optimizer) to a normal Adam optimizer.

Unfortunately, none of this worked. What am I doing wrong?

Comments:
cc @younesbelkada 

Hello @DrishtiShrrrma
You are using a monkey-patched version of Llama modeling script so I can't really know what is going on under the hood.
Flash Attention 2 support is being added https://github.com/huggingface/transformers/pull/25598 if you want to already use it please run:
```bash
pip install git+https://github.com/younesbelkada/transformers.git@add-flash-attn-2
```
And make sure to add `use_flash_attn_2=True` when calling `from_pretrained` as stated here: https://github.com/huggingface/transformers/pull/25598#issue-1857032190
 

I see... Thanks for the help, @younesbelkada!