flash-attn Detection Logic Fails for flash-attn

[Zzhiter]

Problem Description

When using unsloth, I noticed an issue with the flash-attn detection logic. Specifically, unsloth checks for the presence of flash-attn using the following code:

from flash_attn.flash_attn_interface import flash_attn_cuda

However, in latest flash-attn(mine is 2.7.3), the flash_attn_cuda module has been removed and replaced with flash_attn_gpu (see this commit). As a result, unsloth's detection logic fails for flash-attn, incorrectly concluding that flash-attn is not installed or is broken, and falls back to xformers.

Expected Behavior

unsloth should correctly detect the presence of flash-attn regardless of whether the installed version uses flash_attn_cuda or flash_attn_gpu.

Actual Behavior

When flash-attn is installed, unsloth fails to detect it and falls back to xformers, even though flash-attn is properly installed and functional.

Steps to Reproduce

1. Install flash-attn: pip install flash-attn --no-build-isolation
2. Run a training script using unsloth.
3. Observe the logs. unsloth will incorrectly report that flash-attn is not installed or is broken, and will fall back to xformers.

Environment Information

Python Version: 3.10
PyTorch Version: 2.5.1+cu124
CUDA Version: 12.4
flash-attn Version: 2.7.3

[danielhanchen]

Appreciate the fix!

[danielhanchen]

Just pushed it in - hopefully flash-attn works as expected now! For local machines, please update Unsloth via:

pip install --upgrade --no-cache-dir --force-reinstall --no-deps unsloth unsloth_zoo

[ChaJiuZuoban]

This is an awesome issue. It has been of great help to me.