| |
| """ |
| peer_run.py -- generate completions for peer-model NIAH prompts. |
| |
| Backends, tried in this order unless --backend forces one: |
| vllm fast, batched. Cannot load Spark-X2.5 (model_type spark2_5 is |
| custom_code and has no vLLM backend), so it is skipped automatically. |
| hf transformers with trust_remote_code. Works for everything, slower. |
| |
| Prompts are fed as TOKEN IDS, never text: peer_gen.py guarantees an exact token |
| count and a decode->re-encode round trip would silently break it. |
| |
| Greedy decoding, matching our own harness (CP_GEN_GREEDY=1). Generation stops at |
| --max-new tokens; no early stop, because the peer models have no shared |
| end-of-answer convention and truncation is harmless -- scoring is substring |
| match, so trailing chatter costs nothing. |
| |
| usage: |
| peer_run.py --model Qwen/Qwen3-4B-Base --prompts out/niah_single_prompts.jsonl \ |
| --out gen/qwen3-4b_niah_single.jsonl --max-new 24 |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
|
|
| def load_prompts(p): |
| rows = [] |
| with open(p) as f: |
| for line in f: |
| line = line.strip() |
| if line: |
| rows.append(json.loads(line)) |
| return rows |
|
|
|
|
| def run_vllm(model, rows, max_new, gpu_frac, max_len): |
| from vllm import LLM, SamplingParams |
| llm = LLM(model=model, dtype="bfloat16", trust_remote_code=True, |
| gpu_memory_utilization=gpu_frac, max_model_len=max_len, |
| enforce_eager=False) |
| sp = SamplingParams(temperature=0.0, max_tokens=max_new) |
| outs = llm.generate(prompt_token_ids=[r["ids"] for r in rows], |
| sampling_params=sp) |
| return [o.outputs[0].text for o in outs] |
|
|
|
|
| def run_hf(model, rows, max_new, batch): |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
| tok = AutoTokenizer.from_pretrained(model, trust_remote_code=True) |
| m = AutoModelForCausalLM.from_pretrained( |
| model, dtype=torch.bfloat16, device_map="cuda", trust_remote_code=True) |
| m.eval() |
| texts = [] |
| |
| |
| for s in range(0, len(rows), batch): |
| chunk = rows[s:s + batch] |
| ids = torch.tensor([r["ids"] for r in chunk], device="cuda") |
| with torch.no_grad(): |
| out = m.generate(input_ids=ids, max_new_tokens=max_new, |
| do_sample=False, num_beams=1, |
| pad_token_id=getattr(tok, "pad_token_id", None) |
| or getattr(tok, "eos_token_id", 0)) |
| for j in range(len(chunk)): |
| texts.append(tok.decode(out[j][ids.shape[1]:], skip_special_tokens=True)) |
| done = min(s + batch, len(rows)) |
| print(f" {done}/{len(rows)}", flush=True) |
| return texts |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--prompts", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--max-new", type=int, default=24) |
| ap.add_argument("--backend", default="auto", choices=["auto", "vllm", "hf"]) |
| ap.add_argument("--gpu-frac", type=float, default=0.90) |
| ap.add_argument("--batch", type=int, default=1, help="hf backend only") |
| args = ap.parse_args() |
|
|
| rows = load_prompts(args.prompts) |
| if not rows: |
| sys.exit(f"no prompts in {args.prompts}") |
| plen = len(rows[0]["ids"]) |
| max_len = plen + args.max_new + 8 |
| |
| |
| |
| try: |
| from transformers import AutoConfig |
| mp = getattr(AutoConfig.from_pretrained(args.model, trust_remote_code=True), |
| "max_position_embeddings", None) |
| if mp and max_len > mp: |
| print(f" clamping max_model_len {max_len} -> {mp} (model's native window)") |
| max_len = mp |
| if plen + args.max_new > mp: |
| sys.exit(f" FATAL: prompt {plen} + max_new {args.max_new} exceeds " |
| f"the model's window {mp}; lower --prompt-len") |
| except Exception as e: |
| print(f" (could not read max_position_embeddings: {type(e).__name__})") |
| print(f" {args.model} {len(rows)} prompts x {plen} tok max_new={args.max_new}") |
|
|
| backend = args.backend |
| if backend == "auto": |
| backend = "hf" |
| try: |
| import vllm |
| |
| backend = "hf" if "spark" in args.model.lower() else "vllm" |
| except ImportError: |
| pass |
| print(f" backend: {backend}") |
|
|
| t0 = time.time() |
| if backend == "vllm": |
| try: |
| texts = run_vllm(args.model, rows, args.max_new, args.gpu_frac, max_len) |
| except Exception as e: |
| print(f" vLLM failed ({type(e).__name__}: {str(e)[:140]}) -- falling back to hf") |
| texts = run_hf(args.model, rows, args.max_new, args.batch) |
| else: |
| texts = run_hf(args.model, rows, args.max_new, args.batch) |
| dt = time.time() - t0 |
|
|
| os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) |
| with open(args.out, "w") as f: |
| for r, t in zip(rows, texts): |
| f.write(json.dumps({"i": r["i"], "generated": t}) + "\n") |
| print(f" wrote {args.out} ({len(texts)} rows, {dt:.0f}s, " |
| f"{dt/max(1,len(texts)):.2f}s/prompt)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
