PR #45524: utils: handle flash_attn missing from importlib packages_distributions without crashing

• Repository: huggingface/transformers
• Author: SAY-5
• State: open | merged: False
• Link: https://github.com/huggingface/transformers/pull/45524

Description (problem / solution / changelog)

Fixes #45520.

is_flash_attn_2_available, is_flash_attn_3_available, is_flash_attn_4_available, and is_flash_attn_greater_or_equal all do two checks:

is_available, _ = _is_package_available("flash_attn", return_version=True)
is_available = is_available and "flash-attn" in [
    pkg.replace("_", "-") for pkg in PACKAGE_DISTRIBUTION_MAPPING["flash_attn"]
]

Step 1 uses importlib.util.find_spec, which returns a spec if any flash_attn import is findable — an editable install, a namespace package, a bundled shim, or a stub module sitting under another project. Step 2 then assumes that every findable import name also has an entry in importlib.metadata.packages_distributions().

That assumption does not hold in practice. On Python 3.13 with ComfyUI setups (as reported in #45520), and more generally in any environment where the flash_attn import is resolvable via a non-pip source, packages_distributions() has no "flash_attn" key. Because the list comprehension is evaluated before the in operator, short-circuit evaluation of the outer and does not protect us — the KeyError fires during transformers import and takes down the whole process before any model is loaded.

Swap the four raising subscripts for .get(name, []). If the name is missing from the distribution map we simply conclude that the requested flash-attention flavour is not properly installed — which is the answer is_flash_attn_*_available() would have returned anyway — instead of raising. The inner helper _is_package_available already wraps the same subscript in a try/except, so we are only making the outer call sites match that contract.

Changed files

• src/transformers/utils/import_utils.py (modified, +4/-4)