spb/localvm-research Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : localvm-research4#  File      : experiments/micro/expF_error_accumulation/benchmark.py5#  Purpose   : Layer-sensitivity map — degrade-one / repair-one / repair-top-k6#              (does layer-restricted escalation cut bytes-per-escalation?)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 F — error accumulation / layer sensitivity (charter §9.F).1516Per depth-group 4-bit degradation and repair on Qwen3-1.7B, teacher-forced17over reference greedy trajectories.1819Usage:20    .venv/bin/python benchmark.py [--per-domain 8] [--gen-tokens 128] [--groups 7]21"""2223from __future__ import annotations2425import argparse26import json27import sys28import time29from datetime import datetime, timezone30from pathlib import Path3132import mlx.core as mx33import mlx.nn as nn34import numpy as np35from mlx_lm import load3637REPO_ROOT = Path(__file__).resolve().parents[3]38sys.path.insert(0, str(REPO_ROOT / "benchmarks"))39sys.path.insert(0, str(REPO_ROOT / "src"))40from hardware_manifest import collect_manifest  # noqa: E40241from localvm.quality.decision_stats import greedy_generate, kl_ref_vs, teacher_forced_stats  # noqa: E4024243GROUP_SIZE = 6444BITS = 4454647def layer_index(path: str) -> int | None:48    parts = path.split(".")49    for i, p in enumerate(parts):50        if p == "layers" and i + 1 < len(parts) and parts[i + 1].isdigit():51            return int(parts[i + 1])52    return None535455def quantized_weights_by_layer(model) -> dict[str, tuple[int, mx.array]]:56    """{param_path: (layer_idx, 4-bit-dequantized bf16 weight)} for all57    divisible Linear layers inside transformer blocks."""58    out = {}59    for path, module in model.named_modules():60        li = layer_index(path)61        if li is None or not isinstance(module, nn.Linear):62            continue63        if module.weight.shape[-1] % GROUP_SIZE != 0:64            continue65        w = module.weight.astype(mx.float32)66        qw, sc, bi = mx.quantize(w, group_size=GROUP_SIZE, bits=BITS)67        deq = mx.dequantize(qw, sc, bi, group_size=GROUP_SIZE, bits=BITS).astype(mx.bfloat16)68        mx.eval(deq)69        out[path] = (li, deq)70    return out717273def apply_config(model, qweights: dict, originals: dict, degrade_layers: set[int]) -> None:74    """Set each eligible Linear to 4-bit dequant if its layer ∈ degrade_layers,75    else restore the original bf16 weight."""76    for path, module in model.named_modules():77        if path in qweights:78            li, deq = qweights[path]79            module.weight = deq if li in degrade_layers else originals[path]808182def evaluate(model, trajectories, ref_stats) -> dict:83    agrees, kls = [], []84    for t, ref in zip(trajectories, ref_stats):85        qs = teacher_forced_stats(model, t["full_ids"], t["start"])86        ref_next = np.array(t["full_ids"][t["start"]:])87        agrees.append((qs["argmax"] == ref_next).astype(np.int8))88        kls.append(kl_ref_vs(qs["logprobs"], ref["logprobs"]))89    return {90        "agreement_rate": float(np.concatenate(agrees).mean()),91        "mean_kl": float(np.mean(np.concatenate(kls))),92    }939495def main() -> None:96    ap = argparse.ArgumentParser()97    ap.add_argument("--model", default="mlx-community/Qwen3-1.7B-bf16")98    ap.add_argument("--gen-tokens", type=int, default=128)99    ap.add_argument("--per-domain", type=int, default=8)100    ap.add_argument("--groups", type=int, default=7)101    args = ap.parse_args()102103    domains = json.loads((REPO_ROOT / "benchmarks/datasets/eval_prompts.json").read_text())["domains"]104    print(f"loading {args.model} …", flush=True)105    model, tokenizer = load(args.model)106    n_layers = len(model.model.layers)107    bounds = np.linspace(0, n_layers, args.groups + 1).astype(int)108    groups = [set(range(bounds[i], bounds[i + 1])) for i in range(args.groups)]109110    trajectories = []111    t0 = time.time()112    for domain, plist in domains.items():113        for prompt in plist[: args.per_domain]:114            ids = tokenizer.apply_chat_template(115                [{"role": "user", "content": prompt}], add_generation_prompt=True)116            gen = greedy_generate(model, tokenizer, ids, args.gen_tokens)117            if len(gen) >= 8:118                trajectories.append({"domain": domain, "full_ids": list(ids) + gen, "start": len(ids)})119    print(f"{len(trajectories)} trajectories in {time.time()-t0:.0f}s", flush=True)120    ref_stats = [teacher_forced_stats(model, t["full_ids"], t["start"]) for t in trajectories]121122    print("precomputing 4-bit weights …", flush=True)123    qweights = quantized_weights_by_layer(model)124    originals = {p: m.weight for p, m in model.named_modules() if p in qweights}125    all_layers = set(range(n_layers))126127    runs: dict[str, dict] = {}128129    def run(tag: str, degrade: set[int]) -> dict:130        apply_config(model, qweights, originals, degrade)131        r = evaluate(model, trajectories, ref_stats)132        r["degraded_layers"] = sorted(degrade)133        runs[tag] = r134        print(f"  {tag:>24}: agree={r['agreement_rate']:.4f} KL={r['mean_kl']:.4f}", flush=True)135        return r136137    print("all-4-bit floor:", flush=True)138    floor = run("all_4bit", all_layers)139    print("degrade-one (rest bf16):", flush=True)140    for gi, g in enumerate(groups):141        run(f"degrade_g{gi}_L{min(g)}-{max(g)}", g)142    print("repair-one (rest 4-bit):", flush=True)143    for gi, g in enumerate(groups):144        run(f"repair_g{gi}_L{min(g)}-{max(g)}", all_layers - g)145146    # repair-top-k by measured repair value147    lost = 1.0 - floor["agreement_rate"]148    repair_value = {149        gi: runs[f"repair_g{gi}_L{min(g)}-{max(g)}"]["agreement_rate"] - floor["agreement_rate"]150        for gi, g in enumerate(groups)151    }152    order = sorted(repair_value, key=repair_value.get, reverse=True)153    print("repair-top-k (best groups bf16):", flush=True)154    for k in (2, 3):155        keep = set().union(*(groups[gi] for gi in order[:k]))156        run(f"repair_top{k}_groups_{sorted(order[:k])}", all_layers - keep)157158    apply_config(model, qweights, originals, set())  # restore159160    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")161    out_dir = REPO_ROOT / "results" / "expF_error_accumulation" / ts162    out_dir.mkdir(parents=True)163    (out_dir / "results.json").write_text(json.dumps({164        "experiment": "expF_error_accumulation",165        "author": "Simon-Pierre Boucher",166        "contact": "contact@spboucher.ai",167        "manifest": collect_manifest(),168        "config": vars(args),169        "bits": BITS, "group_size": GROUP_SIZE,170        "n_layers": n_layers,171        "layer_groups": [sorted(g) for g in groups],172        "agreement_lost_all4bit": lost,173        "repair_value_by_group": {str(k): v for k, v in repair_value.items()},174        "runs": runs,175    }, indent=2))176    print(f"\nwrote {out_dir / 'results.json'}")177178179if __name__ == "__main__":180    main()181