prain: int8/int4 experts quantized once on the GPU, rounds through vLLM's expert_map, [miss] min_weight skips experts not in VRAM; scripts: tune_moe, quality (answer NLL on identical token ids), bench --context, measure --quant

scripts/quality.py

@@ -0,0 +1,89 @@
+"""Answer quality of a served config against reference answers from another config.
+
+    python scripts/quality.py URL make REF.json    greedy reference answers (run on the exact config)
+    python scripts/quality.py URL score REF.json   how likely this server finds those answers
+
+Both work on token ids: the prompt is the chat template with its generation prompt, and `score`
+re-scores the reference answer on exactly the ids it was generated from (prompt logprobs). A
+finished assistant message renders differently from the generation prompt, so re-scoring
+through the chat API skews the logprobs. `nll` is the mean negative log-likelihood per answer
+token (the reference config scores about its own decode logprobs); `same_tokens` is how many
+leading tokens this server's own greedy answer shares with the reference.
+"""
+
+import argparse
+import json
+import urllib.request
+
+PROMPTS = [
+    "Write a Python function that returns the n-th Fibonacci number iteratively.",
+    "Explain what `git rebase -i HEAD~3` does, in three sentences.",
+    "Write a bash one-liner that counts the lines of all .py files below the current directory.",
+    "On Gentoo, how do I rebuild every package that links against a library I just updated?",
+    "What is the difference between a process and a thread? Answer in two sentences.",
+]
+CODE_PROMPT = "Here is some Python code:\n\n{code}\n\nWhat does this code do? Answer in two sentences."
+
+
+def post(url, path, body):
+    request = urllib.request.Request(url + path, json.dumps(body).encode(),
+                                     {"Content-Type": "application/json"})
+    with urllib.request.urlopen(request, timeout=3600) as response:
+        return json.load(response)
+
+
+def greedy(url, model, ids, tokens):
+    choice = post(url, "/v1/completions", {
+        "model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
+        "return_tokens_as_token_ids": True})["choices"][0]
+    return choice["text"], [int(t.split(":")[1]) for t in choice["logprobs"]["tokens"]]
+
+
+def nll(url, model, ids, answer):
+    """Mean negative log-likelihood of the answer tokens after the prompt ids."""
+    entries = post(url, "/v1/completions", {
+        "model": model, "prompt": ids + answer, "max_tokens": 1, "temperature": 0,
+        "prompt_logprobs": 0})["choices"][0]["prompt_logprobs"]
+    return -sum(entries[len(ids) + i][str(t)]["logprob"] for i, t in enumerate(answer)) / len(answer)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("url")
+    parser.add_argument("action", choices=("make", "score"))
+    parser.add_argument("ref")
+    parser.add_argument("--model", default="gemma4")
+    parser.add_argument("--tokens", type=int, default=48, help="answer tokens")
+    parser.add_argument("--code", default="prain/policy.py", help="file for the code question")
+    parser.add_argument("--code-chars", type=int, default=900)
+    args = parser.parse_args()
+    if args.action == "make":
+        with open(args.code) as f:
+            prompts = PROMPTS + [CODE_PROMPT.format(code=f.read()[: args.code_chars])]
+        refs = []
+        for prompt in prompts:
+            ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True,
+                                               "messages": [{"role": "user", "content": prompt}]})
+            ids = ids["tokens"]
+            text, answer = greedy(args.url, args.model, ids, args.tokens)
+            refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text})
+            print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True)
+        with open(args.ref, "w") as f:
+            json.dump(refs, f)
+        return
+    with open(args.ref) as f:
+        refs = json.load(f)
+    scores, same = [], []
+    for ref in refs:
+        scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
+        text, answer = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
+        same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
+                         min(len(answer), len(ref["answer"]))))
+        print(json.dumps({"prompt": ref["prompt"][:40], "nll": round(scores[-1], 3),
+                          "same_tokens": same[-1], "answer": text[:100]}), flush=True)
+    print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3), "same_tokens": same,
+                      "identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))
+
+
+if __name__ == "__main__":
+    main()