feat: add apple silicon GPU acceleration

.github/workflows/tests_preview.yml

@@ -220,6 +220,8 @@ jobs:
     name: Integration / macos-latest
     needs: unit-tests
     runs-on: macos-latest
+    env:
+      PYTORCH_ENABLE_MPS_FALLBACK: 1
     steps:
       - uses: actions/checkout@v4
 

e2e/modeling/test_dpr.py

@@ -707,6 +707,8 @@ def test_dpr_processor_save_load_non_bert_tokenizer(tmp_path: Path, query_and_pa
 
     if torch.cuda.is_available():
         device = torch.device("cuda")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = torch.device("mps")
     else:
         device = torch.device("cpu")
     model = BiAdaptiveModel(
@@ -753,6 +755,8 @@ def test_dpr_processor_save_load_non_bert_tokenizer(tmp_path: Path, query_and_pa
 
     if torch.cuda.is_available():
         device = torch.device("cuda")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = torch.device("mps")
     else:
         device = torch.device("cpu")
     loaded_model = BiAdaptiveModel(
@@ -879,6 +883,8 @@ def test_dpr_processor_save_load_non_bert_tokenizer(tmp_path: Path, query_and_pa
 
     if torch.cuda.is_available():
         device = torch.device("cuda")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        device = torch.device("mps")
     else:
         device = torch.device("cpu")
     model = BiAdaptiveModel(

haystack/environment.py

@@ -106,11 +106,12 @@ def collect_static_system_specs() -> Dict[str, Any]:
 
     try:
         torch_import.check()
+        has_mps = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
         specs.update(
             {
                 "libraries.torch": torch.__version__,
                 "libraries.cuda": torch.version.cuda if torch.cuda.is_available() else False,
-                "hardware.gpus": torch.cuda.device_count() if torch.cuda.is_available() else 0,
+                "hardware.gpus": torch.cuda.device_count() if torch.cuda.is_available() else 1 if has_mps else 0,
             }
         )
     except ImportError:

haystack/modeling/utils.py

@@ -112,6 +112,9 @@ def initialize_device_settings(
             else:
                 devices_to_use = [torch.device("cuda:0")]
                 n_gpu = 1
+        elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            devices_to_use = [torch.device("mps")]
+            n_gpu = 1
         else:
             devices_to_use = [torch.device("cpu")]
             n_gpu = 0
@@ -180,6 +183,7 @@ def all_gather_list(data, group=None, max_size=16384):
         data (Any): data from the local worker to be gathered on other workers
         group (optional): group of the collective
     """
+    # pylint: disable=all
     SIZE_STORAGE_BYTES = 4  # int32 to encode the payload size
 
     enc = pickle.dumps(data)

haystack/preview/components/readers/extractive.py

@@ -111,6 +111,8 @@ def warm_up(self):
         if self.model is None:
             if torch.cuda.is_available():
                 self.device = self.device or "cuda:0"
+            elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                self.device = self.device or "mps:0"
             else:
                 self.device = self.device or "cpu:0"
             self.model = AutoModelForQuestionAnswering.from_pretrained(self.model_name_or_path, token=self.token).to(

haystack/utils/experiment_tracking.py

@@ -236,6 +236,7 @@ def get_or_create_env_meta_data() -> Dict[str, Any]:
     from haystack.telemetry import HAYSTACK_EXECUTION_CONTEXT
 
     global env_meta_data  # pylint: disable=global-statement
+    has_mps = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
     if not env_meta_data:
         env_meta_data = {
             "os_version": platform.release(),
@@ -246,7 +247,7 @@ def get_or_create_env_meta_data() -> Dict[str, Any]:
             "transformers_version": transformers.__version__,
             "torch_version": torch.__version__,
             "torch_cuda_version": torch.version.cuda if torch.cuda.is_available() else 0,
-            "n_gpu": torch.cuda.device_count() if torch.cuda.is_available() else 0,
+            "n_gpu": torch.cuda.device_count() if torch.cuda.is_available() else 1 if has_mps else 0,
             "n_cpu": os.cpu_count(),
             "context": os.environ.get(HAYSTACK_EXECUTION_CONTEXT),
             "execution_env": _get_execution_environment(),

haystack/utils/torch_utils.py

@@ -44,4 +44,6 @@ def get_devices(devices: Optional[List[Union[str, torch.device]]]) -> List[torch
         return [torch.device(device) for device in devices]
     elif torch.cuda.is_available():
         return [torch.device(device) for device in range(torch.cuda.device_count())]
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return [torch.device("mps")]
     return [torch.device("cpu")]

releasenotes/notes/add-apple-silicon-gpu-acceleration-38bf69781a933b95.yaml

@@ -0,0 +1,4 @@
+---
+enhancements:
+  - |
+    Added support for Apple Silicon GPU acceleration through "mps pytorch", enabling better performance on Apple M1 hardware.