File size: 3,322 Bytes
e134e56 | 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 | """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)
|