File size: 5,074 Bytes
5c0a4a8 | 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 | """Full-model parity and measurements on this Mac; emits no simulated hardware results."""
import argparse
import json
import time
from pathlib import Path
import mlx.core as mx
import numpy as np
from solomon_mlx import Solomon
from solomon_mlx._vendor.semantics import listed_probs, p_yes
from solomon_mlx.artifacts import digest, sha256
def probabilities(job, row):
if job["task"] in ("boolean", "multilabel", "entity"):
p = p_yes(row["letter_logits"])
return np.array([p, 1 - p])
n = job["n"] - 2 if job["head_key"].endswith("choiceR") else job["n"]
return listed_probs(row["letter_logits"], n)
def run(model_dir, jobs_path, reference_path, output):
output = Path(output)
if output.exists():
raise FileExistsError("Benchmark outputs are immutable")
jobs = json.loads(Path(jobs_path).read_text())
reference = json.loads(Path(reference_path).read_text())
ref = {r["id"]: r for r in reference["rows"]}
if set(ref) != {j["id"] for j in jobs}:
raise ValueError("Benchmark and reference jobs differ")
mx.reset_peak_memory()
started = time.perf_counter()
model = Solomon.load(model_dir)
load_seconds = time.perf_counter() - started
states = {}
rows = []
prefills = []
try:
for job in jobs:
key = digest(job["parts"])
if key not in states:
state = model.prefill(job["parts"])
states[key] = state
prefills.append(
{
"document": key,
"tokens": state.prefix_tokens,
"seconds": state._data["prefill_seconds"],
"vision_seconds": state._data["vision_seconds"],
"cache_bytes": sum(c.nbytes for c in state._data["cache"]),
}
)
state = states[key]
row = model.engine.ask(
state._data,
job["block"],
job["n"],
job["head_key"],
execution=job.get("execution", "cached"),
taps=job.get("taps", []),
)
p, q = probabilities(job, row), probabilities(job, ref[job["id"]])
row.update(
id=job["id"],
decision_agrees=bool(p.argmax() == q.argmax()),
max_probability_drift=float(np.max(np.abs(p - q))),
max_logit_drift=float(
np.max(np.abs(np.array(row["letter_logits"]) - ref[job["id"]]["letter_logits"]))
),
prefix_ids_exact=state._data["prefix_ids"] == ref[job["id"]]["prefix_ids"],
)
if "token_ids" in row:
row["token_ids_exact"] = row["token_ids"] == ref[job["id"]]["token_ids"]
if row.get("taps"):
row["layer_max_hidden_drift"] = {
k: float(np.max(np.abs(np.array(v) - ref[job["id"]]["taps"][k])))
for k, v in row["taps"].items()
}
rows.append(row)
print(job["id"], row["seconds"], row["decision_agrees"], flush=True)
# Replay checks real prefill and repeated question semantics on the same binding.
first = states[digest(jobs[0]["parts"])]
recipe = output.with_suffix(".replay.json")
first.save(recipe)
with model.replay(recipe) as restored:
b, n, h = jobs[0]["block"], jobs[0]["n"], jobs[0]["head_key"]
replay = model.engine.ask(restored._data, b, n, h)
replay_drift = float(np.max(np.abs(np.array(replay["letter_logits"]) - rows[0]["letter_logits"])))
finally:
for state in states.values():
state.close()
warm = [r["seconds"] for r in rows if r["reused_prefix_tokens"]]
report = {
"runtime": model.identity,
"device": mx.device_info(),
"jobs_sha256": sha256(jobs_path),
"reference_sha256": sha256(reference_path),
"load_seconds": load_seconds,
"prefills": prefills,
"rows": rows,
"warm_question_latency_median_seconds": float(np.median(warm)),
"questions_per_second": len(warm) / sum(warm),
"peak_metal_bytes": mx.get_peak_memory(),
"replay_max_logit_drift": replay_drift,
"decision_agreement": float(np.mean([r["decision_agrees"] for r in rows])),
"scope": "development parity fixtures; not held-out task accuracy or release qualification",
"calibration_status": "uncalibrated",
}
output.write_text(json.dumps(report, indent=2))
return report
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--model", default="models/quality")
p.add_argument("--jobs", default="evaluations/golden-jobs.json")
p.add_argument("--reference", default="evaluations/bf16-reference-1789901869/report.json")
p.add_argument("--output", default="evaluations/bf16-text-benchmark.json")
a = p.parse_args()
run(a.model, a.jobs, a.reference, a.output)
|