106 lines
4.1 KiB
Python
106 lines
4.1 KiB
Python
|
"""Greedy-decode parity for llama.cpp against the same fp32 reference vLLM is
|
||
|
|
scored on.
|
||
|
|
|
||
|
|
The point of the comparison is that "4-bit" is not one thing: llama.cpp's
|
||
|
|
Q4_K_M and our `cpu_int4` spend a similar number of bits per weight but pick
|
||
|
|
different grids, so the speed comparison is only meaningful next to a quality
|
||
|
|
comparison run on identical prompts. llama-server takes a prompt as raw token
|
||
|
|
ids, so both engines see exactly the token sequence the reference did.
|
||
|
|
"""
|
||
|
|
import argparse, json, subprocess, time, urllib.request
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
def post(url, payload, timeout=600):
|
||
|
|
req = urllib.request.Request(
|
||
|
|
url, data=json.dumps(payload).encode(),
|
||
|
|
headers={"Content-Type": "application/json"})
|
||
|
|
with urllib.request.urlopen(req, timeout=timeout) as r:
|
||
|
|
return json.loads(r.read())
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--gguf", required=True)
|
||
|
|
ap.add_argument("--ref", default="ref.json")
|
||
|
|
ap.add_argument("--ref-logits", default="ref_logits.npy")
|
||
|
|
ap.add_argument("--bin", default="/data/zhoutaichang/embedding_infer/llama.cpp/build/bin/llama-server")
|
||
|
|
ap.add_argument("--cores", default="56-87")
|
||
|
|
ap.add_argument("--threads", type=int, default=32)
|
||
|
|
ap.add_argument("--port", type=int, default=18099)
|
||
|
|
ap.add_argument("--ctx", type=int, default=4096)
|
||
|
|
ap.add_argument("--out", default="parity_llamacpp.json")
|
||
|
|
a = ap.parse_args()
|
||
|
|
|
||
|
|
ref = json.load(open(a.ref))
|
||
|
|
ref_logits = np.load(a.ref_logits)
|
||
|
|
n_new = len(ref["new_ids"])
|
||
|
|
|
||
|
|
proc = subprocess.Popen(
|
||
|
|
["taskset", "-c", a.cores, a.bin, "-m", a.gguf, "--port", str(a.port),
|
||
|
|
"-t", str(a.threads), "-c", str(a.ctx), "--no-warmup", "--host", "127.0.0.1"],
|
||
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
||
|
|
url = f"http://127.0.0.1:{a.port}"
|
||
|
|
try:
|
||
|
|
for _ in range(180):
|
||
|
|
try:
|
||
|
|
urllib.request.urlopen(url + "/health", timeout=2)
|
||
|
|
break
|
||
|
|
except Exception:
|
||
|
|
time.sleep(1)
|
||
|
|
else:
|
||
|
|
raise RuntimeError("llama-server did not become healthy")
|
||
|
|
|
||
|
|
res = post(url + "/completion", {
|
||
|
|
"prompt": ref["prompt_ids"], "n_predict": n_new,
|
||
|
|
"temperature": 0.0, "top_k": 1, "n_probs": 20,
|
||
|
|
"cache_prompt": False, "post_sampling_probs": False,
|
||
|
|
})
|
||
|
|
finally:
|
||
|
|
proc.terminate()
|
||
|
|
try:
|
||
|
|
proc.wait(timeout=30)
|
||
|
|
except subprocess.TimeoutExpired:
|
||
|
|
proc.kill()
|
||
|
|
|
||
|
|
probs = res.get("completion_probabilities") or []
|
||
|
|
ids = [p["id"] for p in probs] if probs and "id" in probs[0] else []
|
||
|
|
n = min(n_new, len(ids)) if ids else 0
|
||
|
|
match = sum(int(x == y) for x, y in zip(ref["new_ids"][:n], ids[:n]))
|
||
|
|
|
||
|
|
rep = {
|
||
|
|
"engine": "llama.cpp", "gguf": a.gguf.split("/")[-1],
|
||
|
|
"ref_text": ref["text"], "llamacpp_text": res.get("content", ""),
|
||
|
|
"greedy_token_match": f"{match}/{n}" if n else "n/a",
|
||
|
|
"exact_match": bool(n) and ref["new_ids"][:n] == ids[:n],
|
||
|
|
"ref_top1": ref["top1"],
|
||
|
|
"llamacpp_first_token": ids[0] if ids else None,
|
||
|
|
"first_token_agree": bool(ids) and ids[0] == ref["top1"],
|
||
|
|
}
|
||
|
|
if probs:
|
||
|
|
import math
|
||
|
|
top = (probs[0].get("top_logprobs") or probs[0].get("top_probs")
|
||
|
|
or probs[0].get("probs") or [])
|
||
|
|
first = {}
|
||
|
|
for e in top:
|
||
|
|
if "id" not in e:
|
||
|
|
continue
|
||
|
|
first[int(e["id"])] = (float(e["prob"]) if "prob" in e
|
||
|
|
else math.exp(float(e["logprob"])))
|
||
|
|
ref_top = np.argsort(-ref_logits)[:20]
|
||
|
|
ref_p = np.exp(ref_logits - ref_logits.max())
|
||
|
|
ref_p /= ref_p.sum()
|
||
|
|
pairs = [(int(t), float(ref_p[t]), first[int(t)]) for t in ref_top if int(t) in first]
|
||
|
|
rep["top20_prob_overlap"] = f"{len(pairs)}/20"
|
||
|
|
if pairs:
|
||
|
|
d = [abs(r - v) for _, r, v in pairs]
|
||
|
|
rep["top20_prob_max_abs_diff"] = max(d)
|
||
|
|
rep["top20_prob_mean_abs_diff"] = sum(d) / len(d)
|
||
|
|
print(json.dumps(rep, indent=2, ensure_ascii=False))
|
||
|
|
with open(a.out, "w") as f:
|
||
|
|
json.dump(rep, f, indent=2, ensure_ascii=False)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|