spb/localvm-research Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : localvm-research4#  File      : experiments/micro/expA_weight_concentration/benchmark.py5#  Purpose   : Per-token FFN block-energy concentration (trace shared with expB)6#  Author    : Simon-Pierre Boucher7#  Contact   : contact@spboucher.ai8#  Created   : 2026-08-129#  Modified  : 2026-08-1210#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal11#  License   : All rights reserved (research code)12# =============================================================================13"""Experiment A — weight contribution concentration (charter §9.A).1415Records SwiGLU intermediate-activation energy per 64-neuron block, per token,16per layer, on the bf16 reference model. Writes concentration aggregates to17results.json and the raw block-energy trace (npz) for expB.1819Usage:20    .venv/bin/python benchmark.py [--per-domain 8] [--gen-tokens 128] [--block 64]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  # noqa: E402424344class DownProjRecorder(nn.Module):45    """Wraps a down_proj Linear; records block energy of its input (the SwiGLU46    intermediate activation) for positions in `window`."""4748    def __init__(self, inner: nn.Module, block: int):49        super().__init__()50        self.inner = inner51        self.block = block52        self.window: tuple[int, int] | None = None53        self.block_energy: np.ndarray | None = None54        self.neuron_energy: np.ndarray | None = None5556    def __call__(self, x):57        if self.window is not None:58            a, b = self.window59            h = x[0, a:b].astype(mx.float32)60            e = mx.square(h)61            be = e.reshape(h.shape[0], -1, self.block).sum(axis=-1)62            # normalize per position: raw energies overflow float16 storage,63            # and only relative importance matters for expA/expB64            be = be / (be.sum(axis=-1, keepdims=True) + 1e-12)65            mx.eval(be)66            self.block_energy = np.array(be)67            self.neuron_energy = np.array(e)  # (T, D_int) — reduced by caller68        return self.inner(x)697071def concentration_stats(energy: np.ndarray, fracs=(0.1, 0.2, 0.4, 0.6),72                        targets=(0.90, 0.95, 0.99)) -> dict:73    """energy: (T, N). Returns mean energy captured by top-f fraction and mean74    fraction of units needed to reach target energy."""75    T, N = energy.shape76    srt = np.sort(energy, axis=1)[:, ::-1]77    csum = np.cumsum(srt, axis=1)78    total = csum[:, -1:] + 1e-1279    frac_captured = {}80    for f in fracs:81        k = max(1, int(round(f * N)))82        frac_captured[f] = float(np.mean(csum[:, k - 1] / total[:, 0]))83    needed = {}84    ratio = csum / total85    for t in targets:86        idx = np.argmax(ratio >= t, axis=1) + 187        needed[t] = float(np.mean(idx / N))88    return {"top_frac_energy": frac_captured, "frac_needed_for": needed}899091def main() -> None:92    ap = argparse.ArgumentParser()93    ap.add_argument("--model", default="mlx-community/Qwen3-1.7B-bf16")94    ap.add_argument("--gen-tokens", type=int, default=128)95    ap.add_argument("--per-domain", type=int, default=8)96    ap.add_argument("--block", type=int, default=16,97                    help="trace granularity; analysis also derives 4x-coarser blocks")98    args = ap.parse_args()99100    domains = json.loads((REPO_ROOT / "benchmarks/datasets/eval_prompts.json").read_text())["domains"]101    print(f"loading {args.model} …", flush=True)102    model, tokenizer = load(args.model)103    layers = model.model.layers104    n_layers = len(layers)105106    recorders = []107    for layer in layers:108        rec = DownProjRecorder(layer.mlp.down_proj, args.block)109        layer.mlp.down_proj = rec110        recorders.append(rec)111112    trajectories = []113    t0 = time.time()114    for domain, plist in domains.items():115        for prompt in plist[: args.per_domain]:116            ids = tokenizer.apply_chat_template(117                [{"role": "user", "content": prompt}], add_generation_prompt=True)118            for r in recorders:119                r.window = None  # no recording during generation120            gen = greedy_generate(model, tokenizer, ids, args.gen_tokens)121            if len(gen) >= 8:122                trajectories.append({"domain": domain, "full_ids": list(ids) + gen, "start": len(ids)})123    print(f"{len(trajectories)} trajectories in {time.time()-t0:.0f}s", flush=True)124125    block_traces = []   # per traj: (T, n_layers, n_blocks) f16126    neuron_stats = []   # per traj per layer concentration dicts127    index = []128    for ti, t in enumerate(trajectories):129        a, b = t["start"] - 1, len(t["full_ids"]) - 1130        for r in recorders:131            r.window = (a, b)132        model(mx.array(t["full_ids"])[None])133        per_layer_blocks = np.stack([r.block_energy for r in recorders], axis=1)  # (T, L, B)134        block_traces.append(per_layer_blocks.astype(np.float16))135        neuron_stats.append([concentration_stats(r.neuron_energy) for r in recorders])136        for r in recorders:137            r.neuron_energy = None138        index.append({"traj": ti, "domain": t["domain"], "n_pos": b - a})139        if (ti + 1) % 12 == 0:140            print(f"  traced {ti+1}/{len(trajectories)}", flush=True)141142    all_blocks = np.concatenate(block_traces, axis=0)  # (P, L, B)143    P, L, B = all_blocks.shape144    print(f"trace shape {all_blocks.shape}", flush=True)145146    # expA aggregates at trace granularity and 4x-coarser derived granularity147    coarse = all_blocks.reshape(P, L, B // 4, 4).astype(np.float32).sum(axis=-1)148    per_layer = [concentration_stats(all_blocks[:, li, :].astype(np.float32)) for li in range(L)]149    overall = concentration_stats(all_blocks.reshape(P * L, B).astype(np.float32))150    overall_coarse = concentration_stats(coarse.reshape(P * L, B // 4))151    per_domain = {}152    pos_domain = np.concatenate([[ix["domain"]] * ix["n_pos"] for ix in index])153    for dom in domains:154        sel = all_blocks[pos_domain == dom]155        per_domain[dom] = concentration_stats(sel.reshape(-1, B).astype(np.float32))156    # neuron-granularity mean across trajectories/layers157    neuron_overall = {158        "top_frac_energy": {f: float(np.mean([s["top_frac_energy"][f] for ns in neuron_stats for s in ns]))159                            for f in (0.1, 0.2, 0.4, 0.6)},160        "frac_needed_for": {t: float(np.mean([s["frac_needed_for"][t] for ns in neuron_stats for s in ns]))161                            for t in (0.90, 0.95, 0.99)},162    }163164    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")165    out_dir = REPO_ROOT / "results" / "expA_weight_concentration" / ts166    out_dir.mkdir(parents=True)167    np.savez_compressed(out_dir / "block_energy_trace.npz",168                        blocks=all_blocks,169                        domains=pos_domain,170                        traj_id=np.concatenate([[ix["traj"]] * ix["n_pos"] for ix in index]))171    (out_dir / "results.json").write_text(json.dumps({172        "experiment": "expA_weight_concentration",173        "author": "Simon-Pierre Boucher",174        "contact": "contact@spboucher.ai",175        "manifest": collect_manifest(),176        "config": vars(args),177        "n_positions": int(P), "n_layers": int(L), "n_blocks": int(B),178        "coarse_block_granularity": {"block_size": args.block * 4, "overall": overall_coarse},179        "block_granularity": {"overall": overall,180                              "per_layer": {str(i): s for i, s in enumerate(per_layer)},181                              "per_domain": per_domain},182        "neuron_granularity": neuron_overall,183    }, indent=2, default=float))184    print(f"\nwrote {out_dir}/results.json (+ block_energy_trace.npz for expB)")185    print("overall block-64:", json.dumps(overall, default=float))186    print("neuron-level    :", json.dumps(neuron_overall, default=float))187188189if __name__ == "__main__":190    main()191