| """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)): |
| 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) |
|
|
