| |
| """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 |
| from wsp import protocol as P |
|
|
| 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() |
|
|
