diff --git a/pyproject.toml b/pyproject.toml index 4364c2e33a..ea420e2ffb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -97,6 +97,15 @@ colab_ampere = [ "ninja", "flash-attn", ] +colab_torch211 = [ + "unsloth[cu121_torch211]", +] +colab_ampere_torch211 = [ + "unsloth[cu121_torch211]", + "packaging", + "ninja", + "flash-attn", +] cu118_ampere = [ "unsloth[huggingface]", "bitsandbytes", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 8ddf4d60da..5f35851443 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -30,7 +30,19 @@ major_version, minor_version = torch.cuda.get_device_capability() if major_version >= 8: try: from flash_attn import flash_attn_func - HAS_FLASH_ATTENTION = True + # Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl" + try: + from flash_attn.flash_attn_interface import flash_attn_cuda + HAS_FLASH_ATTENTION = True + except: + logger.warning_once( + "Unsloth: Your Flash Attention 2 installation seems to be broken?\n"\ + "A possible explanation is you have a new CUDA version which isn't\n"\ + "yet compatible with FA2? Please file a ticket to Unsloth or FA2.\n"\ + "We shall now use Xformers instead, which gets a 0.01% performance hit.\n"\ + "We found this negligible impact by benchmarking on 1x A100." + ) + HAS_FLASH_ATTENTION = False except: HAS_FLASH_ATTENTION = False else: