| """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) |
| |
| 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) |
|
|
