File size: 3,943 Bytes
c429240 | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """Measure what changing the number of active MoE experts actually does.
Run against any MoE GGUF. It loads the model once per expert count, times a
fixed generation, and reports throughput. The point is to replace a guess with
a number before touching a serving configuration.
"""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
from llama_cpp import Llama
PROMPT = "Write a Python function that reverses a string."
def measure(
model_path: str,
expert_count: int | None,
tokens: int,
threads: int,
) -> dict:
"""Load the model with an optional expert override and time a generation."""
key = "expert_used_count"
overrides = None
if expert_count is not None:
overrides = {key: expert_count}
kwargs = {"kv_overrides": overrides} if overrides else {}
started = time.monotonic()
try:
model = Llama(
model_path=model_path,
n_ctx=512,
n_threads=threads,
verbose=False,
**kwargs,
)
except Exception as exc:
return {"experts": expert_count, "error": f"{type(exc).__name__}: {exc}"}
load_s = time.monotonic() - started
arch = model.metadata.get("general.architecture", "")
native = model.metadata.get(f"{arch}.{key}")
started = time.monotonic()
out = model.create_chat_completion(
messages=[{"role": "user", "content": PROMPT}],
max_tokens=tokens,
temperature=0.0,
)
gen_s = time.monotonic() - started
produced = out.get("usage", {}).get("completion_tokens", 0)
text = out["choices"][0]["message"]["content"]
del model
return {
"experts": expert_count,
"native_experts": native,
"load_s": round(load_s, 2),
"gen_s": round(gen_s, 2),
"tokens": produced,
"tok_per_s": round(produced / gen_s, 1) if gen_s else None,
"preview": text[:120],
}
def main() -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("model", help="path to a MoE GGUF file")
parser.add_argument(
"--experts",
default="2,8,15",
help="comma-separated expert counts to try (omit for native only)",
)
parser.add_argument("--tokens", type=int, default=128)
parser.add_argument("--threads", type=int, default=4)
parser.add_argument("--json", help="write the results to this JSON file")
args = parser.parse_args()
if not Path(args.model).is_file():
parser.error(f"model not found: {args.model}")
counts: list[int | None] = [None]
counts += [int(c) for c in args.experts.split(",") if c.strip()]
results = []
for count in counts:
label = "native" if count is None else f"{count} experts"
print(f"--- {label}", flush=True)
row = measure(args.model, count, args.tokens, args.threads)
results.append(row)
if "error" in row:
print(f" FAILED: {row['error']}", flush=True)
else:
print(
f" {row['tok_per_s']} tok/s "
f"(load {row['load_s']}s, gen {row['gen_s']}s)",
flush=True,
)
baseline = next((r for r in results if r.get("experts") is None), None)
if baseline and baseline.get("tok_per_s"):
print("\nrelative to native:", flush=True)
for row in results:
if row.get("tok_per_s"):
ratio = row["tok_per_s"] / baseline["tok_per_s"]
print(
f" {str(row['experts']):>4} experts -> x{ratio:.2f} speed",
flush=True,
)
if args.json:
Path(args.json).write_text(
json.dumps(results, indent=2, ensure_ascii=False), encoding="utf-8"
)
print(f"\nwritten to {args.json}", flush=True)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|