embeddings(bge-m3): fix loader crash + head dtype/grad

- import build_from_pretrained_kwargs inside _load_bge_m3 (it's a local import
  of _load_embedding_model, not module-level) — the NameError made every bge-m3
  load fall over so the model silently degraded to the dense-only path.
- cast the sparse/colbert heads to the base model's dtype (bf16) so the head
  matmul over the last hidden state doesn't dtype-mismatch.
- requires_grad_(False) on the heads (inference only) — no autograd graph per
  request.
- log the bge-m3 backend selection at load.
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Mw2KQiswmD69T45fTfjKwW
parent 66f5540b
......@@ -337,6 +337,7 @@ def _load_bge_m3(model_name: str, device, model_config):
(already in the stack) reproduce dense/sparse/colbert exactly."""
import torch
from transformers import AutoTokenizer, AutoModel
from codai.models.hf_loading import build_from_pretrained_kwargs
fp = build_from_pretrained_kwargs(model_config)
tokenizer = AutoTokenizer.from_pretrained(model_name)
......@@ -353,7 +354,11 @@ def _load_bge_m3(model_name: str, device, model_config):
w = sd['weight'] # (out_features, in_features)
lin = torch.nn.Linear(w.shape[1], w.shape[0], bias=('bias' in sd))
lin.load_state_dict(sd)
return lin.to(device).eval()
# Match the base model's dtype (e.g. bf16) so the head matmul over the last
# hidden state doesn't dtype-mismatch; no grad (inference only).
lin = lin.to(device=device, dtype=model.dtype).eval()
lin.requires_grad_(False)
return lin
heads = {}
sp = _load_head('sparse_linear.pt')
......@@ -867,6 +872,7 @@ def _load_embedding_model(model_name: str, device: str, model_config: dict = Non
# sparse_linear head; must come BEFORE the sentence-transformers branch (which would
# otherwise load it dense-only and drop the sparse/colbert heads).
if _is_bge_m3(model_name, trust):
print(f"[embeddings] loading '{model_name}' via bge-m3 multi-vector backend", flush=True)
try:
return _load_bge_m3(model_name, device, model_config)
except Exception as e:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment