spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expH_capture_cost_frontier/implementation/benchmark_real.py5#  Purpose   : Run #3 — capture overhead on a real quantized checkpoint (mlx-lm)6#  Author    : Simon-Pierre Boucher7#  Contact   : contact@spboucher.ai8#  Website   : https://modelmap.io9#  Created   : 2026-08-1210#  Modified  : 2026-08-1211#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal12#  License   : All rights reserved (research code)13# =============================================================================14"""expH run #3 (see hypothesis.md, registered before this run).1516Loads a real 4-bit model via mlx-lm, wraps every decoder layer with a17retaining tap, and measures prefill throughput in three modes:18plain / retain-all-layers / retain + NumPy conversion + mmap write.19This is the first quantized-model activation capture in the project —20the capability the Phase 1 survey found nowhere in Python tooling.21"""2223from __future__ import annotations2425import json26import shutil27import subprocess28import sys29import tempfile30import time31from pathlib import Path3233import numpy as np3435ROOT = Path(__file__).resolve().parents[4]36sys.path.insert(0, str(ROOT / "benchmarks"))37from hardware_manifest import manifest3839MODEL = "mlx-community/Qwen3-0.6B-4bit"40SEQ = 102441N_FWD, WARMUP, REPEATS = 10, 3, 342SEED = 0434445class Tap:46    """Wraps a decoder layer; optionally retains its output."""4748    def __init__(self, layer):49        self.layer = layer50        self.retained = None51        self.enabled = False5253    def __call__(self, *args, **kwargs):54        out = self.layer(*args, **kwargs)55        if self.enabled:56            self.retained = out57        return out5859    def __getattr__(self, name):  # delegate attribute access (e.g. .self_attn)60        return getattr(self.layer, name)616263def main() -> int:64    import mlx.core as mx65    from mlx_lm import load6667    mx.random.seed(SEED)68    model, tokenizer = load(MODEL)69    layers = model.model.layers70    n_layers = len(layers)71    taps = [Tap(l) for l in layers]72    model.model.layers = taps7374    text = ("The internal cartography of local language models requires "75            "systematic measurement of every layer. ") * 6076    tokens = tokenizer.encode(text)[:SEQ]77    x = mx.array([tokens])7879    # infer d_model at runtime (quantized embeddings pack their weight shapes)80    taps[0].enabled = True81    mx.eval(model(x))82    d_model = int(taps[0].retained.shape[-1])83    taps[0].enabled = False84    taps[0].retained = None85    print(f"model={MODEL} layers={n_layers} d_model={d_model} seq={len(tokens)}")8687    workdir = Path(tempfile.mkdtemp(prefix="modelmap_expH3_"))88    store = np.memmap(workdir / "capture.raw", dtype=np.float16, mode="w+",89                      shape=(N_FWD * n_layers * len(tokens), d_model))9091    def run(capture: bool, to_disk: bool) -> float:92        for t in taps:93            t.enabled = capture94            t.retained = None95        row = 096        for _ in range(WARMUP):97            out = model(x)98            mx.eval(out, *[t.retained for t in taps if t.retained is not None])99        t0 = time.perf_counter()100        for _ in range(N_FWD):101            out = model(x)102            retained = [t.retained for t in taps] if capture else []103            mx.eval(out, *[r for r in retained if r is not None])104            if to_disk:105                for r in retained:106                    # model runs bf16 — cast in MLX (numpy has no bfloat16)107                    a = np.array(r.astype(mx.float16), copy=False).reshape(-1, d_model)108                    store[row:row + a.shape[0]] = a109                    row += a.shape[0]110        return (time.perf_counter() - t0) / N_FWD111112    results = []113    try:114        for mode, cap, disk in (("plain", False, False),115                                ("retain", True, False),116                                ("retain+copy+write", True, True)):117            times = [run(cap, disk) for _ in range(REPEATS)]118            results.append({"backend": "mlx-lm", "mode": mode,119                            "s_per_forward": times,120                            "tokens_per_s_mean": len(tokens) / np.mean(times)})121            print(f"  {mode:22s} {np.mean(times)*1000:8.1f} ms/prefill "122                  f"({len(tokens)/np.mean(times):8.0f} tok/s)")123    finally:124        store.flush()125        shutil.rmtree(workdir, ignore_errors=True)126127    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,128                            capture_output=True, text=True, check=False).stdout.strip()129    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())130    outdir = ROOT / "results" / "expH_capture_cost_frontier" / ts131    outdir.mkdir(parents=True)132    doc = {133        "experiment": "expH_capture_cost_frontier",134        "run": 3,135        "scope": "capture overhead on a real 4-bit checkpoint (mlx-lm prefill)",136        "commit": commit,137        "config": {"model": MODEL, "seq": len(tokens), "n_layers": n_layers,138                   "d_model": int(d_model), "n_forwards": N_FWD,139                   "warmup": WARMUP, "repeats": REPEATS, "seed": SEED},140        "manifest": manifest(),141        "compute": results,142    }143    (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n")144    print(f"results -> {outdir / 'results.json'}")145    return 0146147148if __name__ == "__main__":149    sys.exit(main())150