Speed up import transformers by building the package distribution mapping lazily (#50047)

Build the package distribution mapping lazily

importlib.metadata.packages_distributions() stats every file of each
distribution without a top_level.txt (torch >= 2.14 ships none). Build it on
first use, and try the import name as distribution name before it.

Author: qgallouedec

src/transformers/modeling_flash_attention_utils.py

@@ -34,7 +34,7 @@
     is_torch_xpu_available,
     logging,
 )
-from .utils.import_utils import PACKAGE_DISTRIBUTION_MAPPING, is_tracing
+from .utils.import_utils import _package_distribution_mapping, is_tracing
 
 
 logger = logging.get_logger(__name__)
@@ -103,7 +103,8 @@ def is_flash_attn_available():
         "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.get("flash_attn", [])]
+            and "flash-attn"
+            in [pkg.replace("_", "-") for pkg in _package_distribution_mapping().get("flash_attn", [])]
         ),
         "supported_devices": (
             (is_torch_cuda_available, "cuda"),
@@ -126,7 +127,7 @@ def is_flash_attn_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.get("flash_attn_interface", [])]
+            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
@@ -136,7 +137,8 @@ def is_flash_attn_available():
         "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.get("flash_attn", [])]
+            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

src/transformers/utils/import_utils.py

@@ -49,7 +49,9 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
-PACKAGE_DISTRIBUTION_MAPPING = importlib.metadata.packages_distributions()
+@functools.cache
+def _package_distribution_mapping() -> dict[str, list[str]]:
+    return importlib.metadata.packages_distributions()
 
 
 def _candidate_distribution_names(pkg_name: str) -> list[str]:
@@ -61,7 +63,7 @@ def _candidate_distribution_names(pkg_name: str) -> list[str]:
     """
     # Per PEP 503, underscores and hyphens are equivalent in package names.
     normalized_pkg_name = pkg_name.replace("_", "-")
-    distributions = PACKAGE_DISTRIBUTION_MAPPING.get(pkg_name, [])
+    distributions = _package_distribution_mapping().get(pkg_name, [])
     candidates = [
         *(name for name in (normalized_pkg_name, pkg_name) if name in distributions),
         *distributions,
@@ -77,21 +79,26 @@ def _is_package_available(pkg_name: str, return_version: bool = False) -> tuple[
     package_exists = spec is not None
     package_version = "N/A"
     if package_exists and return_version:
-        for distribution_name in _candidate_distribution_names(pkg_name):
-            try:
-                package_version = importlib.metadata.version(distribution_name)
-                break
-            except importlib.metadata.PackageNotFoundError:
-                continue
-        else:
-            # No metadata under any candidate name (editable install without a `dist-info`, for example).
-            # Last resort: importing defeats the lazy imports these checks guard, costing every
-            # `import transformers` the package's whole import tree.
-            package = importlib.import_module(pkg_name)
-            package_version = getattr(package, "__version__", "N/A")
-            # No version + no __file__ means a namespace package (PEP 420) shadowing on sys.path, not a real install.
-            if package_version == "N/A" and getattr(package, "__file__", None) is None:
-                package_exists = False
+        try:
+            # Most packages are distributed under their import name. Try it before building the mapping, which
+            # stats every file of each distribution without a `top_level.txt`, `torch` >= 2.14 among them
+            package_version = importlib.metadata.version(pkg_name)
+        except importlib.metadata.PackageNotFoundError:
+            for distribution_name in _candidate_distribution_names(pkg_name):
+                try:
+                    package_version = importlib.metadata.version(distribution_name)
+                    break
+                except importlib.metadata.PackageNotFoundError:
+                    continue
+            else:
+                # No metadata under any candidate name (editable install without a `dist-info`, for example).
+                # Last resort: importing defeats the lazy imports these checks guard, costing every
+                # `import transformers` the package's whole import tree.
+                package = importlib.import_module(pkg_name)
+                package_version = getattr(package, "__version__", "N/A")
+                # No version + no __file__ means a namespace package (PEP 420) shadowing on sys.path, not a real install.
+                if package_version == "N/A" and getattr(package, "__file__", None) is None:
+                    package_exists = False
         logger.debug(f"Detected {pkg_name} version: {package_version}")
 
     if return_version:
@@ -1227,7 +1234,7 @@ def is_flash_attn_2_available(kernels_fallback_ok: bool = False) -> 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.get("flash_attn", [])
+        pkg.replace("_", "-") for pkg in _package_distribution_mapping().get("flash_attn", [])
     ]
 
     # Only allow versions >= 2.3.3 to avoid very old legacy workarounds that are now 2+ years old
@@ -1260,7 +1267,7 @@ def is_flash_attn_3_available(kernels_fallback_ok: bool = False) -> 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.get("flash_attn_interface", [])
+        pkg.replace("_", "-") for pkg in _package_distribution_mapping().get("flash_attn_interface", [])
     ]
     if is_available and is_torch_cuda_available():
         return True
@@ -1289,7 +1296,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.get("flash_attn", [])
+        pkg.replace("_", "-") for pkg in _package_distribution_mapping().get("flash_attn", [])
     ]
 
     return is_available and is_torch_cuda_available()
@@ -1300,7 +1307,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.get("flash_attn", [])
+        pkg.replace("_", "-") for pkg in _package_distribution_mapping().get("flash_attn", [])
     ]
 
     if not is_available:

tests/utils/test_import_utils.py

@@ -105,17 +105,17 @@ def test_is_package_available_unmapped_distribution_does_not_import():
 
     with (
         patch("transformers.utils.import_utils.importlib.util.find_spec", return_value=object()),
-        patch("transformers.utils.import_utils.PACKAGE_DISTRIBUTION_MAPPING", {}),
+        patch("transformers.utils.import_utils._package_distribution_mapping", return_value={}),
         patch("transformers.utils.import_utils.importlib.metadata.version", return_value="1.2.3") as version,
         patch("transformers.utils.import_utils.importlib.import_module") as import_module,
     ):
         assert _is_package_available(pkg_name, return_version=True) == (True, "1.2.3")
-        version.assert_called_once_with(pkg_name.replace("_", "-"))
+        version.assert_called_once_with(pkg_name)
         import_module.assert_not_called()
 
 
 def test_candidate_distribution_names():
-    with patch("transformers.utils.import_utils.PACKAGE_DISTRIBUTION_MAPPING", {"PIL": ["pillow"]}):
+    with patch("transformers.utils.import_utils._package_distribution_mapping", return_value={"PIL": ["pillow"]}):
         # A mapped import name prefers its distribution, but keeps the import name as a fallback.
         assert _candidate_distribution_names("PIL") == ["pillow", "PIL"]
         # An unmapped import name falls back on itself, normalized first, without duplicates.
@@ -160,7 +160,9 @@ def fake_is_package_available(pkg_name: str, return_version: bool = False) -> tu
     try:
         with (
             patch("transformers.utils.import_utils._is_package_available", side_effect=fake_is_package_available),
-            patch("transformers.utils.import_utils.PACKAGE_DISTRIBUTION_MAPPING", fake_distribution_mapping),
+            patch(
+                "transformers.utils.import_utils._package_distribution_mapping", return_value=fake_distribution_mapping
+            ),
             patch("transformers.utils.import_utils.is_torch_cuda_available", return_value=cuda_available),
             patch("transformers.utils.import_utils.is_torch_mlu_available", return_value=False),
             patch("transformers.utils.import_utils.is_torch_musa_available", return_value=False),

tests/utils/test_modeling_utils.py

@@ -94,7 +94,7 @@
     WEIGHTS_NAME,
 )
 from transformers.utils.import_utils import (
-    PACKAGE_DISTRIBUTION_MAPPING,
+    _package_distribution_mapping,
     is_flash_attn_2_available,
     is_flash_attn_3_available,
     is_flash_attn_4_available,
@@ -3328,14 +3328,16 @@ def test_transformers_sdpa_force_gqa_env_var_overrides_use_gqa_in_sdpa_attention
 
     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
+        # When flash_attn is importable but not present in the 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")
+            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):
+        with patch("transformers.utils.import_utils._package_distribution_mapping", return_value=stripped_map):
+            with patch(
+                "transformers.modeling_flash_attention_utils._package_distribution_mapping", return_value=stripped_map
+            ):
                 self.assertFalse(is_flash_attn_2_available())
                 self.assertFalse(is_flash_attn_3_available())
                 self.assertFalse(is_flash_attn_4_available())