File size: 4,270 Bytes
b16c3a6 | 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 | #!/usr/bin/env python3
"""Decision path through vLLM: accuracy vs. the HF/PyTorch reference, throughput, and a mixed AR+decision call.
python3 scripts/vllm_decisions.py --bundle exports/jev-judge-qwen35-9b-v0.8 --hf-ref reports/hf_probs_9b_t30k.npz \
--out reports/vllm_decisions_9b.json
"""
from __future__ import annotations
import argparse
import json
import os
import sys
import time
import numpy as np
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "src"))
from jev_judge.data import load_split, target_to_slots # noqa: E402
from jev_judge.metrics import fmt_summary, per_sample_table, summarize # noqa: E402
from jev_judge.template import SLOT_RANGES, n_options_for # noqa: E402
def slot_mask_np(df) -> np.ndarray:
m = np.zeros((len(df), 24), dtype=bool)
for i, (k, o) in enumerate(zip(df["kind"], df["options"])):
s, _ = SLOT_RANGES[k]
m[i, s : s + n_options_for(k, list(o))] = True
return m
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--bundle", required=True)
ap.add_argument("--data", default="data")
ap.add_argument("--split", default="test_set_30k")
ap.add_argument("--hf-ref", default=None, help="npz with `probs` [N,24] from the PyTorch path on the same split")
ap.add_argument("--out", required=True)
ap.add_argument("--gpu-mem", type=float, default=0.85)
args = ap.parse_args()
from jev_judge.vllm_engine import VllmJudge
df = load_split(args.data, args.split)
rows = [{"kind": k, "state": s, "question": q, "options": list(o), "family": f}
for k, s, q, o, f in zip(df["kind"], df["state"], df["question"], df["options"], df["family"])]
t0 = time.time()
judge = VllmJudge(args.bundle, gpu_memory_utilization=args.gpu_mem)
load_s = time.time() - t0
judge.decide(rows[:256]) # warm-up (LoRA load, graph capture for these shapes)
t0 = time.time()
p = judge.decide(rows)
dt = time.time() - t0
n_tok = int(df["n_tokens"].sum())
q = np.stack([target_to_slots(k, t) for k, t in zip(df["kind"], df["target"])])
mask = slot_mask_np(df)
s = summarize(per_sample_table(p, q, mask, df))
res = {"bundle": args.bundle, "split": args.split, "n": len(df), "load_s": load_s, "decide_s": dt,
"rows_per_s": len(df) / dt, "prompt_tok_per_s": n_tok / dt, "metrics": s}
print(f"[vLLM decisions] {fmt_summary(s)}")
print(f"[vLLM decisions] {len(df)} rows in {dt:.1f}s -> {len(df)/dt:.0f} rows/s, {n_tok/dt:.0f} prompt tok/s (load {load_s:.0f}s)")
if args.hf_ref and os.path.exists(args.hf_ref):
ref = np.load(args.hf_ref)["probs"]
d = np.abs(p - ref)[mask]
agree = (np.where(mask, p, -1).argmax(1) == np.where(mask, ref, -1).argmax(1)).mean()
res["vs_hf"] = {"max_abs_dp": float(d.max()), "mean_abs_dp": float(d.mean()), "p99_abs_dp": float(np.percentile(d, 99)),
"argmax_agreement": float(agree)}
print(f"[vs PyTorch path] max|Δp| {d.max():.4f} · mean|Δp| {d.mean():.5f} · p99 {np.percentile(d, 99):.4f} · argmax agreement {agree:.4f}")
# one engine call, both heads
from vllm import SamplingParams
msgs = [[{"role": "user", "content": x}] for x in ["In one sentence, what is safety stock?",
"Write a Python one-liner that reverses a string."]]
gen_prompts = [judge.tok.apply_chat_template(m, tokenize=False, add_generation_prompt=True) for m in msgs]
t0 = time.time()
gen_outs, dec = judge.mixed(gen_prompts, SamplingParams(temperature=0.0, max_tokens=48), rows[:64])
res["mixed_call_s"] = time.time() - t0
res["mixed_decisions_match"] = float(np.abs(dec - p[:64]).max())
res["mixed_generation_sample"] = [o.outputs[0].text[:160] for o in gen_outs]
print(f"[mixed] 2 AR generations + 64 decisions in one call: {res['mixed_call_s']:.2f}s; "
f"decisions max|Δp| vs decision-only run {res['mixed_decisions_match']:.2e}")
for t in res["mixed_generation_sample"]:
print(" gen:", t.replace("\n", " ")[:140])
with open(args.out, "w") as f:
json.dump(res, f, indent=1)
print("->", args.out)
if __name__ == "__main__":
main()
|