Flash attention build running forever on colab
Author: BoccheseGiacomo
Pull request: no
Labels: name: bug

### BUG DESCRIPTION

Running on google colab a script to finetune LLAMA 3 8B with flash attention.
This issue is not directly related to transformers but to an extension library: flash attention

**During the installation of the last package "flash-attn" i get the following line in the console running forever:**
`Building wheels for collected packages: flash-attn`
**The issue was not present before october 15 2024 and this installation worked fine**

### System Info
Running on google colab a script to finetune LLAMA 3 8B with flash attention.

Setup of packages:
```
!pip install -U transformers
!pip install -U datasets
!pip install -U accelerate
!pip install -U peft
!pip install -U trl
!pip install -U bitsandbytes
!pip install -U wandb
!pip install -U flash-attn
```

### Who can help?

_No response_

### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

1. go to google colab, and set A100 gpu
2. Setup the following code for downloading packages
```
!pip install -U transformers
!pip install -U datasets
!pip install -U accelerate
!pip install -U peft
!pip install -U trl
!pip install -U bitsandbytes
!pip install -U wandb
!pip install -U flash-attn
```

4. wait

The issue was not present before october 15 2024 and this installation worked fine

### Expected behavior

Building wheels should terminate in 1-2 minutes max, instead it never ends, tried also to wait for 30 minutes.

Comments:
I confirm but it's probably not related to transformers.

FlashAttention takes many hours to build with PyTorch 2.5.
If you downgrade to 2.4, it should work:
!pip install torch=='2.4.1+cu121' torchvision=='0.19.1+cu121' torchaudio=='2.4.1+cu121'  --index-url https://download.pytorch.org/whl/cu121 


> I confirm but it's probably not related to transformers.
> 
> FlashAttention takes many hours to build with PyTorch 2.5. If you downgrade to 2.4, it should work: !pip install torch=='2.4.1+cu121' torchvision=='0.19.1+cu121' torchaudio=='2.4.1+cu121' --index-url https://download.pytorch.org/whl/cu121

Thanks a lot , this has resolved the issue

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.

> FlashAttention takes many hours to build with PyTorch 2.5. If you downgrade to 2.4, it should work:

is there a solution for those of us who need it to run with `torch==2.6` ?

> > FlashAttention takes many hours to build with PyTorch 2.5. If you downgrade to 2.4, it should work:
> 
> is there a solution for those of us who need it to run with `torch==2.6` ?

same question here, tried many other possible solutions but didn't succeed with any of them

edit: downgrading flash-attn==2.7.3 worked for me