From ddaa5f5f1b1259acd0add5c0af7403e16c057062 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= Date: Thu, 22 Jan 2026 16:25:28 +0100 Subject: [PATCH] fix: simplify pytorch model initialization to prevent meta tensor issues (#185) Remove device_map from model_kwargs as it conflicts with CrossEncoder's internal .to(device) call. The low_cpu_mem_usage=False setting alone is sufficient to prevent lazy loading (meta tensors). --- .../hindsight_api/engine/cross_encoder.py | 23 +++++-------------- .../hindsight_api/engine/embeddings.py | 21 ++++------------- 2 files changed, 10 insertions(+), 34 deletions(-) diff --git a/hindsight-api/hindsight_api/engine/cross_encoder.py b/hindsight-api/hindsight_api/engine/cross_encoder.py index d9507787..51957180 100644 --- a/hindsight-api/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api/hindsight_api/engine/cross_encoder.py @@ -132,36 +132,25 @@ class LocalSTCrossEncoder(CrossEncoderModel): logger.info(f"Reranker: initializing local provider with model {self.model_name}") - # Determine device and device_map based on hardware and installed packages. - # When accelerate is installed but no GPU/MPS is available, transformers can - # incorrectly use lazy loading (meta tensors) which fails on .to(device). - # We use device_map="cpu" in that case to force direct CPU loading. + # Determine device based on hardware availability. + # We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors) + # which can cause issues when accelerate is installed but no GPU is available. + # Note: We do NOT use device_map because CrossEncoder internally calls .to(device) + # after loading, which conflicts with accelerate's device_map handling. import torch - try: - import accelerate # type: ignore[import-not-found] # noqa: F401 - - accelerate_available = True - except ImportError: - accelerate_available = False - # Check for GPU (CUDA) or Apple Silicon (MPS) has_gpu = torch.cuda.is_available() or (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()) if has_gpu: device = None # Let sentence-transformers auto-detect GPU/MPS - device_map = None - elif accelerate_available: - device = "cpu" - device_map = "cpu" # Force direct CPU loading to avoid meta tensors else: device = "cpu" - device_map = None self._model = CrossEncoder( self.model_name, device=device, - model_kwargs={"low_cpu_mem_usage": False, "device_map": device_map}, + model_kwargs={"low_cpu_mem_usage": False}, ) # Initialize shared executor (limited workers naturally limits concurrency) diff --git a/hindsight-api/hindsight_api/engine/embeddings.py b/hindsight-api/hindsight_api/engine/embeddings.py index 5402be7c..3ab3aefc 100644 --- a/hindsight-api/hindsight_api/engine/embeddings.py +++ b/hindsight-api/hindsight_api/engine/embeddings.py @@ -129,36 +129,23 @@ class LocalSTEmbeddings(Embeddings): logger.info(f"Embeddings: initializing local provider with model {self.model_name}") - # Determine device and device_map based on hardware and installed packages. - # When accelerate is installed but no GPU/MPS is available, transformers can - # incorrectly use lazy loading (meta tensors) which fails on .to(device). - # We use device_map="cpu" in that case to force direct CPU loading. + # Determine device based on hardware availability. + # We always set low_cpu_mem_usage=False to prevent lazy loading (meta tensors) + # which can cause issues when accelerate is installed but no GPU is available. import torch - try: - import accelerate # type: ignore[import-not-found] # noqa: F401 - - accelerate_available = True - except ImportError: - accelerate_available = False - # Check for GPU (CUDA) or Apple Silicon (MPS) has_gpu = torch.cuda.is_available() or (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()) if has_gpu: device = None # Let sentence-transformers auto-detect GPU/MPS - device_map = None - elif accelerate_available: - device = "cpu" - device_map = "cpu" # Force direct CPU loading to avoid meta tensors else: device = "cpu" - device_map = None self._model = SentenceTransformer( self.model_name, device=device, - model_kwargs={"low_cpu_mem_usage": False, "device_map": device_map}, + model_kwargs={"low_cpu_mem_usage": False}, ) self._dimension = self._model.get_sentence_embedding_dimension()