File size: 3,363 Bytes
ee3e28a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | """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()
|