spb/localvm-research Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : localvm-research4#  File      : experiments/micro/expD_progressive_reconstruction/benchmark.py5#  Purpose   : Residual-ladder progressive weight reconstruction — decision and6#              hidden-state convergence vs cumulative bits (candidate C1 math)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 D — progressive weight reconstruction (charter §9.D).1516Builds base+residual affine-quantized ladders (3/3+3/3+3+3 and 4/4+4 bits),17teacher-forces each cumulative stage over reference greedy trajectories, and18measures decision convergence, two-tier margin-gated policies, and19hidden-state error at several depths.2021Usage:22    .venv/bin/python benchmark.py [--model mlx-community/Qwen3-1.7B-bf16]23        [--gen-tokens 128] [--per-domain 8]24"""2526from __future__ import annotations2728import argparse29import json30import sys31import time32from datetime import datetime, timezone33from pathlib import Path3435import mlx.core as mx36import mlx.nn as nn37import numpy as np38from mlx_lm import load3940REPO_ROOT = Path(__file__).resolve().parents[3]41sys.path.insert(0, str(REPO_ROOT / "benchmarks"))42sys.path.insert(0, str(REPO_ROOT / "src"))43from hardware_manifest import collect_manifest  # noqa: E40244from localvm.quality.decision_stats import (  # noqa: E40245    auroc, escalation_curve, greedy_generate, kl_ref_vs, teacher_forced_stats,46)4748GROUP = 64495051def quantizable(m) -> bool:52    return isinstance(m, nn.Linear) and m.weight.shape[-1] % GROUP == 0535455def residual_ladder_weights(model, ladder: list[int]) -> list[dict[str, mx.array]]:56    """For each quantizable Linear, build cumulative dequantized weights for57    each stage of `ladder` (bits per stage). Returns a list (one per stage) of58    {param_path: bf16 weight} replacements. Memory: one bf16 copy per stage59    per layer is materialized lazily at apply time; here we keep the per-stage60    cumulative tensors (float32 accumulation, cast to bf16)."""61    stages = [dict() for _ in ladder]62    for path, module in model.named_modules():63        if not quantizable(module):64            continue65        w = module.weight.astype(mx.float32)66        acc = mx.zeros_like(w)67        err = w68        for k, bits in enumerate(ladder):69            qw, scales, biases = mx.quantize(err, group_size=GROUP, bits=bits)70            deq = mx.dequantize(qw, scales, biases, group_size=GROUP, bits=bits)71            acc = acc + deq72            err = w - acc73            stages[k][path] = acc.astype(mx.bfloat16)74            mx.eval(stages[k][path])75    return stages767778def apply_weights(model, replacement: dict[str, mx.array]) -> dict[str, mx.array]:79    """Swap Linear weights in place; returns the originals for restoration."""80    originals = {}81    for path, module in model.named_modules():82        if path in replacement:83            originals[path] = module.weight84            module.weight = replacement[path]85    return originals868788def hidden_state_errors(model, ref_hidden: dict, full_ids: list[int], start: int,89                        depths: list[int]) -> dict[int, float]:90    """Relative L2 error of hidden states vs reference at given layer indices."""91    h = capture_hidden(model, full_ids, start, depths)92    out = {}93    for d in depths:94        r, q = ref_hidden[d], h[d]95        out[d] = float(np.linalg.norm(q - r) / (np.linalg.norm(r) + 1e-9))96    return out979899def capture_hidden(model, full_ids: list[int], start: int, depths: list[int]) -> dict:100    """Hidden states (post-layer) at selected depths for predicted positions.101    Replicates the inner transformer loop manually (instance-level __call__102    monkey-patching does not intercept Python's type-level dunder dispatch)."""103    from mlx_lm.models.base import create_attention_mask104105    inner = model.model106    h = inner.embed_tokens(mx.array(full_ids)[None])107    mask = create_attention_mask(h, None)108    result = {}109    for i, layer in enumerate(inner.layers):110        h = layer(h, mask, cache=None)111        if i in depths:112            t = h[0, start - 1 : len(full_ids) - 1].astype(mx.float32)113            mx.eval(t)114            result[i] = np.array(t)115    return result116117118def two_tier_policy(lo: dict, hi: dict, ref_next: np.ndarray, taus: list[float]) -> list[dict]:119    """Policy: take lo's decision when its margin >= tau, else hi's decision."""120    out = []121    for tau in taus:122        esc = lo["margin"] < tau123        decision = np.where(esc, hi["argmax"], lo["argmax"])124        out.append({125            "tau": tau,126            "escalated_frac": float(esc.mean()),127            "policy_agreement": float((decision == ref_next).mean()),128        })129    return out130131132def main() -> None:133    ap = argparse.ArgumentParser()134    ap.add_argument("--model", default="mlx-community/Qwen3-1.7B-bf16")135    ap.add_argument("--gen-tokens", type=int, default=128)136    ap.add_argument("--per-domain", type=int, default=8)137    ap.add_argument("--hidden-trajectories", type=int, default=8)138    args = ap.parse_args()139140    domains = json.loads((REPO_ROOT / "benchmarks/datasets/eval_prompts.json").read_text())["domains"]141142    print(f"loading reference {args.model} …", flush=True)143    model, tokenizer = load(args.model)144    n_layers = len(model.model.layers)145    depths = [max(0, round(n_layers * f) - 1) for f in (0.25, 0.5, 0.75, 1.0)]146147    trajectories = []148    t0 = time.time()149    for domain, plist in domains.items():150        for prompt in plist[: args.per_domain]:151            ids = tokenizer.apply_chat_template(152                [{"role": "user", "content": prompt}], add_generation_prompt=True)153            gen = greedy_generate(model, tokenizer, ids, args.gen_tokens)154            if len(gen) >= 8:155                trajectories.append({"domain": domain, "full_ids": list(ids) + gen, "start": len(ids)})156    print(f"{len(trajectories)} reference trajectories in {time.time()-t0:.0f}s", flush=True)157158    ref_stats = [teacher_forced_stats(model, t["full_ids"], t["start"]) for t in trajectories]159    hidden_subset = trajectories[:: max(1, len(trajectories) // args.hidden_trajectories)][: args.hidden_trajectories]160    ref_hidden = [capture_hidden(model, t["full_ids"], t["start"], depths) for t in hidden_subset]161162    ladders = {"A_base3": [3, 3, 3], "B_base4": [4, 4]}163    results: dict[str, list] = {}164    overhead_bits = 32 / GROUP * 2  # bf16 scales + biases per group per stage165166    for name, ladder in ladders.items():167        print(f"building ladder {name} {ladder} …", flush=True)168        stages = residual_ladder_weights(model, ladder)169        stage_records = []170        for k, replacement in enumerate(stages):171            originals = apply_weights(model, replacement)172            rows_margin, rows_agree, rows_kl, rows_argmax, doms = [], [], [], [], []173            for t, ref in zip(trajectories, ref_stats):174                qs = teacher_forced_stats(model, t["full_ids"], t["start"])175                ref_next = np.array(t["full_ids"][t["start"]:])176                rows_margin.append(qs["margin"])177                rows_argmax.append(qs["argmax"])178                rows_agree.append((qs["argmax"] == ref_next).astype(np.int8))179                rows_kl.append(kl_ref_vs(qs["logprobs"], ref["logprobs"]))180                doms.append(t["domain"])181            hid = [hidden_state_errors(model, rh, t["full_ids"], t["start"], depths)182                   for rh, t in zip(ref_hidden, hidden_subset)]183            apply_weights(model, originals)184185            margins = np.concatenate(rows_margin)186            agrees = np.concatenate(rows_agree)187            cum_bits = sum(ladder[: k + 1]) + overhead_bits * (k + 1)188            rec = {189                "stage": k,190                "ladder_bits": ladder[: k + 1],191                "cumulative_bits_per_param": round(cum_bits, 2),192                "agreement_rate": float(agrees.mean()),193                "mean_kl": float(np.mean(np.concatenate(rows_kl))),194                "auroc": auroc(-margins, 1 - agrees),195                "escalation_curve": escalation_curve(margins, agrees),196                "hidden_rel_err_by_depth": {197                    str(d): float(np.mean([h[d] for h in hid])) for d in depths198                },199                "_margins": margins, "_argmax": np.concatenate(rows_argmax),200            }201            for pt in rec["escalation_curve"]:202                if pt["residual_disagree"] <= 0.01:203                    rec["escalation_frac_for_99pct"] = pt["escalated_frac"]204                    break205            stage_records.append(rec)206            print(f"  stage {k} ({rec['cumulative_bits_per_param']} bits): "207                  f"agree={rec['agreement_rate']:.4f} KL={rec['mean_kl']:.4f} "208                  f"auroc={rec['auroc']:.3f}", flush=True)209        results[name] = stage_records210211    # two-tier margin-gated policies between consecutive stages212    ref_next_all = np.concatenate([np.array(t["full_ids"][t["start"]:]) for t in trajectories])213    taus = [0.25, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]214    policies = {}215    for name, recs in results.items():216        for k in range(len(recs) - 1):217            lo = {"margin": recs[k]["_margins"], "argmax": recs[k]["_argmax"]}218            hi = {"argmax": recs[k + 1]["_argmax"]}219            policies[f"{name}_stage{k}_to_{k+1}"] = two_tier_policy(lo, hi, ref_next_all, taus)220    for recs in results.values():221        for r in recs:222            del r["_margins"], r["_argmax"]223224    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")225    out_dir = REPO_ROOT / "results" / "expD_progressive_reconstruction" / ts226    out_dir.mkdir(parents=True)227    (out_dir / "results.json").write_text(json.dumps({228        "experiment": "expD_progressive_reconstruction",229        "author": "Simon-Pierre Boucher",230        "contact": "contact@spboucher.ai",231        "manifest": collect_manifest(),232        "config": vars(args),233        "group_size": GROUP,234        "scale_overhead_bits_per_param_per_stage": overhead_bits,235        "n_trajectories": len(trajectories),236        "hidden_depth_layers": depths,237        "ladders": {k: v for k, v in results.items()},238        "two_tier_policies": policies,239    }, indent=2))240    print(f"\nwrote {out_dir / 'results.json'}")241242243if __name__ == "__main__":244    main()245