feat: add apple silicon GPU acceleration

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

@@ -104,13 +104,14 @@ def collect_static_system_specs() -> Dict[str, Any]:
     except ImportError:
         specs["libraries.transformers"] = False
 
+    has_mps = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
     try:
         torch_import.check()
         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

haystack/preview/components/readers/extractive.py

@@ -101,6 +101,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")]