File size: 5,901 Bytes
a12e613 | 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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """Sequential-batch eval harness for a single-GPU box: batches of samples go teacher (concurrent
vLLM) -> king reads -> challenger reads -> compare, keeping states files within the TTL and the
shm budget. Same scoring as the eval service.
.venv/bin/python eval_seq.py --dataset pick20.jsonl --teacher DIR --king DIR --challenger DIR
[--samples N] [--batch 8] [--workers 8] [--out results.jsonl]
"""
from __future__ import annotations
import argparse
import json
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from eval_service import dataset, metrics, run, servers
from eval_service.run import Distance, _render, generate, load_tokenizer
def aggregate(sums: list, role: str) -> Distance:
rows = [by_role[role] for by_role in sums if role in by_role]
if not rows:
return Distance(kl=float("inf"), dist=float("inf"))
return Distance(
kl=sum(s.kl for s in rows) / sum(s.predictions for s in rows),
dist=sum(s.layer_mean_dist for s in rows) / max(sum(s.activations for s in rows), 1),
)
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--dataset", type=Path, required=True)
ap.add_argument("--teacher", type=Path, required=True)
ap.add_argument("--king", type=Path, required=True)
ap.add_argument("--challenger", type=Path, required=True)
ap.add_argument("--samples", type=int, default=20)
ap.add_argument("--batch", type=int, default=8)
ap.add_argument("--workers", type=int, default=8)
ap.add_argument("--out", type=Path, default=Path("eval_seq.jsonl"))
args = ap.parse_args()
samples = dataset.load(args.dataset).samples[: args.samples]
tokenizer = load_tokenizer(args.teacher)
models = {"teacher": args.teacher, "king": args.king, "challenger": args.challenger}
device = "cuda:0"
heads = metrics.heads(models, [device])
all_sums = []
records_out = []
try:
for b0 in range(0, len(samples), args.batch):
batch = samples[b0 : b0 + args.batch]
records = []
groups = servers.serve({"teacher": args.teacher})
with ThreadPoolExecutor(args.workers) as pool:
def one(i, sample, prompt):
answer, tpath = generate(groups["teacher"], prompt, i, sample.id)
for _ in range(120): # the connector writes states asynchronously
if Path(tpath).exists():
break
time.sleep(0.5)
keep = Path("/root/states") / Path(tpath).name
keep.write_bytes(Path(tpath).read_bytes())
Path(str(keep) + ".lock").touch()
return answer, str(keep)
futures = {i: (s, p, pool.submit(one, b0 + i, s, p)) for i, (s, p) in enumerate((sample, _render(tokenizer, sample.messages, sample.tools)) for sample in batch)}
for i, (sample, prompt, fut) in futures.items():
answer, tpath = fut.result()
records.append({"id": sample.id, "prompt": prompt, "answer": answer, "tpath": tpath})
print(f"teacher {b0+i+1}/{len(samples)}: {sample.id} p{len(prompt)} a{len(answer)}", flush=True)
servers.release("teacher")
for role, dir_ in (("king", args.king), ("challenger", args.challenger)):
groups = servers.serve({role: dir_})
group = groups[role]
for i, rec in enumerate(records):
read, path = group.complete(rec["prompt"], len(rec["answer"]), b0 + i, rec["prompt"] + rec["answer"])
if read != rec["answer"]:
raise SystemExit(f"{role} left the teacher's answer on {rec['id']}")
rec[role] = path
print(f"{role} {b0+i+1}/{len(samples)}: {rec['id']}", flush=True)
servers.release(role)
for rec in records:
read = rec["prompt"] + rec["answer"][:-1]
states = {
role: metrics.load_states(rec[role] if role != "teacher" else rec["tpath"], read, len(rec["prompt"]) - 1)
for role in ("teacher", "king", "challenger")
}
by_role = metrics.compare(states, heads[device], rec["answer"], device)
all_sums.append(by_role)
records_out.append({"id": rec["id"], **{role: s.summary() for role, s in by_role.items()}})
print(
f"scored {rec['id']}: king KL {by_role['king'].kl / by_role['king'].predictions:.4f}"
f" challenger KL {by_role['challenger'].kl / by_role['challenger'].predictions:.4f}",
flush=True,
)
for p in (rec["tpath"], rec["king"], rec["challenger"]):
servers.discard(p)
agg = {role: aggregate(all_sums, role) for role in ("teacher", "king", "challenger")}
for role, d in agg.items():
print(f"{role:12s} KL {d.kl:.6f} distance {d.dist:.5f}")
if agg["king"].kl < float("inf"):
need = Distance(kl=0.015 * agg["king"].kl, dist=0.0084 * agg["king"].dist)
print(
f"margin needed: KL {need.kl:.6f} dist {need.dist:.6f} ;"
f" challenger leads by KL {agg['king'].kl - agg['challenger'].kl:.6f}"
f" ({100 * (agg['king'].kl - agg['challenger'].kl) / agg['king'].kl:.3f}%)"
f" dist {agg['king'].dist - agg['challenger'].dist:.6f}"
f" ({100 * (agg['king'].dist - agg['challenger'].dist) / agg['king'].dist:.3f}%)"
)
args.out.write_text("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in records_out))
finally:
servers.stop_all()
if __name__ == "__main__":
main()
|