artifacts/dixie-flatline/a40_quickbench.py

#!/usr/bin/env python
"""A40 quick-bench for fast quality A/B of gemma-4 variants (transformers, batched).

AIME 2024+2025 (n=60) — the one axis that showed signal, hard reasoning, ROBUST
numeric scorer (integer exact-match, NO model-grader → dodges the inspect MATH bug).
Runs entirely on the A40 (transformers 5.9 + torch 2.8cu128); capability is
hardware/engine-invariant so it composes with the a10g baselines. ~10 min/variant.

Usage: python3 a40_quickbench.py --model /path/or/hf-id [--n 60 --batch 16 --max-new 4096]
"""
from __future__ import annotations
import argparse, json, re, time, torch
from transformers import AutoTokenizer, AutoModelForCausalLM, LogitsProcessor
from datasets import load_dataset


class KeepsetMask(LogitsProcessor):
    """Replicate the cleanstack 128k head-prune at GENERATION time: mask every
    non-keepset token to -inf so it can never be emitted. Lets us isolate the
    head-prune effect on ANY base model (e.g. bf16 ± prune) — codex's clean
    causal separator, no re-quantization needed."""
    def __init__(self, keep_ids, vocab_size):
        block = torch.ones(vocab_size, dtype=torch.bool)
        block[torch.tensor(sorted(set(int(i) for i in keep_ids)))] = False
        self.block = block

    def __call__(self, input_ids, scores):
        b = self.block.to(scores.device)
        if b.shape[0] != scores.shape[-1]:  # pad/trim to actual logit width
            nb = torch.ones(scores.shape[-1], dtype=torch.bool, device=scores.device)
            nb[: b.shape[0]] = b[: scores.shape[-1]] if b.shape[0] >= scores.shape[-1] else b
            b = nb
        scores[:, b] = float("-inf")
        return scores


def load_keepset(path):
    d = json.load(open(path))
    if isinstance(d, dict):
        for k in ("keep", "ids", "keepset", "token_ids", "keep_ids"):
            if k in d:
                return d[k]
        # dict of {token: id} or {id: ...}
        return [int(x) for x in d.keys()] if all(str(x).lstrip("-").isdigit() for x in d.keys()) else list(d.values())
    return d


def extract_answer(text: str):
    """AIME answers are integers 0-999. Prefer \\boxed{...}, else last integer."""
    boxed = re.findall(r"\\boxed\{\s*(-?\d+)\s*\}", text)
    if boxed:
        try: return int(boxed[-1]) % 1000
        except ValueError: pass
    nums = re.findall(r"-?\d+", text.replace(",", ""))
    return (int(nums[-1]) % 1000) if nums else None


AIME_SETS = {
    # original 60: AIME 2024 (I+II) + 2025 (I+II)
    "core": [("Maxwell-Jia/AIME_2024", "train", "Problem", "Answer"),
             ("math-ai/aime25", "test", "problem", "answer")],
    # ~120: AIMO validation (2022+2023+2024, 90) + 2025 (30). Drops Maxwell-Jia 2024 to
    # avoid duplicating the AIMO 2024 problems. Larger n → tighter CI for the paired test.
    "expanded": [("AI-MO/aimo-validation-aime", "train", "problem", "answer"),
                 ("math-ai/aime25", "test", "problem", "answer")],
    # ~963: every AIME 1983-2024 (933) + 2025 (30). 8x n. NOTE: old years (pre-~2010) are
    # likely memorized → concordant pairs that add n but little paired POWER; the recent
    # years carry the discriminating signal. Still a strict superset, free to run.
    "full": [("di-zhang-fdu/AIME_1983_2024", "train", "Question", "Answer"),
             ("math-ai/aime25", "test", "problem", "answer")],
}


def load_aime(n: int, which: str = "core"):
    items = []
    for repo, split, pk, ak in AIME_SETS[which]:
        ds = load_dataset(repo, split=split)
        for d in ds:
            try:                                    # AIME answers are ints 0-999; skip malformed rows
                items.append((d[pk], int(str(d[ak]).strip()) % 1000))
            except (ValueError, TypeError):
                continue
    return items if n <= 0 else items[:n]


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--n", type=int, default=60)
    ap.add_argument("--batch", type=int, default=16)
    ap.add_argument("--max-new", type=int, default=4096)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--label", default="")
    ap.add_argument("--aime-set", default="core", choices=list(AIME_SETS), help="core(60)|expanded(~120)")
    ap.add_argument("--dump", default="", help="write per-problem {idx,answer,pred,correct} JSON for paired/McNemar analysis")
    ap.add_argument("--check", action="store_true", help="load the AIME set, print count+samples+answer-extraction, exit (no model)")
    ap.add_argument("--keepset", default="", help="path to pck04_keepset.json → apply 128k head-prune mask")
    ap.add_argument("--greedy", action="store_true", help="deterministic decode (clean A/B isolation, no sampling noise)")
    ap.add_argument("--gptq", action="store_true",
                    help="load a GPTQModel-saved checkpoint (honors mixed/dynamic bits; AutoModel's GPTQ path may not)")
    args = ap.parse_args()
    torch.manual_seed(args.seed)

    if args.check:  # validate-first: confirm the AIME set loads + answer extraction works, no model
        items = load_aime(args.n, args.aime_set)
        print(f"[check] aime-set={args.aime_set} loaded n={len(items)}", flush=True)
        for j in (0, len(items) // 2, len(items) - 1):
            p, a = items[j]
            print(f"[check] idx={j} ans={a} prompt[:90]={p[:90]!r}", flush=True)
        return

    tok = AutoTokenizer.from_pretrained(args.model)
    if tok.pad_token_id is None:
        tok.pad_token = tok.eos_token
    tok.padding_side = "left"
    t0 = time.time()
    if args.gptq:
        # GPTQModel.load honors per-module `dynamic` bit configs; AutoModel's GPTQ loader
        # may silently ignore them. Wrapper proxies .generate()/.eval(); underlying HF
        # model is .model (for get_output_embeddings).
        from gptqmodel import GPTQModel
        model = GPTQModel.load(args.model)
    else:
        model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16,
                                                     device_map="cuda", attn_implementation="sdpa")
    model.eval()
    print(f"[bench] loaded {model.__class__.__name__} in {time.time()-t0:.1f}s", flush=True)

    logits_processor = None
    if args.keepset:
        keep = load_keepset(args.keepset)
        vocab = model.get_output_embeddings().weight.shape[0]
        logits_processor = [KeepsetMask(keep, vocab)]
        print(f"[bench] keepset-mask ON: {len(set(int(i) for i in keep))} kept / {vocab} vocab", flush=True)

    items = load_aime(args.n, args.aime_set)
    per_problem = []  # paired-analysis record: one row per problem, in fixed order
    correct, total, t0 = 0, 0, time.time()
    for i in range(0, len(items), args.batch):
        chunk = items[i:i + args.batch]
        prompts = [tok.apply_chat_template([{"role": "user", "content": p}],
                   enable_thinking=True, add_generation_prompt=True, tokenize=False)
                   for p, _ in chunk]
        enc = tok(prompts, return_tensors="pt", padding=True, add_special_tokens=False).to("cuda")
        gen_kw = dict(max_new_tokens=args.max_new, pad_token_id=tok.pad_token_id,
                      logits_processor=logits_processor)
        if args.greedy:
            gen_kw.update(do_sample=False)
        else:
            gen_kw.update(do_sample=True, temperature=1.0, top_p=0.95, top_k=64)
        with torch.no_grad():
            out = model.generate(**enc, **gen_kw)
        gens = tok.batch_decode(out[:, enc.input_ids.shape[1]:], skip_special_tokens=True)
        for (_, ans), g in zip(chunk, gens):
            pred = extract_answer(g)
            ok = pred is not None and pred == ans % 1000
            per_problem.append({"idx": total, "answer": ans % 1000, "pred": pred, "correct": int(ok)})
            correct += int(ok); total += 1
        print(f"[bench] {total}/{len(items)} running_acc={correct/total:.3f} "
              f"({time.time()-t0:.0f}s)", flush=True)

    if args.dump:
        json.dump({"label": args.label or args.model, "aime_set": args.aime_set,
                   "greedy": bool(args.greedy), "per_problem": per_problem}, open(args.dump, "w"))
        print(f"[bench] dumped per-problem → {args.dump}", flush=True)

    acc = correct / total
    se = (acc * (1 - acc) / total) ** 0.5
    print(f"=====QUICKBENCH=====", flush=True)
    print(f"model={args.label or args.model} AIME24+25 n={total} acc={acc:.4f} "
          f"stderr={se:.4f} wall={time.time()-t0:.0f}s", flush=True)


if __name__ == "__main__":
    main()