| """HF-stack decode speed probe (go/no-go for the HF port of our method). |
| Loads one model (bf16 / compressed-tensors / GPTQ -- auto-detected from the |
| checkpoint), generates greedily on image prompts, reports decode tok/s with |
| prefill subtracted via a 1-token call. transformers-native generate, KV cache. |
| """ |
| import argparse, io, json, statistics, time |
| import torch |
| from PIL import Image |
| import pyarrow.parquet as pq |
|
|
| from stage0_gonogo import build_inputs |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--tok-model", default="") |
| ap.add_argument("--parquet", required=True) |
| ap.add_argument("--pids-from", required=True) |
| ap.add_argument("--n", type=int, default=3) |
| ap.add_argument("--max-new", type=int, default=256) |
| ap.add_argument("--img-max-side", type=int, default=1024) |
| ap.add_argument("--bf16", action="store_true", help="force bf16 dtype (unquantized)") |
| args = ap.parse_args() |
|
|
| from transformers import AutoProcessor, AutoModelForImageTextToText |
| proc = AutoProcessor.from_pretrained(args.tok_model or args.model) |
| kw = dict(device_map="cuda:0") |
| if args.bf16: |
| kw["dtype"] = torch.bfloat16 |
| try: |
| model = AutoModelForImageTextToText.from_pretrained(args.model, **kw) |
| except TypeError: |
| kw.pop("dtype", None) |
| if args.bf16: |
| kw["torch_dtype"] = torch.bfloat16 |
| model = AutoModelForImageTextToText.from_pretrained(args.model, **kw) |
| model.eval() |
| print(f"[init] {args.model} loaded, mem={torch.cuda.memory_allocated()/2**30:.1f}GiB", |
| flush=True) |
|
|
| pids = [] |
| for line in open(args.pids_from): |
| d = json.loads(line) |
| if "pid" in d: |
| pids.append(d["pid"]) |
| if len(pids) >= args.n: |
| break |
| tbl = pq.read_table(args.parquet, columns=["pid", "query", "decoded_image"]) |
| rowmap = {r["pid"]: r for r in tbl.to_pylist()} |
|
|
| speeds = [] |
| for pid in pids: |
| row = rowmap[pid] |
| img = Image.open(io.BytesIO(row["decoded_image"]["bytes"])).convert("RGB") |
| if max(img.size) > args.img_max_side: |
| sc = args.img_max_side / max(img.size) |
| img = img.resize((int(img.width * sc), int(img.height * sc))) |
| enc = build_inputs(proc, row["query"], img) |
| enc = {k: (v.to("cuda:0") if hasattr(v, "to") else v) for k, v in enc.items()} |
| with torch.inference_mode(): |
| model.generate(**enc, max_new_tokens=4, do_sample=False) |
| torch.cuda.synchronize() |
| t0 = time.time() |
| model.generate(**enc, max_new_tokens=1, do_sample=False) |
| torch.cuda.synchronize() |
| t_pref = time.time() - t0 |
| t0 = time.time() |
| out = model.generate(**enc, max_new_tokens=args.max_new, do_sample=False) |
| torch.cuda.synchronize() |
| t_full = time.time() - t0 |
| ntok = out.shape[1] - enc["input_ids"].shape[1] |
| dec = (ntok - 1) / max(t_full - t_pref, 1e-6) |
| speeds.append(dec) |
| print(f"[{pid}] gen={ntok} prefill={t_pref:.2f}s decode={dec:.1f} tok/s", flush=True) |
| print(f"[agg] model={args.model} n={len(speeds)} " |
| f"decode median={statistics.median(speeds):.1f} mean={statistics.mean(speeds):.1f} tok/s", |
| flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
