I got a message about Flash Attention 2 when I using axolotl full fine tuning mixtral7B x 8
Author: DopeorNope-Lee
Pull request: no

### System Info

```python3
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/opt/conda/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/jovyan/fileviewer/LLM/axolotl/src/axolotl/cli/train.py", line 38, in <module>
    fire.Fire(do_cli)
  File "/home/jovyan/.local/lib/python3.10/site-packages/fire/core.py", line 141, in Fire
    component_trace = _Fire(component, args, parsed_flag_args, context, name)
  File "/home/jovyan/.local/lib/python3.10/site-packages/fire/core.py", line 475, in _Fire
    component, remaining_args = _CallAndUpdateTrace(
  File "/home/jovyan/.local/lib/python3.10/site-packages/fire/core.py", line 691, in _CallAndUpdateTrace
    component = fn(*varargs, **kwargs)
  File "/home/jovyan/fileviewer/LLM/axolotl/src/axolotl/cli/train.py", line 34, in do_cli
    train(cfg=parsed_cfg, cli_args=parsed_cli_args, dataset_meta=dataset_meta)
  File "/home/jovyan/fileviewer/LLM/axolotl/src/axolotl/train.py", line 62, in train
    model, peft_config = load_model(cfg, tokenizer, inference=cli_args.inference)
  File "/home/jovyan/fileviewer/LLM/axolotl/src/axolotl/utils/models.py", line 464, in load_model
    raise err
  File "/home/jovyan/fileviewer/LLM/axolotl/src/axolotl/utils/models.py", line 453, in load_model
    model = AutoModelForCausalLM.from_pretrained(
  File "/home/jovyan/.local/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py", line 566, in from_pretrained
    return model_class.from_pretrained(
  File "/home/jovyan/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 3444, in from_pretrained
    config = cls._autoset_attn_implementation(
  File "/home/jovyan/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 1302, in _autoset_attn_implementation
    cls._check_and_enable_flash_attn_2(
  File "/home/jovyan/.local/lib/python3.10/site-packages/transformers/modeling_utils.py", line 1401, in _check_and_enable_flash_attn_2
    raise ImportError(f"{preface} Flash Attention 2 is not available. {install_message}")
ImportError: FlashAttention2 has been toggled on, but it cannot be used due to the following error: Flash Attention 2 is not available. Please refer to the documentation of https://huggingface.co/docs/transformers/perf_infer_gpu_one#flashattention-2 to install Flash Attention 2.
```

I got a message using axolotl for full fine tuning mixtral 7B.

However, I encountered an errormessage.

How can I fix it?

### Who can help?

_No response_

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

https://github.com/OpenAccess-AI-Collective/axolotl
I used this examples for full fine-tuning Mixtral-7B

### Expected behavior

could you help me fix this error when I fine-tune the model??

Comments:
Hi @DopeorNope-Lee, thanks for raising this issue! 

Could you share your running environment: run `transformers-cli env` in the terminal and copy-paste the output? 

Use version `flash_attn==2.3.4`
and on the axo cfg `flash_attention_2: true`

Hi @fblgit, thanks for providing this information. To be able to debug, we'll information about the transformers version and any relevant libraries (obtained by running `transformers-cli env`). 

Could you also provide a minimal reproducer? The link provided in the issue description just goes directly to a repo, rather than a specific script. We get many issues and PRs, so that we can get to all of them in a timely manner, we need you to help us help you. In this case, we'll need the smallest amount of code that can be run to reproduce the error. 

same error.

> Hi @DopeorNope-Lee, thanks for raising this issue!
> 
> Could you share your running environment: run `transformers-cli env` in the terminal and copy-paste the output?

Same error, and I also cannot successfully install the flash_attn2. The transformer-cli env is:
**- `transformers` version: 4.32.0.dev0
- Platform: Linux-5.15.0-1042-azure-x86_64-with-glibc2.17
- Python version: 3.8.13
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- Accelerate version: 0.21.0
- Accelerate config:    not found
- PyTorch version (GPU?): 1.12.1+rocm5.2.3 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>**
Could I know how to fix it? 

Hi @SUNJIMENG, thanks for providing the environment information. Could you also provide a minimal code snippet to reproduce the issue? 

Flash Attention was added for mistral in v4.34 and for mixstral in v4.36. I'd first suggest upgrading to the latest version of transformers `pip install -U transformers` and re-running your code. 

@fblgit from his advice I can fix the error.

@SUNJIMENG could you remove your flash-attn and reinstall with `pip install -U flash-attn`?

Facing same issue, Followed

```
pip install -U transformers
pip install -U flash-attn
```

Still same error

@BakingBrains If you're also experiencing this error, could you provide your environment information - run`transformers-cli env` - and a snippet we can run to reproduce the error? 

@amyeroberts the problem was with the CUDA drivers, I re-installed, now it's fine.
Thank you.

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.

@BakingBrains can you share how you re-installed the CUDA drivers? what commands did you run? I am running into the same error and running the following command did not fix it for me. 

`pip install -U flash-attn` 

I had `os.environ["CUDA_VISIBLE_DEVICES"] = ''`, just ensure torch.cuda.is_available() is also true. The error message is not super helpful

> [@BakingBrains](https://github.com/BakingBrains) can you share how you re-installed the CUDA drivers? what commands did you run? I am running into the same error and running the following command did not fix it for me.
> 
> `pip install -U flash-attn`

I know this is quite old now, but I did this to remove the drivers from my machine (Ubuntu 22.04):

```sh
sudo apt-get --purge remove '*nvidia*'
sudo apt-get --purge remove '*cuda*'
sudo apt-get autoremove
sudo reboot
```

and then to download the drivers go to [this link](https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&Distribution=Ubuntu&target_version=22.04)