скрипт заміру швидкості

scripts/bench_tps.py

@@ -0,0 +1,128 @@
+"""Швидкість декоду для трьох шляхів: transformers / onnxruntime / llama.cpp.
+
+Міряємо ОДНЕ й те саме: генерацію одного рядка за раз (так модель і працює —
+речення за реченням), greedy, без препроцесингу. Рахуємо згенеровані токени,
+поділені на час саме генерації.
+
+    python bench_tps.py <тека_моделі> [--gguf-url http://127.0.0.1:8903]
+
+llama.cpp не міряємо самі — беремо його власні timings, вони точніші за наш
+секундомір навколо HTTP.
+"""
+import argparse
+import json
+import time
+import urllib.request
+
+PROMPTS = [
+    "Зустріч о 08:05 у кабінеті 12.",
+    "Збитки оцінили в 7 000 000 000 доларів.",
+    "Пиши на support@shop.com.ua або дзвони 067 123 45 67.",
+    "Борг сягнув 2 625 414 042 726 гривень.",
+    "Марафон стартує 12 жовтня о 8:15, дистанції 5, 10 і 21,1 км.",
+    "У специфікації вказано роз'єм тип C та частоту 60 Гц.",
+]
+WARMUP = 2
+
+
+def run_hf(model_dir, device):
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    tok = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_dir, dtype=torch.bfloat16 if device == "cuda" else torch.float32)
+    model = model.to(device).eval()
+
+    def gen(text):
+        ids = tok.apply_chat_template([{"role": "user", "content": text}],
+                                      add_generation_prompt=True,
+                                      return_tensors="pt").to(device)
+        t0 = time.perf_counter()
+        with torch.inference_mode():
+            out = model.generate(ids, max_new_tokens=192, do_sample=False)
+        if device == "cuda":
+            torch.cuda.synchronize()
+        return out.shape[-1] - ids.shape[-1], time.perf_counter() - t0
+
+    return gen
+
+
+def run_onnx(model_dir, provider="CPUExecutionProvider"):
+    from optimum.onnxruntime import ORTModelForCausalLM
+    from transformers import AutoTokenizer
+    import onnxruntime as ort
+    tok = AutoTokenizer.from_pretrained(model_dir)
+    # без явного provider optimum мовчки бере CPU, навіть коли CUDA доступна.
+    # io-binding на CUDA для цього графа падає (cudaErrorIllegalAddress на
+    # embed_tokens/Mul), тож на GPU його вимикаємо.
+    kw = {"provider": provider}
+    if provider != "CPUExecutionProvider":
+        kw["use_io_binding"] = False
+    model = ORTModelForCausalLM.from_pretrained(model_dir, subfolder="onnx", **kw)
+    prov = model.model.get_providers()[0] if hasattr(model, "model") else "?"
+    print(f"    onnxruntime провайдер: {prov} "
+          f"(доступні: {ort.get_available_providers()})")
+
+    def gen(text):
+        ids = tok.apply_chat_template([{"role": "user", "content": text}],
+                                      add_generation_prompt=True,
+                                      return_tensors="pt")
+        t0 = time.perf_counter()
+        out = model.generate(ids, max_new_tokens=192, do_sample=False)
+        return out.shape[-1] - ids.shape[-1], time.perf_counter() - t0
+
+    return gen
+
+
+def bench(name, gen):
+    for p in PROMPTS[:WARMUP]:
+        gen(p)
+    ntok = secs = 0
+    for p in PROMPTS:
+        n, s = gen(p)
+        ntok += n
+        secs += s
+    print(f"  {name:24} {ntok/secs:7.1f} ток/с   "
+          f"({ntok} токенів за {secs:.2f} с)")
+    return ntok / secs
+
+
+def bench_gguf(url):
+    ntok = secs = 0
+    for p in PROMPTS:
+        body = {"messages": [{"role": "user", "content": p}],
+                "temperature": 0, "top_k": 1, "n_predict": 192,
+                "cache_prompt": False}
+        req = urllib.request.Request(f"{url}/v1/chat/completions",
+                                     data=json.dumps(body, ensure_ascii=False).encode(),
+                                     headers={"Content-Type": "application/json"})
+        with urllib.request.urlopen(req, timeout=120) as r:
+            d = json.loads(r.read())
+        t = d.get("timings") or d.get("usage", {})
+        ntok += t.get("predicted_n", 0)
+        secs += t.get("predicted_ms", 0) / 1000
+    print(f"  {'llama.cpp (GGUF F16)':24} {ntok/secs:7.1f} ток/с   "
+          f"({ntok} токенів за {secs:.2f} с)")
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser()
+    ap.add_argument("model_dir")
+    ap.add_argument("--gguf-url", default=None)
+    ap.add_argument("--skip", default="", help="через кому: hf_cuda,hf_cpu,onnx")
+    args = ap.parse_args()
+    skip = set(args.skip.split(","))
+
+    print("\n=== швидкість декоду, один рядок за раз, greedy ===")
+    if "hf_cuda" not in skip:
+        bench("transformers (bf16, GPU)", run_hf(args.model_dir, "cuda"))
+    if "onnx_gpu" not in skip:
+        bench("onnxruntime (fp32, GPU)",
+              run_onnx(args.model_dir, "CUDAExecutionProvider"))
+    if "onnx" not in skip:
+        bench("onnxruntime (fp32, CPU)", run_onnx(args.model_dir))
+    if "hf_cpu" not in skip:
+        bench("transformers (fp32, CPU)", run_hf(args.model_dir, "cpu"))
+    if args.gguf_url:
+        bench_gguf(args.gguf_url)
+    print("BENCH_DONE")