Add inference settings for model loading: prioritize CUDA, MPS, and CPU. Update runner and pipeline to utilize new settings for dtype and device placement.

tensegrity/bench/runner.py

@@ -33,6 +33,7 @@ from dataclasses import dataclass, field, asdict
 from pathlib import Path
 
 from tensegrity.bench.tasks import TaskSample, TaskConfig, TASK_REGISTRY, load_task_samples
+from tensegrity.torch_device import inference_load_settings
 
 logger = logging.getLogger(__name__)
 
@@ -164,17 +165,19 @@ class EvalRunner:
             return
 
         from transformers import AutoTokenizer, AutoModelForCausalLM
-        import torch
 
+        dtype, device_map, move_to = inference_load_settings()
         logger.info(f"Loading model {self.model_name}...")
         self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         if self._tokenizer.pad_token is None:
             self._tokenizer.pad_token = self._tokenizer.eos_token
         self._model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto" if torch.cuda.is_available() else None,
+            torch_dtype=dtype,
+            device_map=device_map,
         )
+        if move_to is not None:
+            self._model = self._model.to(move_to)
         self._model.eval()
         logger.info("Model loaded.")
 

tensegrity/graft/pipeline.py

@@ -30,6 +30,7 @@ from tensegrity.graft.logit_bias import (
     StaticLogitBiasBuilder,
     GraftState,
 )
+from tensegrity.torch_device import inference_load_settings
 
 logger = logging.getLogger(__name__)
 
@@ -96,15 +97,17 @@ class HybridPipeline:
             return
         
         from transformers import AutoTokenizer, AutoModelForCausalLM
-        import torch
-        
+
+        dtype, device_map, move_to = inference_load_settings()
         logger.info(f"Loading {self.model_name}...")
         self._tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         self._model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            device_map="auto" if torch.cuda.is_available() else None,
+            torch_dtype=dtype,
+            device_map=device_map,
         )
+        if move_to is not None:
+            self._model = self._model.to(move_to)
         
         # Build vocabulary grounding
         if self._hypothesis_keywords:

tensegrity/torch_device.py

@@ -0,0 +1,27 @@
+"""
+Pick inference dtype and placement for transformers models.
+
+Preference order: CUDA (device_map auto) → Apple MPS → CPU.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Tuple
+
+
+def inference_load_settings() -> Tuple[Any, Optional[str], Optional[Any]]:
+    """
+    Returns (torch_dtype, device_map, move_to_device).
+
+    - CUDA: float16, device_map=\"auto\", move_to_device=None
+    - MPS: float16, device_map=None, move_to_device=torch.device(\"mps\")
+    - CPU: float32, device_map=None, move_to_device=None
+    """
+    import torch
+
+    if torch.cuda.is_available():
+        return torch.float16, "auto", None
+    mps = getattr(torch.backends, "mps", None)
+    if mps is not None and mps.is_available():
+        return torch.float16, None, torch.device("mps")
+    return torch.float32, None, None