| """AIME25 avg@k 评测。 |
| |
| 针对 non-thinking 模型: prompt 必须用 enable_thinking=False 渲染 |
| (尾部会带 <think>\n\n</think>\n\n)。见 MEMORY: nothink-flag-for-1p7b-sft。 |
| |
| 答案抽取按优先级: \boxed{} -> "Answer: X" -> 文末最后一个整数。 |
| AIME 答案恒为 0-999 的整数, 比较时按 int 比。 |
| """ |
| import argparse |
| import json |
| import os |
| import re |
| import sys |
| import time |
| from collections import Counter |
|
|
| p = argparse.ArgumentParser() |
| p.add_argument("--model", required=True) |
| p.add_argument("--data", default="/home/haojinw2/datasets/aime25/test.jsonl") |
| p.add_argument("--out", required=True) |
| p.add_argument("--n", type=int, default=16) |
| p.add_argument("--max-new", type=int, default=8192) |
| p.add_argument("--ctx", type=int, default=10240) |
| p.add_argument("--tp", type=int, default=1) |
| p.add_argument("--temp", type=float, default=0.7) |
| p.add_argument("--top-p", type=float, default=0.8) |
| p.add_argument("--top-k", type=int, default=20) |
| p.add_argument("--seed", type=int, default=0) |
| p.add_argument("--enable-thinking", type=lambda s: s.lower() == "true", default=False, |
| help="False = 关闭 think(non-thinking 模型必须用 False)") |
| p.add_argument("--gpu-util", type=float, default=0.85) |
| a = p.parse_args() |
|
|
|
|
| def log(*x): |
| print(f"[{time.strftime('%H:%M:%S')}]", *x, flush=True) |
|
|
|
|
| INSTR = "Please reason step by step, and put your final answer within \\boxed{}." |
|
|
| BOXED = re.compile(r"\\boxed\s*\{([^{}]*)\}") |
| ANSW = re.compile(r"[Aa]nswer\s*[:：]\s*\$?(-?\d+)") |
| INTS = re.compile(r"-?\d+") |
|
|
|
|
| def _safe_int(s): |
| """AIME 答案恒为 0-999。超长数字是模型退化的产物(见过 28000 位), |
| 直接 int() 会撞 Python 3.11 的 int_max_str_digits=4300 抛 ValueError。""" |
| if s is None: |
| return None |
| t = s.lstrip("-") |
| if len(t) > 10 or not t: |
| return None |
| try: |
| return int(s) |
| except ValueError: |
| return None |
|
|
|
|
| def extract(text): |
| m = BOXED.findall(text) |
| if m: |
| digits = INTS.findall(m[-1]) |
| if digits: |
| v = _safe_int(digits[-1]) |
| if v is not None: |
| return v |
| m = ANSW.findall(text) |
| if m: |
| v = _safe_int(m[-1]) |
| if v is not None: |
| return v |
| m = INTS.findall(text) |
| if m: |
| return _safe_int(m[-1]) |
| return None |
|
|
|
|
| def _kill_descendants(): |
| """递归杀掉本进程的所有子孙(vLLM worker 是独立进程, 不杀会变孤儿占满卡)。""" |
| import signal |
| import subprocess |
| me = os.getpid() |
| try: |
| out = subprocess.run(["ps", "-eo", "pid,ppid", "--no-headers"], |
| capture_output=True, text=True, timeout=20).stdout |
| except Exception: |
| return |
| children = {} |
| for ln in out.splitlines(): |
| parts = ln.split() |
| if len(parts) != 2: |
| continue |
| try: |
| pid, ppid = int(parts[0]), int(parts[1]) |
| except ValueError: |
| continue |
| children.setdefault(ppid, []).append(pid) |
| doomed, stack = [], [me] |
| while stack: |
| for c in children.get(stack.pop(), []): |
| doomed.append(c) |
| stack.append(c) |
| for pid in reversed(doomed): |
| try: |
| os.kill(pid, signal.SIGKILL) |
| except OSError: |
| pass |
|
|
|
|
| def _hard_exit(code=0): |
| sys.stdout.flush() |
| sys.stderr.flush() |
| _kill_descendants() |
| os._exit(code) |
|
|
|
|
| def _die(exc_type, exc, tb): |
| import traceback as _tb |
| _tb.print_exception(exc_type, exc, tb) |
| _hard_exit(1) |
|
|
|
|
| sys.excepthook = _die |
|
|
| rows = [json.loads(l) for l in open(a.data) if l.strip()] |
| log(f"{len(rows)} 题 x n={a.n}, enable_thinking={a.enable_thinking}") |
|
|
| from transformers import AutoTokenizer |
| from vllm import LLM, SamplingParams |
|
|
| tok = AutoTokenizer.from_pretrained(a.model, trust_remote_code=True) |
| prompts = [] |
| for r in rows: |
| msgs = [{"role": "user", "content": r["problem"].strip() + "\n\n" + INSTR}] |
| prompts.append(tok.apply_chat_template( |
| msgs, add_generation_prompt=True, tokenize=False, |
| enable_thinking=a.enable_thinking)) |
| log("prompt 尾部抽检: " + repr(prompts[0][-60:])) |
|
|
| def _fit_tp(model_path, tp): |
| """attn heads 必须能被 TP 整除, 否则 vLLM 起不来(R1-Distill-1.5B 是 12 heads)。""" |
| try: |
| cfg = json.load(open(os.path.join(model_path, "config.json"))) |
| h = cfg.get("num_attention_heads") |
| except Exception: |
| return tp |
| if not h or h % tp == 0: |
| return tp |
| for cand in range(tp, 0, -1): |
| if h % cand == 0: |
| log(f"注意: {h} 个 attention head 不能被 TP={tp} 整除, 自动降为 TP={cand}") |
| return cand |
| return 1 |
|
|
|
|
| a.tp = _fit_tp(a.model, a.tp) |
| llm = LLM(model=a.model, tensor_parallel_size=a.tp, |
| gpu_memory_utilization=a.gpu_util, max_model_len=a.ctx, |
| trust_remote_code=True, seed=a.seed) |
| |
| |
| sp = SamplingParams(n=a.n, temperature=a.temp, top_p=a.top_p, top_k=a.top_k, |
| max_tokens=a.max_new) |
|
|
| t0 = time.time() |
| outs = llm.generate(prompts, sp) |
| log(f"生成 {time.time()-t0:.0f}s") |
|
|
| recs, n_correct, n_tot, n_trunc, lens = [], 0, 0, 0, [] |
| per_problem = [] |
| for r, out in zip(rows, outs): |
| gt = int(str(r["answer"]).strip()) |
| hits, samples = 0, [] |
| for o in out.outputs: |
| pred = extract(o.text) |
| ok = (pred is not None and pred == gt) |
| hits += ok |
| n_correct += ok |
| n_tot += 1 |
| n_trunc += (o.finish_reason == "length") |
| lens.append(len(o.token_ids)) |
| samples.append({"pred": pred, "ok": bool(ok), |
| "finish": o.finish_reason, "len": len(o.token_ids), |
| "text": o.text}) |
| per_problem.append({"id": r.get("id"), "gt": gt, "hits": hits, "n": len(out.outputs)}) |
| recs.append({"id": r.get("id"), "problem": r["problem"], "gt": gt, |
| "hits": hits, "samples": samples}) |
|
|
| os.makedirs(os.path.dirname(os.path.abspath(a.out)) or ".", exist_ok=True) |
| with open(a.out, "w") as f: |
| for rec in recs: |
| f.write(json.dumps(rec, ensure_ascii=False) + "\n") |
|
|
| avg = n_correct / n_tot |
| solved_any = sum(1 for x in per_problem if x["hits"] > 0) |
| solved_all = sum(1 for x in per_problem if x["hits"] == x["n"]) |
| summary = { |
| "model": a.model, "data": a.data, "n": a.n, |
| "enable_thinking": a.enable_thinking, |
| "temperature": a.temp, "top_p": a.top_p, "top_k": a.top_k, |
| "max_new": a.max_new, "seed": a.seed, |
| "problems": len(rows), "samples": n_tot, |
| f"avg@{a.n}": avg, |
| "solved_any": solved_any, "solved_all": solved_all, |
| "truncated_rate": n_trunc / n_tot, |
| "resp_len_mean": sum(lens) / len(lens), |
| "resp_len_max": max(lens), |
| "hits_hist": dict(Counter(x["hits"] for x in per_problem)), |
| } |
| with open(a.out.replace(".jsonl", "_summary.json"), "w") as f: |
| json.dump(summary, f, indent=2, ensure_ascii=False) |
|
|
| print() |
| print(json.dumps(summary, indent=2, ensure_ascii=False), flush=True) |
| print() |
| print(f"AIME25 avg@{a.n} = {avg:.4f} ({n_correct}/{n_tot})", flush=True) |
| print(f" 解出>=1次 {solved_any}/{len(rows)}, 全对 {solved_all}/{len(rows)}", flush=True) |
| print(f" 截断率 {n_trunc/n_tot:.3f}, 平均长度 {sum(lens)/len(lens):.0f}", flush=True) |
|
|
| |
| |
| _hard_exit(0) |
|
|
