Feature Extraction
sentence-transformers
Safetensors
embedding_gemma2
sentence-similarity
embeddings
text-embeddings
multimodal
image-text-retrieval
audio-embeddings
semantic-search
retrieval
quantized
4-bit precision
int4
mixed-precision
low-memory
gemma
embeddinggemma
multimodal-embeddings
cross-modal-retrieval
clip-alternative
quantization
post-training-quantization
cpu
edge
8-bit precision
| """qpack3: int4 group-wise packing for the EmbeddingGemma 2 text and vision towers. | |
| Scheme: per group of 16 input columns, asymmetric min-max 4-bit codes (two codes per byte, low nibble | |
| first). The per-group scale and minimum are stored as 8-bit affine integers per row with fp16 row | |
| bounds. The audio tower is kept: its block linears are packed the same way. The vision position table is cut to the first 2,520 rows per axis, | |
| which is exact for the default 280 soft tokens (at most 2,520 patches). Do not raise max_soft_tokens. | |
| build_package(model, out_dir, side_dir) quantize a loaded EmbeddingGemma2Model in place and save it | |
| load_packed3(path, dtype) rebuild the module tree on empty weights and attach packed tensors | |
| load_st(path, dtype, device) SentenceTransformer around the packed model | |
| Requires transformers 5.19.0 and sentence-transformers >= 6.1. | |
| """ | |
| import contextlib | |
| import json | |
| import os | |
| import re | |
| import shutil | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from safetensors.torch import load_file, save_file | |
| BITS, G, POS_ROWS = 4, 16, 2520 | |
| LEAVE = {"model.safetensors", "quant_config.json", "qpack3.py"} | |
| def _enc_meta(t): | |
| """(rows, groups) fp16 -> uint8 codes plus per-row fp16 lo and hi.""" | |
| x = t.float() | |
| lo = x.amin(1, keepdim=True).half() | |
| hi = x.amax(1, keepdim=True).half() | |
| span = (hi.float() - lo.float()).clamp(min=1e-6) | |
| q = ((x - lo.float()) / span * 255).round().clamp(0, 255).to(torch.uint8) | |
| return q, lo, hi | |
| def _dec_meta(q, lo, hi, dtype): | |
| span = (hi.float() - lo.float()).clamp(min=1e-6) | |
| return (q.float() / 255 * span + lo.float()).half().to(dtype) | |
| CLIP_GRID = (1.0, 0.97, 0.94, 0.91, 0.88, 0.85, 0.82, 0.79, 0.76) | |
| def _clip_range(Wg, wmin, wmax, qmax): | |
| """Per group, pick the shrink of the min-max range that minimizes squared rounding error.""" | |
| center, half = (wmin + wmax) / 2, (wmax - wmin) / 2 | |
| best_err = best_lo = best_hi = None | |
| for f in CLIP_GRID: | |
| lo, hi = center - half * f, center + half * f | |
| sc = ((hi - lo) / qmax).half().float().clamp(min=1e-4) | |
| z = lo.half().float() | |
| q = ((Wg - z.unsqueeze(-1)) / sc.unsqueeze(-1)).round().clamp(0, qmax) | |
| err = ((q * sc.unsqueeze(-1) + z.unsqueeze(-1) - Wg) ** 2).sum(-1) | |
| if best_err is None: | |
| best_err, best_lo, best_hi = err, lo, hi | |
| else: | |
| better = err < best_err | |
| best_err = torch.where(better, err, best_err) | |
| best_lo, best_hi = torch.where(better, lo, best_lo), torch.where(better, hi, best_hi) | |
| return best_lo, best_hi | |
| def quantize(W, g=G, bits=4, chunk=32768, clip=False): | |
| """W (rows, cols) -> dict of packed uint8 codes and 8-bit metadata, chunked over rows.""" | |
| rows, cols = W.shape | |
| assert cols % g == 0 and cols % 2 == 0, (rows, cols) | |
| qmax = 2 ** bits - 1 | |
| parts = {k: [] for k in ("packed", "sq", "slo", "shi", "zq", "zlo", "zhi")} | |
| for a in range(0, rows, chunk): | |
| Wg = W[a:a + chunk].float().reshape(-1, cols // g, g) | |
| wmin, wmax = Wg.amin(-1), Wg.amax(-1) | |
| if clip: | |
| wmin, wmax = _clip_range(Wg, wmin, wmax, qmax) | |
| scale = ((wmax - wmin) / qmax).half().clamp(min=1e-4) | |
| sq, slo, shi = _enc_meta(scale) | |
| zq, zlo, zhi = _enc_meta(wmin.half()) | |
| s = _dec_meta(sq, slo, shi, torch.float32).unsqueeze(-1) | |
| z = _dec_meta(zq, zlo, zhi, torch.float32).unsqueeze(-1) | |
| q = ((Wg - z) / s).round().clamp(0, qmax).to(torch.uint8).reshape(-1, cols) | |
| parts["packed"].append((q[:, 0::2] | (q[:, 1::2] << 4)).contiguous() if bits == 4 else q.contiguous()) | |
| for k, v in zip(("sq", "slo", "shi", "zq", "zlo", "zhi"), (sq, slo, shi, zq, zlo, zhi)): | |
| parts[k].append(v) | |
| return {k: torch.cat(v) for k, v in parts.items()} | |
| def _empty(m, rows, cols, g, bits=4): | |
| mk = lambda *s, dt: torch.empty(*s, dtype=dt, device="meta") | |
| m.bits = bits | |
| m.register_buffer("packed", mk(rows, cols // 2 if bits == 4 else cols, dt=torch.uint8)) | |
| for k in ("sq", "zq"): | |
| m.register_buffer(k, mk(rows, cols // g, dt=torch.uint8)) | |
| for k in ("slo", "shi", "zlo", "zhi"): | |
| m.register_buffer(k, mk(rows, 1, dt=torch.float16)) | |
| def _deq(m, idx, dtype): | |
| pick = (lambda t: t) if idx is None else (lambda t: t[idx]) | |
| s = _dec_meta(pick(m.sq), pick(m.slo), pick(m.shi), dtype).unsqueeze(-1) | |
| z = _dec_meta(pick(m.zq), pick(m.zlo), pick(m.zhi), dtype).unsqueeze(-1) | |
| p = pick(m.packed) | |
| q = torch.stack([p & 15, p >> 4], -1).reshape(p.shape[0], -1) if m.bits == 4 else p | |
| r = q.shape[0] | |
| return (q.to(dtype).reshape(r, -1, m.g) * s + z).reshape(r, -1) | |
| class QLinear3(nn.Module): | |
| def __init__(self, in_features, out_features, bias, g=G, bits=4): | |
| super().__init__() | |
| self.in_features, self.out_features, self.g = in_features, out_features, g | |
| _empty(self, out_features, in_features, g, bits) | |
| if bias: | |
| self.register_buffer("bias", torch.empty(out_features, dtype=torch.float32, device="meta")) | |
| else: | |
| self.bias = None | |
| def forward(self, x): | |
| W = _deq(self, None, x.dtype) | |
| return F.linear(x, W, None if self.bias is None else self.bias.to(x.dtype)) | |
| class QEmbedding3(nn.Module): | |
| def __init__(self, num_embeddings, embedding_dim, padding_idx, embed_scale, g=G, bits=4): | |
| super().__init__() | |
| self.num_embeddings, self.embedding_dim, self.g = num_embeddings, embedding_dim, g | |
| self.padding_idx, self.embed_scale = padding_idx, float(embed_scale) | |
| self.compute_dtype = torch.float32 | |
| _empty(self, num_embeddings, embedding_dim, g, bits) | |
| def forward(self, ids): | |
| flat = ids.reshape(-1) | |
| W = _deq(self, flat, self.compute_dtype) | |
| return W.reshape(*ids.shape, self.embedding_dim) * self.embed_scale | |
| def bits_for(path, int8_layers, audio_bits=4): | |
| """int8 for the layers listed per tower, audio_bits for audio, int4 elsewhere.""" | |
| if path.startswith("audio_tower."): | |
| return audio_bits | |
| tower = path.split(".")[0] | |
| m = re.search(r"layers\.(\d+)\.", path) | |
| if m: | |
| return 8 if int(m.group(1)) in int8_layers.get(tower, ()) else 4 | |
| return 8 if "other" in int8_layers.get(tower, ()) else 4 | |
| def want_linear(path, _mod): | |
| if path.startswith("language_model."): | |
| return True | |
| if path.startswith("vision_tower."): | |
| return bool(re.search(r"layers\.\d+\.", path)) and _mod.in_features % G == 0 | |
| if path.startswith("audio_tower."): | |
| return _mod.in_features % G == 0 and _mod.in_features * _mod.out_features >= 65536 | |
| return False | |
| def want_embedding(path, mod): | |
| return path.startswith("language_model.") and mod.num_embeddings >= 100_000 | |
| def swap(root, make_lin, make_emb, prefix=""): | |
| for name, child in list(root.named_children()): | |
| path = f"{prefix}{name}" | |
| if isinstance(child, nn.Embedding) and want_embedding(path, child): | |
| setattr(root, name, make_emb(path, child)) | |
| elif isinstance(child, nn.Linear) and want_linear(path, child): | |
| setattr(root, name, make_lin(path, child)) | |
| else: | |
| swap(child, make_lin, make_emb, path + ".") | |
| def build_package(model, out_dir, side_dir, audio_bits=4, int8_layers=None): | |
| int8_layers = int8_layers or {} | |
| """Quantize `model` (EmbeddingGemma2Model) in place, then write the package to out_dir.""" | |
| scales = {} | |
| def make_lin(_p, lin): | |
| b = bits_for(_p, int8_layers, audio_bits) | |
| m = QLinear3(lin.in_features, lin.out_features, lin.bias is not None, bits=b) | |
| for k, v in quantize(lin.weight.data, bits=b).items(): | |
| setattr(m, k, v) | |
| if lin.bias is not None: | |
| m.bias = lin.bias.data.float() | |
| return m | |
| def make_emb(path, emb): | |
| scales[path] = float(getattr(emb, "embed_scale", 1.0)) | |
| m = QEmbedding3(emb.num_embeddings, emb.embedding_dim, emb.padding_idx, scales[path]) | |
| for k, v in quantize(emb.weight.data).items(): | |
| setattr(m, k, v) | |
| return m | |
| swap(model, make_lin, make_emb) | |
| pe = model.vision_tower.patch_embedder | |
| pe.position_embedding_table = nn.Parameter(pe.position_embedding_table.data[:, :POS_ROWS].clone(), | |
| requires_grad=False) | |
| os.makedirs(out_dir, exist_ok=True) | |
| sd = {k: v.detach().cpu().contiguous() for k, v in model.state_dict().items()} | |
| save_file(sd, os.path.join(out_dir, "model.safetensors"), metadata={"format": "pt"}) | |
| json.dump({"format": 3, "bits": BITS, "audio_bits": audio_bits, "int8_layers": int8_layers, "group": G, "position_rows": POS_ROWS, "embed_scale": scales, | |
| "meta": "uint8 affine per row with fp16 lo and hi", "audio": f"tower block linears int{audio_bits}; convolutions, norms and embed_audio bf16", | |
| "text": "linears and token table int4", "vision": "block linears int4", | |
| "kept_bf16": "patch input projection, position table, embed_vision, norms"}, | |
| open(os.path.join(out_dir, "quant_config.json"), "w"), indent=2) | |
| for root, _dirs, files in os.walk(side_dir): | |
| rel = os.path.relpath(root, side_dir) | |
| if rel.startswith(".") and rel != ".": | |
| continue | |
| for f in files: | |
| if f.endswith((".safetensors", ".bin", ".gguf", ".onnx", ".md", ".gitattributes")) or f in LEAVE: | |
| continue | |
| dst = os.path.join(out_dir, rel if rel != "." else "") | |
| os.makedirs(dst, exist_ok=True) | |
| shutil.copy(os.path.join(root, f), dst) | |
| shutil.copy(os.path.abspath(__file__), os.path.join(out_dir, "qpack3.py")) | |
| return sd | |
| _LAZY_TOWERS = ("audio_tower.", "embed_audio.") | |
| _LAZY = (".packed", ".sq", ".zq", ".slo", ".shi", ".zlo", ".zhi") | |
| def _read_lazy(wfile, dtype): | |
| """Memory-map the audio tower so text and image workloads never make it resident; read everything | |
| else normally (memory-mapping the hot towers slowed the forward pass in a test).""" | |
| import struct | |
| import warnings | |
| import numpy as np | |
| from safetensors import safe_open | |
| with open(wfile, "rb") as f: | |
| n = struct.unpack("<Q", f.read(8))[0] | |
| hdr = json.loads(f.read(n)) | |
| base = 8 + n | |
| dts = {"U8": np.uint8, "F16": np.float16} | |
| sd = {} | |
| with safe_open(wfile, framework="pt") as f: | |
| for k in f.keys(): | |
| h = hdr[k] | |
| if k.startswith(_LAZY_TOWERS) and k.endswith(_LAZY) and h["dtype"] in dts: | |
| s0, _ = h["data_offsets"] | |
| arr = np.memmap(wfile, dtype=dts[h["dtype"]], mode="r", offset=base + s0, shape=tuple(h["shape"])) | |
| with warnings.catch_warnings(): | |
| warnings.simplefilter("ignore") | |
| sd[k] = torch.from_numpy(arr) | |
| else: | |
| v = f.get_tensor(k) | |
| sd[k] = v.to(dtype) if (dtype != torch.bfloat16 and v.dtype == torch.bfloat16) else v | |
| return sd | |
| def load_packed3(path, dtype=torch.bfloat16): | |
| from accelerate import init_empty_weights | |
| from transformers import AutoConfig | |
| from transformers.models.embedding_gemma2.modeling_embedding_gemma2 import EmbeddingGemma2Model | |
| qc = json.load(open(os.path.join(path, "quant_config.json"))) | |
| cfg = AutoConfig.from_pretrained(path) | |
| with init_empty_weights(include_buffers=False): | |
| model = EmbeddingGemma2Model(cfg) | |
| swap(model, | |
| lambda _p, l: QLinear3(l.in_features, l.out_features, l.bias is not None, | |
| bits=bits_for(_p, {k: set(v) for k, v in qc.get("int8_layers", {}).items()}, qc.get("audio_bits", 4))), | |
| lambda p, e: QEmbedding3(e.num_embeddings, e.embedding_dim, e.padding_idx, qc["embed_scale"][p])) | |
| pe = model.vision_tower.patch_embedder | |
| pe.position_embedding_table = nn.Parameter( | |
| torch.empty(2, qc["position_rows"], pe.position_embedding_table.shape[-1], device="meta"), | |
| requires_grad=False) | |
| sd = _read_lazy(os.path.join(path, "model.safetensors"), dtype) | |
| missing, unexpected = model.load_state_dict(sd, strict=False, assign=True) | |
| assert not unexpected, unexpected[:6] | |
| left = [n for n, t in list(model.named_parameters()) + list(model.named_buffers()) if t.is_meta] | |
| assert not left, f"still on meta: {left[:6]} missing={missing[:6]}" | |
| for m in model.modules(): | |
| if isinstance(m, QEmbedding3): | |
| m.compute_dtype = dtype | |
| return model.eval() | |
| def _return_model(model): | |
| """Make from_pretrained hand back an already built model, so SentenceTransformer reads no weights.""" | |
| import transformers.modeling_utils as mu | |
| from transformers.models.auto import auto_factory as af | |
| targets = [mu.PreTrainedModel, af._BaseAutoModelClass] | |
| saved = [(t, vars(t)["from_pretrained"]) for t in targets] | |
| try: | |
| for t in targets: | |
| t.from_pretrained = classmethod(lambda cls, *a, **k: model) | |
| yield | |
| finally: | |
| for t, orig in saved: | |
| t.from_pretrained = orig | |
| def load_st(path, dtype=torch.bfloat16, device="cpu"): | |
| from sentence_transformers import SentenceTransformer | |
| model = load_packed3(path, dtype) | |
| with _return_model(model): | |
| st = SentenceTransformer(path, device=device) | |
| return st | |