| |
| |
| |
| """Quantization formats: FP16 / BF16 / INT8 / INT4. |
| |
| All quantizers are implemented here from first principles (no external |
| quantizer libraries). Every quantize() returns (payload, meta) and every |
| meta carries enough information for an exact dequantize() round-trip. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
|
|
| def to_fp16(W: np.ndarray): |
| W = np.asarray(W, dtype=np.float32) |
| return W.astype(np.float16), {"fmt": "fp16"} |
|
|
|
|
| def from_fp16(q: np.ndarray) -> np.ndarray: |
| return np.asarray(q, dtype=np.float32) |
|
|
|
|
| def to_bf16(W: np.ndarray): |
| """Round-to-nearest-even float32 -> bfloat16 (stored as uint16).""" |
| a = np.asarray(W, dtype=np.float32) |
| i = a.view(np.int32).copy() |
| finite = np.isfinite(a) |
| lsb = (i >> 16) & 1 |
| i[finite] += (0x7FFF + lsb[finite]).astype(np.int32) |
| return (i >> 16).astype(np.uint16), {"fmt": "bf16"} |
|
|
|
|
| def from_bf16(q: np.ndarray) -> np.ndarray: |
| return (np.asarray(q, dtype=np.int32) << 16).view(np.float32) |
|
|
|
|
| def quantize_int8(W: np.ndarray): |
| """Asymmetric per-tensor INT8.""" |
| W = np.asarray(W, dtype=np.float64) |
| wmin, wmax = float(W.min()), float(W.max()) |
| if wmax == wmin: |
| return np.zeros(W.shape, dtype=np.int8), { |
| "fmt": "int8", "scale": 1.0, "zero": 0.0} |
| scale = (wmax - wmin) / 255.0 |
| q = np.clip(np.round((W - wmin) / scale), 0, 255).astype(np.int64) - 128 |
| return q.astype(np.int8), {"fmt": "int8", "scale": scale, "zero": wmin} |
|
|
|
|
| def dequantize_int8(q: np.ndarray, meta: dict) -> np.ndarray: |
| return (q.astype(np.float64) + 128.0) * meta["scale"] + meta["zero"] |
|
|
|
|
| def quantize_int4(W: np.ndarray, group: int = 32): |
| """Group-wise asymmetric INT4, two nibbles packed per byte (low nibble |
| first). Groups run along the last axis; ragged tail is zero-padded and |
| the pad is recorded so dequantize() restores the exact shape.""" |
| W = np.asarray(W, dtype=np.float64) |
| shape = list(W.shape) |
| last = shape[-1] |
| rows = int(np.prod(shape[:-1])) if len(shape) > 1 else 1 |
| flat = W.reshape(rows, last) |
| pad = (-last) % group |
| if pad: |
| flat = np.concatenate([flat, np.zeros((rows, pad))], axis=1) |
| G = flat.reshape(-1, group) |
| gmin = G.min(axis=1) |
| gmax = G.max(axis=1) |
| scale = np.maximum(gmax - gmin, 1e-300) / 15.0 |
| q = np.clip(np.round((G - gmin[:, None]) / scale[:, None]), 0, 15).astype(np.uint8) |
| packed = (q[:, 0::2] | (q[:, 1::2] << 4)).astype(np.uint8) |
| meta = {"fmt": "int4", "group": group, "shape": shape, |
| "scales": scale.astype(np.float64), "zeros": gmin.astype(np.float64), |
| "rows": rows, "last": last} |
| return packed, meta |
|
|
|
|
| def dequantize_int4(packed: np.ndarray, meta: dict) -> np.ndarray: |
| group, rows, last = meta["group"], meta["rows"], meta["last"] |
| p = np.asarray(packed).reshape(-1, group // 2) |
| q = np.empty((p.shape[0], group), dtype=np.float64) |
| q[:, 0::2] = p & 0x0F |
| q[:, 1::2] = (p >> 4) & 0x0F |
| G = q * np.asarray(meta["scales"])[:, None] + np.asarray(meta["zeros"])[:, None] |
| flat = G.reshape(rows, -1)[:, :last] |
| return flat.reshape(meta["shape"]).astype(np.float32) |
|
|
|
|
| def quantize_tensor(W: np.ndarray, fmt: str): |
| """Dispatch: fmt in {fp32, fp16, bf16, int8, int4}.""" |
| if fmt == "fp32": |
| return np.asarray(W, dtype=np.float32), {"fmt": "fp32"} |
| if fmt == "fp16": |
| return to_fp16(W) |
| if fmt == "bf16": |
| return to_bf16(W) |
| if fmt == "int8": |
| return quantize_int8(W) |
| if fmt == "int4": |
| return quantize_int4(W) |
| raise ValueError(f"unknown format {fmt!r}") |
|
|
|
|
| def dequantize_tensor(q: np.ndarray, meta: dict) -> np.ndarray: |
| fmt = meta["fmt"] |
| if fmt == "fp32": |
| return np.asarray(q, dtype=np.float32) |
| if fmt == "fp16": |
| return from_fp16(q) |
| if fmt == "bf16": |
| return from_bf16(q) |
| if fmt == "int8": |
| return dequantize_int8(q, meta).astype(np.float32) |
| if fmt == "int4": |
| return dequantize_int4(q, meta) |
| raise ValueError(f"unknown format {fmt!r}") |
|
|
|
|
| def payload_nbytes(q: np.ndarray, fmt: str) -> int: |
| return int(np.asarray(q).nbytes) |
|
|
|
|
| |
| |
| |
| |
|
|
| def quantize_int4_per_tensor(W: np.ndarray): |
| """Whole-tensor asymmetric INT4: ONE scale/zero over the entire tensor.""" |
| W = np.asarray(W, dtype=np.float64) |
| shape = list(W.shape) |
| flat = W.ravel() |
| gmin, gmax = float(flat.min()), float(flat.max()) |
| scale = max(gmax - gmin, 1e-300) / 15.0 |
| q = np.clip(np.round((flat - gmin) / scale), 0, 15).astype(np.uint8) |
| if q.size % 2: |
| q = np.concatenate([q, np.zeros(1, dtype=np.uint8)]) |
| packed = (q[0::2] | (q[1::2] << 4)).astype(np.uint8) |
| meta = {"fmt": "int4", "scheme": "per-tensor", "shape": shape, |
| "scales": np.array([scale]), "zeros": np.array([gmin]), |
| "size": flat.size} |
| return packed, meta |
|
|
|
|
| def dequantize_int4_per_tensor(packed: np.ndarray, meta: dict) -> np.ndarray: |
| q = np.asarray(packed) |
| flat = np.empty(q.size * 2, dtype=np.float64) |
| flat[0::2] = q & 0x0F |
| flat[1::2] = (q >> 4) & 0x0F |
| flat = flat[:meta["size"]] |
| W = flat * float(meta["scales"][0]) + float(meta["zeros"][0]) |
| return W.reshape(meta["shape"]).astype(np.float32) |
|
|
|
|
| def quantize_int4_per_channel(W: np.ndarray): |
| """Per-channel (per-row) asymmetric INT4: one scale/zero per row instead |
| of per group-of-32. Returns (packed, meta) with the same nibble layout |
| as quantize_int4 so dequantize_int4-style unpacking applies.""" |
| W = np.asarray(W, dtype=np.float64) |
| shape = list(W.shape) |
| last = shape[-1] |
| rows = int(np.prod(shape[:-1])) if len(shape) > 1 else 1 |
| flat = W.reshape(rows, last) |
| gmin = flat.min(axis=1) |
| gmax = flat.max(axis=1) |
| scale = np.maximum(gmax - gmin, 1e-300) / 15.0 |
| q = np.clip(np.round((flat - gmin[:, None]) / scale[:, None]), 0, 15) |
| q = q.astype(np.uint8) |
| |
| if last % 2: |
| q = np.concatenate([q, np.zeros((rows, 1), dtype=np.uint8)], axis=1) |
| packed = (q[:, 0::2] | (q[:, 1::2] << 4)).astype(np.uint8) |
| meta = {"fmt": "int4", "scheme": "per-channel", "shape": shape, |
| "scales": scale.astype(np.float64), |
| "zeros": gmin.astype(np.float64), |
| "rows": rows, "last": last} |
| return packed, meta |
|
|
|
|
| def dequantize_int4_per_channel(packed: np.ndarray, meta: dict) -> np.ndarray: |
| rows, last = meta["rows"], meta["last"] |
| p = np.asarray(packed).reshape(rows, -1) |
| q = np.empty((rows, p.shape[1] * 2), dtype=np.float64) |
| q[:, 0::2] = p & 0x0F |
| q[:, 1::2] = (p >> 4) & 0x0F |
| q = q[:, :last] |
| W = q * np.asarray(meta["scales"])[:, None] + np.asarray(meta["zeros"])[:, None] |
| return W.reshape(meta["shape"]).astype(np.float32) |
|
|
|
|
| def search_group_size(W: np.ndarray, |
| groups: tuple = (16, 32, 64, 128)) -> dict: |
| """Search the INT4 group size by MEASURED reconstruction MSE — the |
| winner is the argmin of measured error, never a hard-coded default.""" |
| W = np.asarray(W, dtype=np.float64) |
| results = [] |
| for g in groups: |
| if g < 2: |
| continue |
| q, meta = quantize_int4(W, group=g) |
| back = dequantize_int4(q, meta).astype(np.float64) |
| err = float(np.mean((W - back) ** 2)) |
| results.append({"group": g, "mse": err, |
| "bytes": int(np.asarray(q).nbytes)}) |
| results.sort(key=lambda r: r["mse"]) |
| return {"results": results, "best_group": results[0]["group"], |
| "best_mse": results[0]["mse"]} |
|
|
|
|
| def quant_error_map(W: np.ndarray, Wd: np.ndarray) -> dict: |
| """Instrument quantization error per tensor. All values measured.""" |
| W = np.asarray(W, dtype=np.float64) |
| Wd = np.asarray(Wd, dtype=np.float64) |
| err = W - Wd |
| flat = err.ravel() |
| sd = float(flat.std()) if flat.size else 0.0 |
| per_channel = [] |
| if W.ndim == 2: |
| for r in range(W.shape[0]): |
| e = err[r] |
| per_channel.append(float(np.mean(e ** 2))) |
| |
| per_group = [] |
| if W.ndim >= 1 and W.shape[-1] >= 32: |
| G = err.reshape(-1, 32) |
| per_group = [float(np.mean(g ** 2)) for g in G] |
| denom = float(np.linalg.norm(W)) |
| return { |
| "min": float(err.min()) if err.size else 0.0, |
| "max": float(err.max()) if err.size else 0.0, |
| "mean": float(err.mean()) if err.size else 0.0, |
| "variance": float(err.var()) if err.size else 0.0, |
| "std": sd, |
| "outlier_frac": float(np.mean(np.abs(flat) > 3 * sd)) if sd > 0 else 0.0, |
| "mse": float(np.mean(flat ** 2)) if flat.size else 0.0, |
| "cosine_sim": float(np.dot(W.ravel(), Wd.ravel()) |
| / (np.linalg.norm(W.ravel()) |
| * np.linalg.norm(Wd.ravel()) + 1e-300)), |
| "rel_err": float(np.linalg.norm(err) / (denom + 1e-300)), |
| "per_channel_mse": per_channel, |
| "per_group_mse": per_group, |
| } |
|
|
|
|
| def int4_scheme_report(W: np.ndarray) -> dict: |
| """Benchmark per-tensor / per-channel / group-wise INT4 on one tensor. |
| |
| Returns measured MSE + byte size for each scheme; the caller picks the |
| winner from measurements.""" |
| W = np.asarray(W, dtype=np.float64) |
| out = {} |
| qt, mt = quantize_int4_per_tensor(W) |
| backt = dequantize_int4_per_tensor(qt, mt).astype(np.float64) |
| out["per-tensor"] = {"mse": float(np.mean((W - backt) ** 2)), |
| "bytes": int(np.asarray(qt).nbytes), |
| "meta_bytes": 16} |
| qp, mp = quantize_int4_per_channel(W) |
| backp = dequantize_int4_per_channel(qp, mp).astype(np.float64) |
| out["per-channel"] = {"mse": float(np.mean((W - backp) ** 2)), |
| "bytes": int(np.asarray(qp).nbytes), |
| "meta_bytes": int(mp["scales"].nbytes |
| + mp["zeros"].nbytes)} |
| gs = search_group_size(W) |
| out["group-wise"] = {"mse": gs["best_mse"], |
| "bytes": next(r["bytes"] for r in gs["results"] |
| if r["group"] == gs["best_group"]), |
| "best_group": gs["best_group"], |
| "all": gs["results"]} |
| |
| qi, mi = quantize_int8(W) |
| backi = dequantize_int8(qi, mi).astype(np.float64) |
| out["int8"] = {"mse": float(np.mean((W - backi) ** 2)), |
| "bytes": int(np.asarray(qi).nbytes)} |
| return out |
|
|
|
|
| def mixed_precision_plan(params: dict, size_budget_bytes: int) -> dict: |
| """Assign INT4 vs INT8 per tensor under a physical size budget. |
| |
| Tensors where INT8 buys the most error reduction per extra byte stay |
| INT8; the rest go INT4 (best measured scheme per tensor). Sensitive |
| tensors are thereby protected by measurement, not by fiat. |
| """ |
| rows = [] |
| for k, v in params.items(): |
| v = np.asarray(v) |
| if v.ndim != 2 or v.size < 64: |
| rows.append({"tensor": k, "fmt": "fp32", |
| "bytes": int(v.nbytes), "gain_per_byte": 0.0, |
| "int4_mse": 0.0, "int8_mse": 0.0}) |
| continue |
| rep = int4_scheme_report(v) |
| best4 = min(rep["per-tensor"], rep["per-channel"], |
| rep["group-wise"], key=lambda r: r["mse"]) |
| b4 = best4["bytes"] + best4.get("meta_bytes", 0) |
| b8 = rep["int8"]["bytes"] |
| gain = max(0.0, best4["mse"] - rep["int8"]["mse"]) |
| gpb = gain / max(1, b8 - b4) |
| rows.append({"tensor": k, "fmt": "int4", "bytes": b4, |
| "gain_per_byte": gpb, "int4_mse": best4["mse"], |
| "int8_mse": rep["int8"]["mse"], |
| "int4_scheme": "per-tensor" |
| if best4 is rep["per-tensor"] else |
| ("per-channel" if best4 is rep["per-channel"] |
| else f"group-{rep['group-wise']['best_group']}")}) |
| |
| total4 = sum(r["bytes"] for r in rows) |
| cands = sorted([r for r in rows if r["fmt"] == "int4" |
| and r["gain_per_byte"] > 0], |
| key=lambda r: -r["gain_per_byte"]) |
| plan, used = {}, total4 |
| for r in rows: |
| plan[r["tensor"]] = "int4" |
| for r in cands: |
| v = np.asarray(params[r["tensor"]]) |
| qi, _ = quantize_int8(v) |
| b8 = int(np.asarray(qi).nbytes) |
| if used - r["bytes"] + b8 <= size_budget_bytes: |
| plan[r["tensor"]] = "int8" |
| used = used - r["bytes"] + b8 |
| return {"plan": plan, "total_bytes": used, |
| "budget_bytes": int(size_budget_bytes), |
| "within_budget": used <= size_budget_bytes, |
| "n_int8": sum(1 for v in plan.values() if v == "int8"), |
| "n_int4": sum(1 for v in plan.values() if v == "int4"), |
| "rows": rows} |
|
|
|
|
| def dequantize_dispatch(q: np.ndarray, meta: dict) -> np.ndarray: |
| """Dequantize honoring per-tensor / per-channel int4 metadata.""" |
| if meta.get("fmt") == "int4" and meta.get("scheme") == "per-tensor": |
| return dequantize_int4_per_tensor(q, meta) |
| if meta.get("fmt") == "int4" and meta.get("scheme") == "per-channel": |
| return dequantize_int4_per_channel(q, meta) |
| return dequantize_tensor(q, meta) |
|
|
|
|
| def recovery_train(s_qparams: dict, s_cfg, t_params: dict, t_cfg, |
| *, steps: int = 30, batch: int = 64, lr: float = 1e-3, |
| seed: int = 0) -> dict: |
| """Quantization-aware recovery (spec section 9): |
| |
| INT4-dequantized student -> distillation (task + logit KL) |
| vs the teacher -> requantize INT4 -> evaluate. |
| |
| Returns measured ppl before/after and the gap delta. Improvement is |
| reported ONLY if measured; otherwise the (non-)result stands. |
| """ |
| from .model import (Adam, backward, ce_loss, eval_metrics, forward, |
| gen_data, kl_loss) |
| rng = np.random.default_rng(seed) |
| |
| params = {k: v.astype(np.float32) for k, v in s_qparams.items()} |
| before = eval_metrics(params, s_cfg, batches=4, batch=batch, seed=911) |
| opt = Adam(params, lr=lr) |
| for _ in range(steps): |
| ids = gen_data(rng, batch, s_cfg) |
| tl, _ = forward(t_params, t_cfg, ids[:, :-1]) |
| sl, cache = forward(params, s_cfg, ids[:, :-1]) |
| lt, dl_t = ce_loss(sl, ids[:, 1:]) |
| lk, dl_k = kl_loss(sl, tl, tau=2.0) |
| grads = backward(params, s_cfg, cache, dl_t + 0.5 * dl_k) |
| opt.step(params, grads) |
| after = eval_metrics(params, s_cfg, batches=4, batch=batch, seed=911) |
| |
| re_q = {} |
| for k, v in params.items(): |
| if v.ndim == 2 and v.size >= 64: |
| q, m = quantize_int4(v) |
| re_q[k] = dequantize_int4(q, m).astype(np.float32) |
| else: |
| re_q[k] = v |
| re_eval = eval_metrics(re_q, s_cfg, batches=4, batch=batch, seed=911) |
| return { |
| "steps": steps, |
| "ppl_before": float(np.exp(before["loss"])), |
| "ppl_after_train": float(np.exp(after["loss"])), |
| "ppl_after_requant": float(np.exp(re_eval["loss"])), |
| "gap_before": float(np.exp(before["loss"]) - np.exp(after["loss"])), |
| "gap_delta": float(np.exp(before["loss"]) |
| - np.exp(re_eval["loss"])), |
| "improved": bool(np.exp(re_eval["loss"]) < np.exp(before["loss"])), |
| } |
|
|
