Spaces:
Sleeping
Sleeping
File size: 4,617 Bytes
3cb2282 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 | """
Load a sentence-transformers model in fp32 and repair what transformers v5 breaks.
Shared by eval_baseline.py, submission.py, finetune.py, precompute_corpus.py and the
Hugging Face Space. space/model_loading.py is a byte-for-byte copy of this file
(tests/test_space_sync.py checks it), so keep it free of repo-internal imports.
"""
from __future__ import annotations
import contextlib
import logging
from typing import Any
import torch
import transformers
from sentence_transformers import SentenceTransformer
logger = logging.getLogger("model_loading")
def restore_nonpersistent_buffers(st_model: SentenceTransformer) -> tuple[int, int]:
"""Recompute non-persistent buffers that transformers v5 leaves uninitialised.
v5 builds the model on the meta device, refills every non-persistent buffer
with torch.empty_like, and relies on `_init_weights` to restore them. Remote-code
models compute such buffers in __init__ (rotary inv_freq / cos / sin caches,
position_ids, attention norm_factor), and their `_init_weights` never touches
them, so they hold garbage after loading. NomicBert (CodeRankEmbed) then
silently gets wrong RoPE frequencies, and Alibaba new-impl models fail with a
CUDA index-out-of-bounds assert (Alibaba-NLP/new-impl discussion #14).
Fix: build a throwaway copy from the config (real construction, weight init
skipped) and copy its buffers over. Returns (restored, differed).
"""
from transformers import PreTrainedModel
hf = next((m for m in st_model.modules() if isinstance(m, PreTrainedModel)), None)
if hf is None or not hasattr(hf, "named_non_persistent_buffers"): # v4 never trashes them
return 0, 0
loaded = dict(hf.named_non_persistent_buffers())
if not loaded:
return 0, 0
try:
from transformers.initialization import no_init_weights
except ImportError:
no_init_weights = contextlib.nullcontext
with torch.device("cpu"), no_init_weights():
fresh = dict(type(hf)(hf.config).named_non_persistent_buffers())
restored = differed = 0
for name, buf in loaded.items():
new = fresh.get(name)
if new is None or new.shape != buf.shape:
logger.warning("Cannot restore buffer %s (missing or shape mismatch in fresh model)", name)
continue
new = new.to(device=buf.device, dtype=buf.dtype)
differed += not torch.equal(new, buf)
parent, _, attr = name.rpartition(".")
hf.get_submodule(parent).register_buffer(attr, new, persistent=False)
restored += 1
return restored, differed
def _is_nomic_bert(model_name: str, trust_remote_code: bool) -> bool:
try:
cfg = transformers.AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code)
except Exception as e: # let SentenceTransformer raise the real loading error
logger.warning("Could not read config of %s: %s", model_name, e)
return False
return cfg.model_type == "nomic_bert"
def load_st_model(model_name: str, device: str, trust_remote_code: bool = True) -> tuple[SentenceTransformer, dict[str, int]]:
"""Load in fp32 and repair non-persistent buffers.
To pin a Hub commit, snapshot_download(repo, revision=...) and pass the local dir:
NomicBert's remote from_pretrained fetches Hub weights without a revision.
"""
# fp32: some checkpoints are stored in fp16/bf16, and transformers v5
# would otherwise load them in that dtype (slow and lossy on CPU).
model_kwargs: dict[str, Any] = {"torch_dtype": torch.float32}
# NomicBert's remote from_pretrained loads Hub ids through
# state_dict_from_pretrained(safe_serialization=kwargs.get("safe_serialization", False)),
# which only looks for pytorch_model.bin(.index.json); our Hub repo (like
# nomic-ai/CodeRankEmbed) has only model.safetensors. Local dirs take a separate
# branch that ignores the flag. Other architectures may reject the kwarg, so gate it.
if _is_nomic_bert(model_name, trust_remote_code):
model_kwargs["safe_serialization"] = True
model = SentenceTransformer(
model_name,
device=device,
trust_remote_code=trust_remote_code,
model_kwargs=model_kwargs,
)
restored, differed = restore_nonpersistent_buffers(model)
logger.info("Restored %d non-persistent buffers (%d had uninitialised values after load)",
restored, differed)
return model, {"restored": restored, "differed": differed}
|