spb/localvm-research Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : localvm-research4#  File      : experiments/micro/expG_decision_stability/benchmark.py5#  Purpose   : Joint (cheap-pass margin × agreement) matrix across bit-widths —6#              the Gate-Zero measurement for margin-gated escalation (G02/G23)7#  Author    : Simon-Pierre Boucher8#  Contact   : contact@spboucher.ai9#  Created   : 2026-08-1210#  Modified  : 2026-08-1211#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal12#  License   : All rights reserved (research code)13# =============================================================================14"""Experiment G — decision stability (charter §9.G).1516Generates greedy continuations with a bf16 reference model, teacher-forces17low-bit quantized variants over the same sequences, and records per-position18margin/agreement/KL. Outputs the joint matrix, AUROC of margin as a19disagreement detector, and escalation curves.2021Usage:22    .venv/bin/python benchmark.py [--model mlx-community/Qwen3-1.7B-bf16]23        [--gen-tokens 128] [--bits 2,3,4,8] [--per-domain 8]24"""2526from __future__ import annotations2728import argparse29import gc30import json31import sys32import time33from datetime import datetime, timezone34from pathlib import Path3536import mlx.core as mx37import mlx.nn as nn38import numpy as np39from mlx_lm import load4041REPO_ROOT = Path(__file__).resolve().parents[3]42sys.path.insert(0, str(REPO_ROOT / "benchmarks"))43from hardware_manifest import collect_manifest  # noqa: E402444546def greedy_generate(model, tokenizer, prompt_ids: list[int], n_tokens: int) -> list[int]:47    """Greedy generation without sampling helpers — deterministic, no cache reuse48    across prompts. Returns generated token ids."""49    tokens = list(prompt_ids)50    generated = []51    from mlx_lm.models.cache import make_prompt_cache5253    cache = make_prompt_cache(model)54    inp = mx.array(tokens)[None]55    for _ in range(n_tokens):56        logits = model(inp, cache=cache)57        nxt = int(mx.argmax(logits[0, -1]).item())58        if nxt == tokenizer.eos_token_id:59            break60        generated.append(nxt)61        inp = mx.array([[nxt]])62    return generated636465def teacher_forced_stats(model, full_ids: list[int], start: int) -> dict:66    """Forward the full sequence once; return per-position stats for positions67    predicting tokens at indices [start, len(full_ids)) — i.e., logits at68    positions start-1 .. len-2."""69    logits = model(mx.array(full_ids)[None])[0]  # (T, V)70    sel = logits[start - 1 : len(full_ids) - 1].astype(mx.float32)71    top2 = mx.topk(sel, 2, axis=-1)  # values sorted ascending in MLX topk72    argmax = mx.argmax(sel, axis=-1)73    logprobs = sel - mx.logsumexp(sel, axis=-1, keepdims=True)74    mx.eval(top2, argmax, logprobs)75    v = np.array(top2)76    margin = v[:, 1] - v[:, 0] if v[0, 1] >= v[0, 0] else v[:, 0] - v[:, 1]77    return {78        "margin": np.abs(margin),79        "argmax": np.array(argmax),80        # float16 storage: 48 trajectories × (128, ~152k vocab) would be ~4 GB81        # in float32; KL is computed in float32 at use time.82        "logprobs": np.array(logprobs).astype(np.float16),83    }848586def auroc(scores: np.ndarray, labels: np.ndarray) -> float:87    """AUROC of `scores` (higher = predicted positive) for binary labels.88    Here: score = -margin (low margin should predict disagreement=1)."""89    pos, neg = scores[labels == 1], scores[labels == 0]90    if len(pos) == 0 or len(neg) == 0:91        return float("nan")92    order = np.argsort(np.concatenate([pos, neg]), kind="mergesort")93    ranks = np.empty(len(order)); ranks[order] = np.arange(1, len(order) + 1)94    # average ranks for ties95    allv = np.concatenate([pos, neg])96    sorted_v = allv[order]97    i = 098    while i < len(sorted_v):99        j = i100        while j + 1 < len(sorted_v) and sorted_v[j + 1] == sorted_v[i]:101            j += 1102        if j > i:103            ranks[order[i : j + 1]] = ranks[order[i : j + 1]].mean()104        i = j + 1105    r_pos = ranks[: len(pos)].sum()106    return float((r_pos - len(pos) * (len(pos) + 1) / 2) / (len(pos) * len(neg)))107108109def escalation_curve(margins: np.ndarray, agree: np.ndarray, points: int = 200) -> list[dict]:110    """For threshold τ over margins: escalate tokens with margin < τ (assume the111    escalated decision becomes correct). Report escalated fraction vs residual112    disagreement (disagreements with margin ≥ τ)."""113    qs = np.quantile(margins, np.linspace(0, 1, points))114    out, n = [], len(margins)115    for tau in qs:116        esc = margins < tau117        residual = np.sum((~esc) & (agree == 0)) / n118        out.append({"tau": float(tau), "escalated_frac": float(esc.mean()),119                    "residual_disagree": float(residual)})120    return out121122123def main() -> None:124    ap = argparse.ArgumentParser()125    ap.add_argument("--model", default="mlx-community/Qwen3-1.7B-bf16")126    ap.add_argument("--gen-tokens", type=int, default=128)127    ap.add_argument("--bits", default="2,3,4,8")128    ap.add_argument("--per-domain", type=int, default=8)129    ap.add_argument("--group-size", type=int, default=64)130    args = ap.parse_args()131    bits_list = [int(b) for b in args.bits.split(",")]132133    prompts_file = REPO_ROOT / "benchmarks" / "datasets" / "eval_prompts.json"134    domains = json.loads(prompts_file.read_text())["domains"]135136    print(f"loading reference {args.model} …", flush=True)137    model, tokenizer = load(args.model)138139    # -------- pass 1: reference greedy trajectories + reference stats140    trajectories = []  # {domain, prompt_ids, full_ids, start}141    t0 = time.time()142    for domain, plist in domains.items():143        for prompt in plist[: args.per_domain]:144            msgs = [{"role": "user", "content": prompt}]145            ids = tokenizer.apply_chat_template(msgs, add_generation_prompt=True)146            gen = greedy_generate(model, tokenizer, ids, args.gen_tokens)147            if len(gen) < 8:148                continue149            trajectories.append({"domain": domain, "full_ids": list(ids) + gen, "start": len(ids)})150        print(f"  generated {domain}", flush=True)151    print(f"reference generation done in {time.time() - t0:.0f}s "152          f"({len(trajectories)} trajectories)", flush=True)153154    ref_stats = [teacher_forced_stats(model, t["full_ids"], t["start"]) for t in trajectories]155156    # -------- pass 2: quantized variants, teacher-forced on the same ids157    per_bits: dict[int, dict] = {}158    for bits in bits_list:159        print(f"quantizing to {bits}-bit (group {args.group_size}) …", flush=True)160        del model161        gc.collect(); mx.clear_cache()162        model, _ = load(args.model)163        nn.quantize(model, group_size=args.group_size, bits=bits,164                    class_predicate=lambda p, m: isinstance(m, nn.Linear)165                    and m.weight.shape[-1] % args.group_size == 0)166        rows = []167        for t, ref in zip(trajectories, ref_stats):168            qs = teacher_forced_stats(model, t["full_ids"], t["start"])169            ref_next = np.array(t["full_ids"][t["start"]:])  # actual (=ref argmax) tokens170            agree = (qs["argmax"] == ref_next).astype(np.int8)171            # KL(ref||q) per position172            ref_lp = ref["logprobs"].astype(np.float32)173            kl = np.sum(np.exp(ref_lp) * (ref_lp - qs["logprobs"].astype(np.float32)), axis=-1)174            rows.append({"domain": t["domain"], "margin": qs["margin"],175                         "agree": agree, "kl": kl})176        margins = np.concatenate([r["margin"] for r in rows])177        agrees = np.concatenate([r["agree"] for r in rows])178        kls = np.concatenate([r["kl"] for r in rows])179        disagree = 1 - agrees180        stats = {181            "bits": bits,182            "n_positions": int(len(margins)),183            "agreement_rate": float(agrees.mean()),184            "mean_kl_ref_q": float(np.mean(kls)),185            "auroc_margin_predicts_disagreement": auroc(-margins, disagree),186            "median_margin_agree": float(np.median(margins[agrees == 1])),187            "median_margin_disagree": float(np.median(margins[agrees == 0])) if (agrees == 0).any() else None,188            "escalation_curve": escalation_curve(margins, agrees),189            "per_domain": {190                d: {191                    "agreement_rate": float(np.concatenate([r["agree"] for r in rows if r["domain"] == d]).mean()),192                    "auroc": auroc(193                        -np.concatenate([r["margin"] for r in rows if r["domain"] == d]),194                        1 - np.concatenate([r["agree"] for r in rows if r["domain"] == d]),195                    ),196                }197                for d in domains198            },199        }200        # operating point: escalation fraction to reach 99% agreement201        for pt in stats["escalation_curve"]:202            if pt["residual_disagree"] <= 0.01:203                stats["escalation_frac_for_99pct"] = pt["escalated_frac"]204                break205        per_bits[bits] = stats206        print(f"  {bits}-bit: agree={stats['agreement_rate']:.4f} "207              f"AUROC={stats['auroc_margin_predicts_disagreement']:.3f} "208              f"esc@99%={stats.get('escalation_frac_for_99pct', 'n/a')}", flush=True)209210    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")211    out_dir = REPO_ROOT / "results" / "expG_decision_stability" / ts212    out_dir.mkdir(parents=True)213    payload = {214        "experiment": "expG_decision_stability",215        "author": "Simon-Pierre Boucher",216        "contact": "contact@spboucher.ai",217        "manifest": collect_manifest(),218        "config": vars(args),219        "n_trajectories": len(trajectories),220        "results_by_bits": {str(k): v for k, v in per_bits.items()},221    }222    (out_dir / "results.json").write_text(json.dumps(payload, indent=2))223    print(f"\nwrote {out_dir / 'results.json'}")224225226if __name__ == "__main__":227    main()228