File size: 3,653 Bytes
8948114 | 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 | """Low-bit weight storage for the char Transformer: matrices -> b-bit (4 or 8) symmetric integers with one fp16
scale per group of G consecutive input weights (per output row); vectors -> fp16. De-quantised to fp32 at load."""
import io, json
import numpy as np, torch
from bhc.xlit_model import XlitTransformer
def _q(v, bits, G):
r, c = v.shape; pad = (-c) % G; vp = np.pad(v, ((0, 0), (0, pad))).reshape(r, -1, G)
L = 2 ** (bits - 1) - 1; base = np.abs(vp).max(2, keepdims=True) / L; base[base == 0] = 1
best_err = None
for al in ((1.0,) if bits >= 8 else (1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65)): # MSE-optimal clipping per group
s_ = (base * al).astype(np.float16).astype(np.float32); q_ = np.clip(np.round(vp / s_), -L, L)
e = ((q_ * s_ - vp) ** 2).sum(2, keepdims=True)
if best_err is None: best_err, sc, q = e, s_, q_
else: m = e < best_err; best_err = np.where(m, e, best_err); sc = np.where(m, s_, sc); q = np.where(m, q_, q)
q = q.astype(np.int8)
return q.reshape(r, -1)[:, :c], sc[..., 0].astype(np.float16)
def export_q(model, path, bits=4, G=32, emb_bits=8):
"""bits for Linear weights; emb_bits for the (small, outlier-prone) embedding/position tables."""
sd = model.state_dict(); meta = {"cfg": model.cfg, "bits": bits, "G": G, "tensors": []}; buf = io.BytesIO()
for k, v in sd.items():
v = v.float().cpu().numpy()
if v.ndim == 2:
b_ = emb_bits if ("emb" in k or "pos" in k) else bits
q, sc = _q(v, b_, G); flat = q.reshape(-1)
if b_ == 4:
u = (flat + 8).astype(np.uint8); u = np.r_[u, np.zeros(len(u) % 2, np.uint8)]; data = (u[0::2] << 4 | u[1::2]).astype(np.uint8).tobytes()
else: data = flat.astype(np.int8).tobytes()
meta["tensors"].append([k, list(v.shape), "q%d" % b_, buf.tell(), len(data)]); buf.write(data); buf.write(sc.tobytes())
else:
meta["tensors"].append([k, list(v.shape), "f16", buf.tell(), 0]); buf.write(v.astype(np.float16).tobytes())
hdr = json.dumps(meta, separators=(",", ":")).encode()
with open(path, "wb") as f: f.write(len(hdr).to_bytes(4, "little")); f.write(hdr); f.write(buf.getvalue())
def load_q(path):
b = open(path, "rb").read(); n = int.from_bytes(b[:4], "little"); meta = json.loads(b[4:4 + n]); data = b[4 + n:]
bits, G = meta["bits"], meta["G"]; m = XlitTransformer(**meta["cfg"]); sd = {}
for k, shape, kind, off, ln in meta["tensors"]:
cnt = int(np.prod(shape))
if kind.startswith("q"):
bits = int(kind[1:]) if len(kind) > 1 else meta["bits"]
raw = np.frombuffer(data, np.uint8 if bits == 4 else np.int8, ln, off)
if bits == 4: q = (np.stack([raw >> 4, raw & 15], 1).reshape(-1)[:cnt].astype(np.int16) - 8)
else: q = raw.astype(np.int16)
r, c = shape; ng = (c + G - 1) // G
sc = np.frombuffer(data, np.float16, r * ng, off + ln).astype(np.float32).reshape(r, ng)
w = q.reshape(r, c).astype(np.float32) * np.repeat(sc, G, 1)[:, :c]
sd[k] = torch.from_numpy(w)
else:
sd[k] = torch.from_numpy(np.frombuffer(data, np.float16, cnt, off).astype(np.float32).reshape(shape).copy())
m.load_state_dict(sd); m.eval(); return m
def load_any(path):
if path.endswith(".int8"):
from bhc.xlit_model import load_int8; return load_int8(path)
if path.endswith(".pt"):
ck = torch.load(path, weights_only=False, map_location="cpu"); m = XlitTransformer(**ck["cfg"]); m.load_state_dict(ck["model"]); m.eval(); return m
return load_q(path)
|