diff --git a/src/transformers/modeling_flash_attention_utils.py b/src/transformers/modeling_flash_attention_utils.py index 32642d71d2a3..c5958f77bb23 100644 --- a/src/transformers/modeling_flash_attention_utils.py +++ b/src/transformers/modeling_flash_attention_utils.py @@ -75,8 +75,10 @@ def is_flash_attn_available(): 2: { "flash_attn_version": 2, "general_availability_check": is_flash_attn_2_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None - and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], + "pkg_availability_check": lambda *args, **kwargs: ( + importlib.util.find_spec("flash_attn") is not None + and "flash-attn" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn", [])] + ), "supported_devices": ( (is_torch_cuda_available, "cuda"), (is_torch_mlu_available, "mlu"), @@ -94,16 +96,21 @@ def is_flash_attn_available(): 3: { "flash_attn_version": 3, "general_availability_check": is_flash_attn_3_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn_interface") is not None - and "flash-attn-3" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"]], + "pkg_availability_check": lambda *args, **kwargs: ( + importlib.util.find_spec("flash_attn_interface") is not None + and "flash-attn-3" + in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn_interface", [])] + ), "supported_devices": ((is_torch_cuda_available, "cuda"),), "cuda_min_major_version": 8, # Ampere }, 4: { "flash_attn_version": 4, "general_availability_check": is_flash_attn_4_available, - "pkg_availability_check": lambda *args, **kwargs: importlib.util.find_spec("flash_attn") is not None - and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]], + "pkg_availability_check": lambda *args, **kwargs: ( + importlib.util.find_spec("flash_attn") is not None + and "flash-attn-4" in [pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn", [])] + ), "supported_devices": ((is_torch_cuda_available, "cuda"),), "cuda_min_major_version": 9, # Hopper }, diff --git a/src/transformers/utils/import_utils.py b/src/transformers/utils/import_utils.py index bfd6f13bb12f..71f63313dec9 100644 --- a/src/transformers/utils/import_utils.py +++ b/src/transformers/utils/import_utils.py @@ -985,7 +985,7 @@ def is_flash_attn_2_available() -> bool: is_available, flash_attn_version = _is_package_available("flash_attn", return_version=True) # FA4 is also distributed under "flash_attn", hence we need to check the naming here is_available = is_available and "flash-attn" in [ - pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"] + pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn", []) ] if not is_available or not (is_torch_cuda_available() or is_torch_mlu_available()): @@ -1004,7 +1004,7 @@ def is_flash_attn_3_available() -> bool: is_available = _is_package_available("flash_attn_interface")[0] # Resolving and ensuring the proper name of FA3 being associated is_available = is_available and "flash-attn-3" in [ - pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn_interface"] + pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn_interface", []) ] return is_available and is_torch_cuda_available() @@ -1016,7 +1016,7 @@ def is_flash_attn_4_available() -> bool: # NOTE: FA2 seems to distribute the `cute` subdirectory even if only FA2 has been installed # -> check for the proper (normalized) distribution name is_available = is_available and "flash-attn-4" in [ - pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"] + pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn", []) ] return is_available and is_torch_cuda_available() @@ -1027,7 +1027,7 @@ def is_flash_attn_greater_or_equal(library_version: str) -> bool: is_available, flash_attn_version = _is_package_available("flash_attn", return_version=True) # FA4 is also distributed under "flash_attn", hence we need to check the naming here is_available = is_available and "flash-attn" in [ - pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"] + pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING.get("flash_attn", []) ] if not is_available: diff --git a/tests/utils/test_modeling_utils.py b/tests/utils/test_modeling_utils.py index dd3f2b74d81e..5c26f44f7736 100644 --- a/tests/utils/test_modeling_utils.py +++ b/tests/utils/test_modeling_utils.py @@ -86,6 +86,7 @@ WEIGHTS_NAME, ) from transformers.utils.import_utils import ( + PACKAGE_DISTRIBUTION_MAPPING, is_flash_attn_2_available, is_flash_attn_3_available, is_flash_attn_4_available, @@ -2893,6 +2894,20 @@ def test_not_available_flash(self): ) self.assertTrue("the package for FlashAttention2 doesn't seem to be installed." in str(cm.exception)) + def test_flash_attn_available_no_keyerror_when_missing_from_distribution_map(self): + # Regression test for https://github.com/huggingface/transformers/issues/45520. + # When flash_attn is importable but not present in PACKAGE_DISTRIBUTION_MAPPING + # (e.g. installed via a non-standard wheel), the availability checks must not raise + # a KeyError; they should simply return False. + stripped_map = { + k: v for k, v in PACKAGE_DISTRIBUTION_MAPPING.items() if k not in ("flash_attn", "flash_attn_interface") + } + with patch("transformers.utils.import_utils.PACKAGE_DISTRIBUTION_MAPPING", stripped_map): + with patch("transformers.modeling_flash_attention_utils.PACKAGE_DISTRIBUTION_MAPPING", stripped_map): + self.assertFalse(is_flash_attn_2_available()) + self.assertFalse(is_flash_attn_3_available()) + self.assertFalse(is_flash_attn_4_available()) + def test_not_available_flash_with_config(self): if is_flash_attn_2_available(): self.skipTest(reason="Please uninstall flash-attn package to run test_not_available_flash")