File size: 5,382 Bytes
84bcfb1 | 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 | #!/usr/bin/env python
"""One vLLM process on one GPU: answer a shard of a problem file, then grade every answer.
Writes <out>.jsonl (one graded record per problem), <out>.texts.jsonl.gz (the full answers)
and <out>.done.json (settings, timing, throughput). A shard "i/k" takes problems i, i+k, ...
"""
import argparse
import gzip
import json
import os
import sys
import time
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from wsp import grading as G # noqa: E402
from wsp import protocol as P # noqa: E402
TYPES = {"temperature": float, "top_p": float, "top_k": int, "presence_penalty": float,
"repetition_penalty": float, "max_tokens": int}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--model", required=True)
ap.add_argument("--data", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--n", type=int, default=1)
ap.add_argument("--shard", default="0/1")
ap.add_argument("--preset", default="think35", choices=sorted(P.PRESETS))
ap.add_argument("--set", action="append", default=[], help="override a sampling value, e.g. temperature=1.0")
ap.add_argument("--gpu-mem", type=float, default=0.92)
ap.add_argument("--max-num-seqs", type=int, default=256)
ap.add_argument("--limit", type=int, default=0)
ap.add_argument("--seed", type=int, default=0, help="per-request seed as in the paper; -1 for none")
ap.add_argument("--with-vision", action="store_true", help="also load the vision encoder (not needed)")
ap.add_argument("--prefix-cache", action="store_true",
help="vLLM prefix caching (off: prompts are ~150 tokens, answers ~30k, and on this hybrid "
"model it keeps extra recurrent-state copies)")
a = ap.parse_args()
cfg = dict(P.PRESETS[a.preset])
for kv in a.set:
k, v = kv.split("=", 1)
cfg[k] = (v.lower() in ("1", "true", "on")) if k == "thinking" else TYPES[k](v)
i, k = map(int, a.shard.split("/"))
with open(a.data) as f:
rows = [json.loads(line) for line in f if line.strip()][i::k]
if a.limit:
rows = rows[:a.limit]
from vllm import LLM, SamplingParams
t0 = time.time()
extra = {} if a.with_vision else {"limit_mm_per_prompt": {"image": 0, "video": 0}}
llm = LLM(model=a.model, dtype="bfloat16", gpu_memory_utilization=a.gpu_mem,
max_model_len=cfg["max_tokens"] + 2048, seed=0, max_num_seqs=a.max_num_seqs,
disable_log_stats=False, enable_prefix_caching=a.prefix_cache, **extra)
tok = llm.get_tokenizer()
prompts = [P.render(tok, r, cfg["thinking"]) for r in rows]
sp = SamplingParams(n=a.n, temperature=cfg["temperature"], top_p=cfg["top_p"], top_k=cfg["top_k"],
presence_penalty=cfg["presence_penalty"], repetition_penalty=cfg["repetition_penalty"],
max_tokens=cfg["max_tokens"], seed=None if a.seed < 0 else a.seed)
t1 = time.time()
outs = llm.generate(prompts, sp)
t2 = time.time()
bad = sum(1 for o in outs for x in o.outputs if x.finish_reason not in ("stop", "length"))
if bad:
raise RuntimeError(f"{bad} generations neither stopped nor reached max_tokens")
ntok = 0
tmp = a.out + ".jsonl.tmp"
with open(tmp, "w") as fr, gzip.open(a.out + ".texts.jsonl.gz", "wt") as ft:
for r, o in zip(rows, outs):
samples = []
for j, s in enumerate(o.outputs):
text = s.text
n_t = len(s.token_ids)
ntok += n_t
rec = {"tokens": n_t, "finish": s.finish_reason, "cut": s.finish_reason == "length",
"closed": "</think>" in text,
"verbatim_loop_end": G.verbatim_loop_end(text),
"repetitive_end": G.repetitive_end(text),
"tail_dup8": round(G.tail_dup8(text), 4)}
if r.get("kind") == "ifeval":
ok, follow = G.ifeval_follow(r, text, cfg["thinking"])
rec.update(correct=bool(ok), follow=follow, pred=None)
else:
pred, nbox = G.extract(text)
rec.update(pred=pred, n_boxes=nbox, correct=G.is_correct(pred, r["answer"]))
samples.append(rec)
ft.write(json.dumps({"id": r["id"], "j": j, "text": text}) + "\n")
fr.write(json.dumps({"id": r["id"], "source": r.get("source"), "answer": r.get("answer"),
"samples": samples}) + "\n")
os.replace(tmp, a.out + ".jsonl")
note = os.path.join(a.model, "edit_note.json")
done = {"model": a.model, "edit_note": json.load(open(note)) if os.path.exists(note) else None,
"data": a.data, "shard": a.shard, "n": a.n, "preset": a.preset, "sampling": cfg, "seed": a.seed,
"prompt_example": prompts[0] if prompts else None,
"problems": len(rows), "samples": len(rows) * a.n, "tokens": ntok,
"load_s": round(t1 - t0, 1), "gen_s": round(t2 - t1, 1),
"tok_per_s": round(ntok / max(t2 - t1, 1e-9)), "grade_s": round(time.time() - t2, 1)}
with open(a.out + ".done.json", "w") as f:
json.dump(done, f, indent=1)
print("DONE", json.dumps(done), flush=True)
if __name__ == "__main__":
main()
|