| """Evaluate the INT4 package against the original FP8 model's outputs (ref.json) under a GPU memory cap. |
| usage: python evaluate.py <model_dir> <cap_gib> <out.json>""" |
| import json, math, sys, time |
| import torch |
|
|
| path, cap, out_path = sys.argv[1], float(sys.argv[2]), sys.argv[3] |
| sys.path.insert(0, path); sys.path.insert(0, __import__("os").path.dirname(__file__)) |
| total = torch.cuda.get_device_properties(0).total_memory |
| torch.cuda.set_per_process_memory_fraction(cap * 2**30 / total) |
| from prism_int4 import Prism |
| from suite import SUITE |
|
|
| t = time.time() |
| model = Prism.from_pretrained(path, gpu_memory_gb=cap) |
| load_s = time.time() - t |
| ref = json.load(open(__import__("os").path.join(__import__("os").path.dirname(__file__), "eval_reference_fp8.json"))) |
|
|
| |
| model.decide(**dict(zip(("state", "questions"), SUITE[0]))) |
| torch.cuda.synchronize(); torch.cuda.reset_peak_memory_stats() |
| times, res = [], [] |
| for s, q in SUITE: |
| b0 = model.streamer.bytes_streamed; t = time.time() |
| res.append(model.decide(state=s, questions=q)); torch.cuda.synchronize() |
| times.append((time.time() - t, model.streamer.bytes_streamed - b0)) |
| peak = torch.cuda.max_memory_allocated() / 2**30 |
|
|
| |
| t = time.time(); res_b = model.decide_batch(SUITE); torch.cuda.synchronize(); batch_s = time.time() - t |
|
|
|
|
| def dist(a): |
| if a["type"] == "yes_no": |
| return {"no": 1 - a["yes"], "yes": a["yes"]} |
| return a["probabilities"] |
|
|
|
|
| rows = [] |
| for r_ref, r_q in zip(ref, res): |
| for name in r_ref: |
| p, qd = dist(r_ref[name]), dist(r_q[name]) |
| keys = list(p) |
| kl = sum(p[k] * math.log(max(p[k], 1e-12) / max(qd[k], 1e-12)) for k in keys) |
| rows.append(dict(top_agree=max(keys, key=p.get) == max(keys, key=qd.get), |
| max_abs=max(abs(p[k] - qd[k]) for k in keys), kl=kl, type=r_ref[name]["type"], |
| ref_top=max(p.values()))) |
| n = len(rows) |
| batch_max_diff = max(abs(dist(a[k])[o] - dist(b[k])[o]) for ra, rb in zip(res, res_b) for k in ra |
| for a, b in [(ra, rb)] for o in dist(ra[k])) |
|
|
| |
| streamed = model.streamer.order |
| recovery = None |
| if len(streamed) > 2: |
| mod = model.backbone.layers[streamed[2]].mlp |
| orig = mod._moe_dispatch |
| def boom(*a, **k): |
| raise RuntimeError("injected failure") |
| mod._moe_dispatch = boom |
| try: |
| model.decide(state=SUITE[1][0], questions=SUITE[1][1]) |
| except RuntimeError: |
| pass |
| mod._moe_dispatch = orig |
| again = model.decide(state=SUITE[0][0], questions=SUITE[0][1]) |
| recovery = max(abs(dist(again[k])[o] - dist(res[0][k])[o]) for k in again for o in dist(again[k])) |
|
|
| summary = dict( |
| cap_gib=cap, load_seconds=round(load_s), peak_gpu_gib=round(peak, 2), |
| resident_layers=len(model.streamer.gpu), streamed_layers=len(streamed), |
| questions=n, top1_agreement=sum(r["top_agree"] for r in rows) / n, |
| mean_max_abs_prob_diff=sum(r["max_abs"] for r in rows) / n, worst_max_abs_prob_diff=max(r["max_abs"] for r in rows), |
| mean_kl=sum(r["kl"] for r in rows) / n, worst_kl=max(r["kl"] for r in rows), |
| top1_disagreements=[dict(r) for r in rows if not r["top_agree"]], |
| median_request_seconds=sorted(x[0] for x in times)[len(times) // 2], |
| gb_streamed_per_forward=round(sorted(x[1] for x in times)[len(times) // 2] / 1e9, 1), |
| batched_all_requests_seconds=round(batch_s, 2), batched_vs_single_max_diff=batch_max_diff, |
| recovery_after_injected_failure_max_diff=recovery) |
| json.dump(dict(summary=summary, results=res), open(out_path, "w"), ensure_ascii=False, indent=1) |
| print(json.dumps(summary, indent=1)) |
|
|
