| """Real-model check: quantize -> load -> WikiText-2 perplexity, for RTN and shaped, uniform INT4 and mixed. |
| usage: python tests/real_models.py <hf_model_dir_or_id> <out_dir> [--group 128] [--budget 4.75] [--windows 32] [--bos ID] [--text wt2_test.txt]""" |
| import argparse, json, os, time, torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| from shapequant import quantize_model, load_quantized, perplexity, wikitext2_windows |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("model"); p.add_argument("out"); p.add_argument("--group", type=int, default=128) |
| p.add_argument("--budget", type=float, default=4.75); p.add_argument("--windows", type=int, default=32) |
| p.add_argument("--bos", type=int, default=None); p.add_argument("--text", default=None); p.add_argument("--device", default="cpu") |
| a = p.parse_args() |
| tok = AutoTokenizer.from_pretrained(a.model) |
| W = wikitext2_windows(tok, 1024, a.windows, a.text) |
| R = {"dense": perplexity(AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.float32).eval().to(a.device), W, a.bos)} |
| for tag, kw in (("rtn_int4", dict(bits=4, method="rtn")), ("shaped_int4", dict(bits=4)), |
| ("rtn_mixed", dict(budget=a.budget, method="rtn")), ("shaped_mixed", dict(budget=a.budget))): |
| t = time.time(); d = os.path.join(a.out, tag) |
| s = quantize_model(a.model, d, group_size=a.group, device=a.device, log=lambda *x: None, **kw) |
| R[tag] = {"ppl": perplexity(load_quantized(d, dtype=torch.float32, device=a.device), W, a.bos), |
| "minutes": (time.time() - t) / 60, "checkpoint_MB": s["checkpoint_bytes"] / 2 ** 20} |
| print(tag, R[tag], flush=True) |
| json.dump(R, open(os.path.join(a.out, "results.json"), "w"), indent=1) |
|
|
