File size: 3,813 Bytes
2d2f883 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """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) # hard cap: allocations beyond it OOM
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")))
# warm-up (Triton compile), then timed suite, one request per call
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
# batched: all requests in one call
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]))
# exception recovery: fail mid-forward inside a streamed layer, then check the next answer is unchanged
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))
|