| """Measure quantization drift against BF16; never used for calibration.""" |
| import gc |
| import fcntl |
| import hashlib |
| import json |
| import math |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| from mlx_lm import load |
|
|
| from release import ROOT, REV, VARIANTS, write_json |
|
|
| gpu_lock = (ROOT / "artifacts/.gpu.lock").open("a") |
| fcntl.flock(gpu_lock, fcntl.LOCK_EX) |
|
|
| corpus_path = ROOT / "evaluation.json" |
| corpus = json.loads(corpus_path.read_text()) |
| corpus_hash = hashlib.sha256(corpus_path.read_bytes()).hexdigest() |
| reference_file = ROOT / "artifacts/bf16-reference-logprobs.safetensors" |
|
|
| if not reference_file.exists(): |
| model, tokenizer = load(str(ROOT / "artifacts/bf16")) |
| reference = {} |
| for i, item in enumerate(corpus): |
| ids = tokenizer.encode(item["text"]) |
| logits = model(mx.array([ids[:-1]])).astype(mx.float32) |
| logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True) |
| mx.eval(logprobs) |
| reference[str(i)] = logprobs |
| del logits |
| mx.save_safetensors(str(reference_file), reference, metadata={"corpus_sha256": corpus_hash, "source_revision": REV}) |
| del model, tokenizer, reference, logprobs |
| gc.collect() |
| mx.clear_cache() |
|
|
| reference, metadata = mx.load(str(reference_file), return_metadata=True) |
| assert metadata["corpus_sha256"] == corpus_hash |
| assert metadata["source_revision"] == REV |
| summary = {} |
| for name in VARIANTS: |
| if not (ROOT / "artifacts" / name / "conversion.json").exists(): |
| continue |
| report_path = ROOT / "reports" / f"{name}.fidelity.json" |
| if report_path.exists(): |
| summary[name] = json.loads(report_path.read_text()) |
| continue |
| rows = [] |
| if name != "bf16": |
| model, tokenizer = load(str(ROOT / "artifacts" / name)) |
| for i, item in enumerate(corpus): |
| teacher = reference[str(i)] |
| if name == "bf16": |
| student = teacher |
| else: |
| ids = tokenizer.encode(item["text"]) |
| logits = model(mx.array([ids[:-1]])).astype(mx.float32) |
| student = logits - mx.logsumexp(logits, axis=-1, keepdims=True) |
| kl = float(mx.mean(mx.sum(mx.exp(teacher) * (teacher - student), axis=-1)).item()) |
| agreement = float(mx.mean(mx.argmax(teacher, axis=-1) == mx.argmax(student, axis=-1)).item()) |
| assert math.isfinite(kl) and kl >= -1e-6 |
| rows.append({"category": item["category"], "tokens": teacher.shape[1], |
| "kl_bf16_to_quant_nats": max(0.0, kl), "top1_agreement": agreement}) |
| count = sum(r["tokens"] for r in rows) |
| report = {"variant": name, "corpus_sha256": corpus_hash, "scored_tokens": count, |
| "mean_kl_bf16_to_quant_nats": sum(r["tokens"] * r["kl_bf16_to_quant_nats"] for r in rows) / count, |
| "top1_agreement": sum(r["tokens"] * r["top1_agreement"] for r in rows) / count, |
| "samples": rows, |
| "scope": "Teacher-forced output-distribution fidelity on a small fixed synthetic corpus; not a capability benchmark"} |
| write_json(report_path, report) |
| summary[name] = report |
| print(name, report["mean_kl_bf16_to_quant_nats"], report["top1_agreement"], flush=True) |
| if name != "bf16": |
| del model, tokenizer, logits |
| del student |
| gc.collect() |
| mx.clear_cache() |
| write_json(ROOT / "reports/fidelity-summary.json", summary) |
|
|
