spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expH_capture_cost_frontier/implementation/benchmark.py5#  Purpose   : Run #1 — storage-format throughput on APFS + hook overhead MLX/MPS6#  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 / Metal / MPS12#  License   : All rights reserved (research code)13# =============================================================================14"""expH run #1. See hypothesis.md (registered before this run).1516Part A — activation-store formats on APFS (warm cache, declared):17    raw np.memmap | safetensors (mmap) | zarr zstd | zarr uncompressed18    sequential write, sequential scan, random-batch reads (SAE-shuffle).19Part B — capture overhead on a synthetic 12-layer transformer (fp16):20    torch-MPS forward hooks vs MLX retained arrays;21    modes: plain / retain-on-device / retain + CPU copy + mmap write.22"""2324from __future__ import annotations2526import json27import shutil28import subprocess29import sys30import tempfile31import time32from pathlib import Path3334import numpy as np3536ROOT = Path(__file__).resolve().parents[4]37sys.path.insert(0, str(ROOT / "benchmarks"))38from hardware_manifest import manifest3940REPEATS = 341ROWS, DIM = 100_000, 4096          # ~0.82 GB fp16 per format42WRITE_CHUNK = 4_096                 # rows per write43BATCH = 4_096                       # rows per random-batch read44N_BATCHES = 2445SEED = 04647# Part B config48L, D_MODEL, N_HEADS, SEQ, BSZ = 12, 1024, 4, 512, 849N_FWD, N_FWD_WRITE, WARMUP = 20, 10, 3505152def _timed(fn, repeats=REPEATS):53    out = []54    for _ in range(repeats):55        t0 = time.perf_counter()56        fn()57        out.append(time.perf_counter() - t0)58    return out596061# ---------------------------------------------------------------- Part A62def bench_storage(workdir: Path) -> list[dict]:63    rng = np.random.default_rng(SEED)64    data = rng.standard_normal((WRITE_CHUNK, DIM)).astype(np.float16)65    batches = [rng.integers(0, ROWS, BATCH) for _ in range(N_BATCHES)]66    total_bytes = ROWS * DIM * 267    batch_bytes = BATCH * DIM * 2 * N_BATCHES68    results = []6970    def record(fmt, op, times, nbytes):71        results.append({72            "format": fmt, "op": op, "bytes": nbytes,73            "seconds": times, "gb_per_s_mean": nbytes / 2**30 / np.mean(times),74        })75        print(f"  {fmt:18s} {op:12s} {nbytes/2**30/np.mean(times):8.2f} GB/s")7677    # ---- raw np.memmap78    p = workdir / "acts.raw"79    def write_raw():80        m = np.memmap(p, dtype=np.float16, mode="w+", shape=(ROWS, DIM))81        for i in range(0, ROWS, WRITE_CHUNK):82            end = min(i + WRITE_CHUNK, ROWS)83            m[i:end] = data[: end - i]84        m.flush(); del m85    record("raw-mmap", "write", _timed(write_raw), total_bytes)86    m = np.memmap(p, dtype=np.float16, mode="r", shape=(ROWS, DIM))87    record("raw-mmap", "seq-scan", _timed(lambda m=m: float(np.asarray(m).sum(dtype=np.float32))), total_bytes)88    record("raw-mmap", "random-batch", _timed(lambda m=m: [m[b].sum(dtype=np.float32) for b in batches]), batch_bytes)89    del m9091    # ---- safetensors (mmap-backed numpy)92    from safetensors import safe_open93    from safetensors.numpy import save_file94    p = workdir / "acts.safetensors"95    full = np.memmap(workdir / "acts.raw", dtype=np.float16, mode="r", shape=(ROWS, DIM))96    def write_st(full=full, p=p):97        save_file({"acts": np.asarray(full)}, str(p))98    record("safetensors", "write", _timed(write_st), total_bytes)99    def open_st():100        return safe_open(str(p), framework="np")101    f = open_st()102    t = f.get_tensor("acts")   # mmap-backed load103    record("safetensors", "seq-scan", _timed(lambda t=t: float(t.sum(dtype=np.float32))), total_bytes)104    record("safetensors", "random-batch", _timed(lambda t=t: [t[b].sum(dtype=np.float32) for b in batches]), batch_bytes)105    del t, f, full106107    # ---- zarr (zstd default) and uncompressed108    import zarr109    for codec, name in ((None, "zarr-uncompressed"), ("default", "zarr-zstd")):110        p = workdir / f"acts_{name}.zarr"111        kwargs = {} if codec == "default" else {"compressors": None}112        def write_zarr(p=p, kwargs=kwargs):113            if p.exists():114                shutil.rmtree(p)115            z = zarr.create_array(store=str(p), shape=(ROWS, DIM), chunks=(WRITE_CHUNK, DIM),116                                  dtype=np.float16, **kwargs)117            for i in range(0, ROWS, WRITE_CHUNK):118                end = min(i + WRITE_CHUNK, ROWS)119                z[i:end] = data[: end - i]120        record(name, "write", _timed(write_zarr), total_bytes)121        z = zarr.open_array(store=str(p), mode="r")122        record(name, "seq-scan", _timed(lambda z=z: float(z[:].sum(dtype=np.float32))), total_bytes)123        record(name, "random-batch", _timed(lambda z=z: [z[np.sort(b)].sum(dtype=np.float32) for b in batches]), batch_bytes)124125    return results126127128# ---------------------------------------------------------------- Part B129def bench_torch_mps(workdir: Path) -> list[dict]:130    import torch131    from torch import nn132    assert torch.backends.mps.is_available(), "MPS required"133    dev, dt = torch.device("mps"), torch.float16134135    class Block(nn.Module):136        def __init__(self):137            super().__init__()138            self.ln1, self.ln2 = nn.LayerNorm(D_MODEL), nn.LayerNorm(D_MODEL)139            self.attn = nn.MultiheadAttention(D_MODEL, N_HEADS, batch_first=True)140            self.mlp = nn.Sequential(nn.Linear(D_MODEL, 4 * D_MODEL), nn.GELU(),141                                     nn.Linear(4 * D_MODEL, D_MODEL))142        def forward(self, x):143            h = self.ln1(x)144            x = x + self.attn(h, h, h, need_weights=False)[0]145            return x + self.mlp(self.ln2(x))146147    torch.manual_seed(SEED)148    model = nn.Sequential(*[Block() for _ in range(L)]).to(dev, dt).eval()149    x = torch.randn(BSZ, SEQ, D_MODEL, device=dev, dtype=dt)150    store = np.memmap(workdir / "torch_capture.raw", dtype=np.float16, mode="w+",151                      shape=(N_FWD_WRITE * L * BSZ * SEQ, D_MODEL))152153    def run(n, capture, to_disk):154        captured, row = [], 0155        hooks = []156        if capture:157            def hook(_m, _i, out):158                captured.append(out)159            hooks = [b.register_forward_hook(hook) for b in model]160        with torch.no_grad():161            for _ in range(WARMUP):162                model(x)163            torch.mps.synchronize()164            t0 = time.perf_counter()165            for _ in range(n):166                captured.clear()167                model(x)168                if to_disk:169                    for c in captured:170                        a = c.to("cpu").numpy().reshape(-1, D_MODEL)171                        store[row:row + a.shape[0]] = a172                        row += a.shape[0]173            torch.mps.synchronize()174            dt_s = time.perf_counter() - t0175        for h in hooks:176            h.remove()177        return dt_s / n178179    out = []180    for mode, cap, disk, n in (("plain", False, False, N_FWD),181                               ("retain", True, False, N_FWD),182                               ("retain+copy+write", True, True, N_FWD_WRITE)):183        times = [run(n, cap, disk) for _ in range(REPEATS)]184        out.append({"backend": "torch-mps", "mode": mode,185                    "s_per_forward": times,186                    "tokens_per_s_mean": BSZ * SEQ / np.mean(times)})187        print(f"  torch-mps {mode:22s} {np.mean(times)*1000:8.1f} ms/fwd")188    return out189190191def bench_mlx(workdir: Path) -> list[dict]:192    import mlx.core as mx193    import mlx.nn as mnn194195    class Block(mnn.Module):196        def __init__(self):197            super().__init__()198            self.ln1, self.ln2 = mnn.LayerNorm(D_MODEL), mnn.LayerNorm(D_MODEL)199            self.attn = mnn.MultiHeadAttention(D_MODEL, N_HEADS)200            self.fc1, self.fc2 = mnn.Linear(D_MODEL, 4 * D_MODEL), mnn.Linear(4 * D_MODEL, D_MODEL)201        def __call__(self, x):202            h = self.ln1(x)203            x = x + self.attn(h, h, h)204            return x + self.fc2(mnn.gelu(self.fc1(self.ln2(x))))205206    mx.random.seed(SEED)207    blocks = [Block() for _ in range(L)]208    for b in blocks:209        b.set_dtype(mx.float16)210    x = mx.random.normal((BSZ, SEQ, D_MODEL)).astype(mx.float16)211    store = np.memmap(workdir / "mlx_capture.raw", dtype=np.float16, mode="w+",212                      shape=(N_FWD_WRITE * L * BSZ * SEQ, D_MODEL))213214    def fwd(capture):215        captured, h = [], x216        for b in blocks:217            h = b(h)218            if capture:219                captured.append(h)220        return h, captured221222    def run(n, capture, to_disk):223        row = 0224        for _ in range(WARMUP):225            out, cap = fwd(capture)226            mx.eval(out, *cap)227        t0 = time.perf_counter()228        for _ in range(n):229            out, cap = fwd(capture)230            mx.eval(out, *cap)231            if to_disk:232                for c in cap:233                    a = np.array(c, copy=False).reshape(-1, D_MODEL)234                    store[row:row + a.shape[0]] = a235                    row += a.shape[0]236        return (time.perf_counter() - t0) / n237238    out = []239    for mode, cap, disk, n in (("plain", False, False, N_FWD),240                               ("retain", True, False, N_FWD),241                               ("retain+copy+write", True, True, N_FWD_WRITE)):242        times = [run(n, cap, disk) for _ in range(REPEATS)]243        out.append({"backend": "mlx", "mode": mode,244                    "s_per_forward": times,245                    "tokens_per_s_mean": BSZ * SEQ / np.mean(times)})246        print(f"  mlx       {mode:22s} {np.mean(times)*1000:8.1f} ms/fwd")247    return out248249250# ---------------------------------------------------------------- main251def main() -> int:252    workdir = Path(tempfile.mkdtemp(prefix="modelmap_expH_"))253    print(f"expH run #1 — workdir {workdir}")254    try:255        print("Part A — storage formats (warm cache):")256        storage = bench_storage(workdir)257        print("Part B — capture overhead:")258        compute = bench_torch_mps(workdir) + bench_mlx(workdir)259    finally:260        shutil.rmtree(workdir, ignore_errors=True)261262    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,263                            capture_output=True, text=True, check=False).stdout.strip()264    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())265    outdir = ROOT / "results" / "expH_capture_cost_frontier" / ts266    outdir.mkdir(parents=True)267    doc = {268        "experiment": "expH_capture_cost_frontier",269        "run": 1,270        "scope": "storage formats (warm cache) + capture overhead, synthetic model",271        "commit": commit,272        "config": {273            "storage": {"rows": ROWS, "dim": DIM, "write_chunk": WRITE_CHUNK,274                        "batch": BATCH, "n_batches": N_BATCHES, "repeats": REPEATS,275                        "cache": "warm (declared limitation; cold pass = run #2)"},276            "compute": {"layers": L, "d_model": D_MODEL, "heads": N_HEADS,277                        "seq": SEQ, "batch": BSZ, "dtype": "float16",278                        "n_forwards": N_FWD, "warmup": WARMUP, "repeats": REPEATS},279            "seed": SEED,280        },281        "manifest": manifest(),282        "storage": storage,283        "compute": compute,284    }285    (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n")286    print(f"results -> {outdir / 'results.json'}")287    return 0288289290if __name__ == "__main__":291    sys.exit(main())292