| """Isolated-process CPU/MPS parity and synchronized warm latency study. |
| |
| Run after the reference evaluation finishes for uncontended timing. Every schema/input |
| gets one untimed warm call on each device before repeated measurements. |
| """ |
|
|
| import argparse |
| import gc |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| from sealcell.config import config |
| from sealcell.model import LocalModel, schema_spec |
| from sealcell.policies import policy |
| from sealcell.serialization import render |
| from sealcell.simulator import Simulator |
| from sealcell.supervisor import Supervisor |
|
|
|
|
| def episode_parity(model): |
| reference = {} |
| source = Path("reports/runs/final-primary/decisions.jsonl") |
| if not source.exists(): |
| return {"status": "not run; CPU reference traces missing"} |
| with source.open() as f: |
| for line in f: |
| row = json.loads(line) |
| if row["episode_id"] > "final-00012": |
| break |
| if row["policy"] in {"hybrid", "direct"}: |
| reference.setdefault((row["episode_id"], row["policy"]), []).append(row) |
| settings = config("evaluation") |
| failures, agreements = [], 0 |
| with Path("reports/mps-episode-parity.jsonl").open("w") as output: |
| for (episode_id, name), cpu_trace in reference.items(): |
| index = int(episode_id.split("-")[1]) |
| s = Simulator( |
| settings["scenarios"][index % len(settings["scenarios"])], |
| settings["splits"]["final"]["seed_start"] + index, |
| settings["splits"]["final"]["recipe"], |
| settings["splits"]["final"]["text_family"], |
| ) |
| sup = Supervisor( |
| s, |
| policy( |
| name, |
| model=model, |
| issue_threshold=settings["issue_threshold"], |
| evidence_threshold=settings["evidence_threshold"], |
| filter_actions=settings.get("filter_actions", False), |
| **settings["conditions"]["primary"], |
| ), |
| ).run() |
|
|
| def signature(trace): |
| return [ |
| ( |
| r["proposal"]["action"], |
| r["executed"], |
| r["gate_rejection"], |
| r["stage_after"], |
| r["time_after"], |
| (r["proposal"]["model"] or {}).get("input_hash"), |
| { |
| h: v["selected_label"] |
| for h, v in (r["proposal"]["model"] or {}).get("heads", {}).items() |
| }, |
| ) |
| for r in trace |
| ] |
|
|
| equal = signature(cpu_trace) == signature(sup.trace) |
| agreements += equal |
| if not equal: |
| failures.append({"episode_id": episode_id, "policy": name}) |
| output.write( |
| json.dumps( |
| { |
| "episode_id": episode_id, |
| "policy": name, |
| "matches_cpu": equal, |
| "trace": sup.trace, |
| } |
| ) |
| + "\n" |
| ) |
| return {"episodes": len(reference), "matching_trajectories": agreements, "failures": failures} |
|
|
|
|
| parser = argparse.ArgumentParser() |
| parser.add_argument("--repeats", type=int, default=3) |
| parser.add_argument("--output", default="reports/device-benchmark.json") |
| args = parser.parse_args() |
| if not torch.backends.mps.is_available(): |
| raise SystemExit("MPS unavailable") |
| source = json.loads(Path("data/fixtures/playground-probes.json").read_text()) |
| cases = [{"id": c["id"], "text": c["text"], "schema": source["schema"]} for c in source["cases"]] |
| for line in Path("data/fixtures/development.jsonl").read_text().splitlines(): |
| case = json.loads(line) |
| cases.append({"id": case["id"], "text": case["text"], "schema": schema_spec()}) |
| for name in config("scenarios"): |
| observation = Simulator(name, 42).observe() |
| for direct in [False, True]: |
| cases.append( |
| { |
| "id": f"initial-{name}-{direct}", |
| "text": render(observation), |
| "schema": schema_spec(direct), |
| } |
| ) |
| result = { |
| "precision": "float32", |
| "fallback_enabled": False, |
| "repeats": args.repeats, |
| "timing": "synchronized, one untimed warm call for each exact input/schema shape", |
| "records": {}, |
| "metadata": {}, |
| } |
| for device in ["cpu", "mps"]: |
| model = LocalModel(device=device) |
| result["metadata"][device] = model.metadata |
| records = [] |
| for i, case in enumerate(cases): |
| warm = model.classify(case["text"], case["schema"]) |
| if warm["error"]: |
| raise RuntimeError(f"{device}/{case['id']}: {warm['error']}") |
| for repeat in range(args.repeats): |
| record = model.classify(case["text"], case["schema"]) |
| records.append({"case_id": case["id"], "repeat": repeat, **record}) |
| if i % 10 == 0: |
| print(device, i + 1, "/", len(cases), flush=True) |
| result["records"][device] = records |
| if device == "mps": |
| result["episode_parity"] = episode_parity(model) |
| result["mps_memory_bytes"] = { |
| "allocated": torch.mps.current_allocated_memory(), |
| "driver": torch.mps.driver_allocated_memory(), |
| } |
| del model |
| gc.collect() |
| if device == "mps": |
| torch.mps.empty_cache() |
| cpu, mps = result["records"]["cpu"], result["records"]["mps"] |
| max_difference = 0.0 |
| max_logit_difference = 0.0 |
| agree = head_count = errors = 0 |
| for a, b in zip(cpu, mps): |
| assert (a["case_id"], a["repeat"], a["schema_hash"], a["input_hash"]) == ( |
| b["case_id"], |
| b["repeat"], |
| b["schema_hash"], |
| b["input_hash"], |
| ) |
| if a["error"] or b["error"]: |
| errors += 1 |
| continue |
| for name, head in a["heads"].items(): |
| other = b["heads"][name] |
| head_count += 1 |
| agree += head["selected_label"] == other["selected_label"] |
| for label, value in head["probabilities"].items(): |
| max_difference = max(max_difference, abs(value - other["probabilities"][label])) |
| max_logit_difference = max( |
| max_logit_difference, abs(head["logits"][label] - other["logits"][label]) |
| ) |
| summary = { |
| "cases": len(cases), |
| "head_comparisons": head_count, |
| "selected_label_agreements": agree, |
| "errors": errors, |
| "max_absolute_probability_difference": max_difference, |
| "max_absolute_logit_difference": max_logit_difference, |
| } |
| for device, records in result["records"].items(): |
| summary[device] = { |
| key: dict( |
| zip( |
| ["p50", "p95"], |
| [float(v) for v in np.quantile([r[key] for r in records], [0.5, 0.95])], |
| ) |
| ) |
| for key in ["total_ms", "inference_ms"] |
| } |
| summary["median_total_speedup"] = ( |
| summary["cpu"]["total_ms"]["p50"] / summary["mps"]["total_ms"]["p50"] |
| ) |
| summary["parity_passed"] = errors == 0 and agree == head_count and max_difference < 1e-3 |
| result["summary"] = summary |
| Path(args.output).write_text(json.dumps(result, indent=2) + "\n") |
| print(json.dumps(summary, indent=2)) |
|
|
