holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

fix: improve pytorch model initialization to prevent meta tensor issues (#180)

Browse source 
* fix: prevent meta tensor issues when accelerate is installed without GPU

When accelerate is installed but no GPU is available, transformers can
incorrectly use lazy loading (meta tensors) which fails when
sentence-transformers tries to move the model to a device.

The fix checks hardware and installed packages to determine the right
loading strategy:
- GPU available: device=None, device_map=None (auto-detect GPU)
- No GPU + accelerate: device='cpu', device_map='cpu' (force CPU loading)
- No GPU + no accelerate: device='cpu', device_map=None (normal CPU)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

* fix: add filelock for model initialization in parallel tests

When pytest-xdist runs multiple workers in parallel, they all try to
load models from the HuggingFace cache simultaneously, causing race
conditions and intermittent meta tensor errors.

Added filelock around embeddings and cross_encoder initialization in
conftest.py, similar to how pg0 database setup is serialized. Models
are now pre-initialized in the fixture before being passed to tests.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

* fix: add MPS support for macOS Apple Silicon

Extend GPU detection to include Apple MPS backend in addition to CUDA.
This ensures macOS users with Apple Silicon use MPS acceleration
instead of being incorrectly routed to the CPU fallback path.
This commit is contained in:
Nicolò Boschi 2026-01-20 14:15:51 +01:00 committed by GitHub
parent 16b85a4faa
commit 0bf85a3435
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 111 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
hindsight-api/hindsight_api/engine/cross_encoder.py
View file

@ -130,17 +130,38 @@ class LocalSTCrossEncoder(CrossEncoderModel):
"Install it with: pip install sentence-transformers"
)
# Note: We use CPU even when GPU/MPS is available because:
# 1. The reranker model (MiniLM) is tiny (~22M params)
# 2. Batch sizes are small (~100-200 pairs)
# 3. Data transfer overhead to GPU outweighs compute benefit
# 4. CPU inference is actually faster for this workload
logger.info(f"Reranker: initializing local provider with model {self.model_name}")
# Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate.
# Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized.
# 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.
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,
model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
device=device,
model_kwargs={"low_cpu_mem_usage": False, "device_map": device_map},
)
# Initialize shared executor (limited workers naturally limits concurrency)

32
hindsight-api/hindsight_api/engine/embeddings.py
View file

@ -128,11 +128,37 @@ class LocalSTEmbeddings(Embeddings):
)
logger.info(f"Embeddings: initializing local provider with model {self.model_name}")
# Disable lazy loading (meta tensors) which causes issues with newer transformers/accelerate
# Setting low_cpu_mem_usage=False and device_map=None ensures tensors are fully materialized
# 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.
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,
model_kwargs={"low_cpu_mem_usage": False, "device_map": None},
device=device,
model_kwargs={"low_cpu_mem_usage": False, "device_map": device_map},
)
self._dimension = self._model.get_sentence_embedding_dimension()

57
hindsight-api/tests/conftest.py
View file

@ -116,16 +116,65 @@ def llm_config():
@pytest.fixture(scope="session")
def embeddings():
def embeddings(tmp_path_factory, worker_id):
"""
Session-scoped embeddings fixture with filelock to prevent race conditions.
return LocalSTEmbeddings()
When pytest-xdist runs multiple workers in parallel, they all try to load
models from the HuggingFace cache simultaneously, which can cause race
conditions and meta tensor errors. We use a filelock to serialize model
initialization across workers.
"""
# Get shared temp dir for coordination between xdist workers
if worker_id == "master":
root_tmp_dir = tmp_path_factory.getbasetemp()
else:
root_tmp_dir = tmp_path_factory.getbasetemp().parent
lock_file = root_tmp_dir / "embeddings_init.lock"
emb = LocalSTEmbeddings()
# Serialize model initialization across workers
with filelock.FileLock(str(lock_file)):
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(emb.initialize())
finally:
loop.close()
return emb
@pytest.fixture(scope="session")
def cross_encoder():
def cross_encoder(tmp_path_factory, worker_id):
"""
Session-scoped cross-encoder fixture with filelock to prevent race conditions.
return LocalSTCrossEncoder()
When pytest-xdist runs multiple workers in parallel, they all try to load
models from the HuggingFace cache simultaneously, which can cause race
conditions and meta tensor errors. We use a filelock to serialize model
initialization across workers.
"""
# Get shared temp dir for coordination between xdist workers
if worker_id == "master":
root_tmp_dir = tmp_path_factory.getbasetemp()
else:
root_tmp_dir = tmp_path_factory.getbasetemp().parent
lock_file = root_tmp_dir / "cross_encoder_init.lock"
ce = LocalSTCrossEncoder()
# Serialize model initialization across workers
with filelock.FileLock(str(lock_file)):
loop = asyncio.new_event_loop()
try:
loop.run_until_complete(ce.initialize())
finally:
loop.close()
return ce
@pytest.fixture(scope="session")
def query_analyzer():