| """Quant quality vs a reference model, on held-out code text. |
| |
| ref: python eval_kl.py ref MODEL OUT.npz [GLOB ...] -> saves top-K log-probs per token |
| cmp: python eval_kl.py cmp MODEL OUT.npz [GLOB ...] -> PPL, KL(ref||model), top-1 agreement |
| |
| KL uses the reference top-K tokens plus one bucket for the rest of the vocab. |
| """ |
| import glob, sys, time |
|
|
| import mlx.core as mx |
| import numpy as np |
| from mlx_lm import load |
|
|
| K, SEQ, NSEQ = 32, 1024, 48 |
|
|
|
|
| def corpus(tok, sources): |
| files = sorted({f for g in sources for f in glob.glob(g, recursive=True) |
| if "node_modules" not in f and "/." not in f and "venv" not in f}) |
| ids = [] |
| for f in files: |
| ids += tok.encode(open(f, errors="ignore").read(), add_special_tokens=False) |
| if len(ids) >= SEQ * NSEQ: |
| break |
| assert len(ids) >= SEQ * NSEQ, f"only {len(ids)} tokens" |
| return np.array(ids[: SEQ * NSEQ]).reshape(NSEQ, SEQ), len(files) |
|
|
|
|
| def logprobs(model, seq): |
| logits = model(mx.array(seq)[None])[0].astype(mx.float32) |
| return logits - mx.logsumexp(logits, axis=-1, keepdims=True) |
|
|
|
|
| def patch_uint8_packing(): |
| |
| import mlx_lm.utils as u |
| raw_load = mx.load |
|
|
| def load_u32(f, *a, **k): |
| w = raw_load(f, *a, **k) |
| return {n: v.view(mx.uint32) if v.dtype == mx.uint8 else v for n, v in w.items()} |
|
|
| u.mx.load = load_u32 |
|
|
|
|
| def main(mode, path, out, *globs): |
| sources = globs or ["**/*.py"] |
| if "flint-int4" in path or "/rec/models/" in path: |
| patch_uint8_packing() |
| model, tok = load(path) |
| data, nfiles = corpus(tok, sources) |
| t0 = time.time() |
| if mode == "ref": |
| top_i, top_lp, tgt_lp = [], [], [] |
| for n, seq in enumerate(data): |
| lp = logprobs(model, seq) |
| print(f" ref seq {n+1}/{NSEQ} {time.time()-t0:.0f}s", flush=True) |
| idx = mx.argpartition(-lp, K, axis=-1)[:, :K] |
| top_i.append(np.array(idx)); top_lp.append(np.array(mx.take_along_axis(lp, idx, -1))) |
| tgt_lp.append(np.array(mx.take_along_axis(lp[:-1], mx.array(seq[1:])[:, None], -1))[:, 0]) |
| np.savez(out, top_i=np.stack(top_i), top_lp=np.stack(top_lp), tgt_lp=np.stack(tgt_lp), data=data) |
| print(f"ref {path}: ppl {np.exp(-np.mean(tgt_lp)):.4f} from {nfiles} files, {time.time()-t0:.0f}s") |
| return |
| ref = np.load(out) |
| assert (ref["data"] == data).all(), "tokenized corpus differs from the reference run" |
| kls, agree, tgt = [], [], [] |
| for n, seq in enumerate(data): |
| lp = logprobs(model, seq) |
| ri, rlp = mx.array(ref["top_i"][n]), mx.array(ref["top_lp"][n]) |
| qlp = mx.take_along_axis(lp, ri, -1) |
| p, q = mx.exp(rlp), mx.exp(qlp) |
| p_rest = mx.maximum(1 - p.sum(-1), 1e-9); q_rest = mx.maximum(1 - q.sum(-1), 1e-9) |
| kl = (p * (rlp - qlp)).sum(-1) + p_rest * (mx.log(p_rest) - mx.log(q_rest)) |
| kls.append(np.array(kl)) |
| print(f" cmp seq {n+1}/{NSEQ} {time.time()-t0:.0f}s", flush=True) |
| ref_top1 = ri[mx.arange(SEQ), mx.argmax(rlp, -1)] |
| agree.append(np.array(mx.argmax(lp, -1) == ref_top1)) |
| tgt.append(np.array(mx.take_along_axis(lp[:-1], mx.array(seq[1:])[:, None], -1))[:, 0]) |
| kls = np.concatenate(kls) |
| print(f"cmp {path}: ppl {np.exp(-np.mean(tgt)):.4f} (ref {np.exp(-np.mean(ref['tgt_lp'])):.4f}) " |
| f"| KL mean {kls.mean():.4f} p99 {np.percentile(kls, 99):.4f} " |
| f"| top-1 agree {np.concatenate(agree).mean()*100:.2f}% | {time.time()-t0:.0f}s") |
|
|
|
|
| if __name__ == "__main__": |
| main(*sys.argv[1:]) |
|
|
