Interpretability twin: load the deployed model.bin into brainscope

brainscope_adapter/README.md

@@ -0,0 +1,38 @@
+# The ESP32's model, under brainscope
+
+Loads the exact model this repo flashes to an ESP32-S3 into
+[brainscope](https://github.com/moudrkat/brainscope) - logit lens, attention,
+per-layer activity, live. Not the training checkpoint: the int4 weights are
+dequantized straight out of `artifacts/tinystories/model.bin`, the same bytes
+the board mmaps from flash. `verify_vs_c.py` proves the twin against the C
+runtime (`runtime/llm.h`) that ships to the device - last-position logits agree
+to ~1e-5, the same fp32-from-int4 path `verify.c` gates before flashing.
+
+brainscope itself is untouched: the PLE architecture is registered with
+transformers in-process and served through brainscope's public CLI.
+
+```bash
+scripts/fetch_model.sh tinystories                      # if artifacts/ is empty
+$BRAINSCOPE_PY brainscope_adapter/build_hf.py           # model.bin -> hf twin
+$BRAINSCOPE_PY brainscope_adapter/verify_vs_c.py        # gate vs the C runtime
+$BRAINSCOPE_PY brainscope_adapter/serve.py              # brainscope on :8010
+```
+
+`$BRAINSCOPE_PY` is any python with `torch`, `transformers` and brainscope
+importable. `serve.py` looks for a brainscope checkout at
+`~/projekty/brainscope`; with brainscope pip-installed, the path insert is
+simply unused.
+
+The model is not a chat model - it continues text. The tokenizer's chat
+template therefore concatenates all message contents verbatim, which turns
+brainscope's chat box into a continue-the-story box: type an opening, watch
+6 layers x 96 dims write the rest at full visibility.
+
+| file | role |
+|---|---|
+| `ple_bin.py` | parse + dequantize `model.bin` (int4 groups, fp16 scales) |
+| `ple_hf.py` | the PLE architecture in brainscope's expected skeleton |
+| `build_hf.py` | write `artifacts/tinystories/hf/` + smoke sample |
+| `dump_logits.c` | C-runtime logits for a prompt, via `runtime/llm.h` |
+| `verify_vs_c.py` | twin-vs-C gate + KV-cache parity |
+| `serve.py` | register the architecture, hand over to brainscope |

brainscope_adapter/build_hf.py

@@ -0,0 +1,101 @@
+"""Build a transformers-loadable twin of the deployed ESP32 model.
+
+Reads artifacts/tinystories/{model.bin,tokenizer.json} - the exact files
+deploy.sh flashes - and writes artifacts/tinystories/hf/ with the dequantized
+fp32 weights in the brainscope-compatible skeleton. Ends with a short greedy
+sample as a smoke test.
+
+Run from the repo root with a python that has torch + transformers:
+    python brainscope_adapter/build_hf.py [--artifacts artifacts/tinystories]
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+from ple_bin import load_model_bin
+from ple_hf import PLEConfig, PLETinyLMForCausalLM, bin_to_hf_key
+
+
+def build(artifacts: Path):
+    cfg_bin, sd_bin = load_model_bin(artifacts / "model.bin")
+    print(f"model.bin: Vin={cfg_bin.vocab_size} Vout={cfg_bin.out_vocab} "
+          f"D={cfg_bin.d_model} L={cfg_bin.n_layers} H={cfg_bin.n_heads} "
+          f"F={cfg_bin.ffn_hidden} P={cfg_bin.ple_dim} seq={cfg_bin.seq_len}")
+    if not cfg_bin.tied_head:
+        raise SystemExit("model.bin is untied; this builder handles the tied-head layout")
+
+    cfg = PLEConfig(
+        vocab_size=cfg_bin.vocab_size, out_vocab=cfg_bin.out_vocab,
+        hidden_size=cfg_bin.d_model, num_hidden_layers=cfg_bin.n_layers,
+        num_attention_heads=cfg_bin.n_heads, ffn_hidden=cfg_bin.ffn_hidden,
+        ple_dim=cfg_bin.ple_dim, seq_len=cfg_bin.seq_len,
+        rope_theta=cfg_bin.rope_theta,
+        eos_token_id=0, pad_token_id=0,
+    )
+    model = PLETinyLMForCausalLM(cfg)
+    sd = {bin_to_hf_key(k): torch.from_numpy(v) for k, v in sd_bin.items()}
+    # Tied head: the first out_vocab rows of the (dequantized) embedding.
+    sd["lm_head.weight"] = sd["model.embed_tokens.weight"][: cfg.out_vocab].clone()
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    missing = [m for m in missing if not m.endswith((".cos", ".sin"))]
+    if missing or unexpected:
+        raise SystemExit(f"state dict mismatch: missing={missing} unexpected={unexpected}")
+    model.eval()
+    return model
+
+
+def save(model, artifacts: Path, out: Path):
+    from transformers import PreTrainedTokenizerFast
+
+    out.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(out)
+    tok = PreTrainedTokenizerFast(
+        tokenizer_file=str(artifacts / "tokenizer.json"),
+        eos_token="<|endoftext|>", pad_token="<|endoftext|>")
+    # Not a chat model: the "chat" is the story so far, concatenated verbatim,
+    # so brainscope's chat box behaves as a continue-the-story box.
+    tok.chat_template = "{% for message in messages %}{{ message['content'] }}{% endfor %}"
+    tok.save_pretrained(out)
+    print(f"wrote {out}")
+    return tok
+
+
+def sample(model, tok, prompt="Once upon a time", n=40):
+    ids = tok(prompt, return_tensors="pt").input_ids
+    past = None
+    out_ids = []
+    with torch.no_grad():
+        feed = ids
+        for _ in range(n):
+            o = model(input_ids=feed, past_key_values=past, use_cache=True)
+            past = o.past_key_values
+            nxt = int(o.logits[0, -1].argmax())
+            if nxt == 0:
+                break
+            out_ids.append(nxt)
+            feed = torch.tensor([[nxt]])
+    print(f"smoke sample: {prompt!r} -> {tok.decode(out_ids)!r}")
+
+
+if __name__ == "__main__":
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--artifacts", type=Path,
+                    default=Path(__file__).resolve().parents[1] / "artifacts" / "tinystories")
+    ap.add_argument("--out", type=Path, default=None,
+                    help="output dir (default: <artifacts>/hf)")
+    ap.add_argument("--zero-ple-table", action="store_true",
+                    help="write the twin with the flash-resident 25M-param PLE "
+                         "table zeroed - the 'flash unplugged' ablation")
+    args = ap.parse_args()
+    model = build(args.artifacts)
+    if args.zero_ple_table:
+        with torch.no_grad():
+            model.model.ple_table.weight.zero_()
+        print("PLE table zeroed: the flash-resident parameters are unplugged")
+    out = args.out or args.artifacts / "hf"
+    tok = save(model, args.artifacts, out)
+    sample(model, tok)

brainscope_adapter/dump_logits.c

@@ -0,0 +1,52 @@
+// Print the C runtime's last-position logits for a prompt given as token ids.
+// The brainscope twin (build_hf.py) must reproduce these numbers; verify_vs_c.py
+// runs both sides and diffs them. Reuses runtime/llm.h unmodified - the same
+// portable inference verify.c gates before anything touches the board.
+//
+//   cc -O3 -Wall -Wextra -I runtime -o dump_logits
+//     brainscope_adapter/dump_logits.c -lm
+//   ./dump_logits artifacts/tinystories/model.bin 1 500 1000 200 42 777 13 99
+#include <stdio.h>
+#include <stdlib.h>
+#include "llm.h"
+
+static uint8_t *read_file(const char *path, size_t *n) {
+  FILE *f = fopen(path, "rb");
+  if (!f) { perror(path); exit(1); }
+  fseek(f, 0, SEEK_END); *n = ftell(f); fseek(f, 0, SEEK_SET);
+  uint8_t *b = malloc(*n);
+  if (fread(b, 1, *n, f) != *n) { fprintf(stderr, "short read\n"); exit(1); }
+  fclose(f); return b;
+}
+
+int main(int argc, char **argv) {
+  if (argc < 3) {
+    fprintf(stderr, "usage: %s <model.bin> <token id> [id ...]\n", argv[0]);
+    return 2;
+  }
+  size_t n;
+  uint8_t *buf = read_file(argv[1], &n);
+  Model m;
+  if (llm_load(buf, &m)) { fprintf(stderr, "bad magic\n"); return 1; }
+
+  int D = m.c.dim, L = m.c.n_layers, P = m.c.ple_dim, F = m.c.ffn,
+      V = m.out_vocab, S = m.c.seq_len;
+  Scratch s;
+  s.x = malloc(D * 4); s.h = malloc((F > D ? F : D) * 4);
+  s.qkv = malloc(3 * D * 4); s.att = malloc(D * 4);
+  s.g1 = malloc(F * 4); s.g2 = malloc((P > F ? P : F) * 4);
+  s.ple = malloc(L * P * 4); s.tmpP = malloc(L * P * 4); s.trow = malloc(L * P * 4);
+  s.logits = malloc(V * 4);
+  s.scores = malloc(S * 4);
+  s.kcache = malloc((size_t)L * S * D * 4);
+  s.vcache = malloc((size_t)L * S * D * 4);
+
+  int plen = argc - 2;
+  for (int i = 0; i < plen; i++) {
+    int id = atoi(argv[2 + i]);
+    if (id < 0 || id >= m.c.vocab) { fprintf(stderr, "id %d out of range\n", id); return 1; }
+    llm_forward(&m, id, i, &s);
+  }
+  for (int i = 0; i < V; i++) printf("%.6f\n", s.logits[i]);
+  return 0;
+}

brainscope_adapter/extract_direction.py

@@ -0,0 +1,40 @@
+"""Extract a steering direction for the ESP32 twin - brainscope untouched.
+
+Same trick as serve.py: register the PLE architecture, then hand over to
+brainscope's own extract CLI. Defaults bake in the demo: the "dark" story-mood
+direction from mood_pairs.jsonl at layer 3 (of 6), written next to the twin.
+
+    ~/projekty/brainscope/.venv/bin/python brainscope_adapter/extract_direction.py
+    ~/projekty/brainscope/.venv/bin/python brainscope_adapter/serve.py \
+        --directions artifacts/tinystories/dirs.json
+"""
+
+import sys
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+ROOT = HERE.parent
+BRAINSCOPE_REPO = Path.home() / "projekty" / "brainscope"
+
+sys.path.insert(0, str(HERE))
+if BRAINSCOPE_REPO.is_dir():
+    sys.path.insert(0, str(BRAINSCOPE_REPO))
+
+import ple_hf  # noqa: F401  (registers ple-tinylm with transformers)
+
+from brainscope import extract  # noqa: E402
+
+defaults = {
+    "--model": str(ROOT / "artifacts" / "tinystories" / "hf"),
+    "--pairs": str(HERE / "mood_pairs.jsonl"),
+    "--layer": "3",
+    "--name": "dark",
+    "--out": str(ROOT / "artifacts" / "tinystories" / "dirs.json"),
+}
+extra = sys.argv[1:]
+argv = ["extract"]
+for flag, value in defaults.items():
+    if flag not in extra:
+        argv += [flag, value]
+sys.argv = argv + extra
+extract.main()

brainscope_adapter/mood_pairs.jsonl

@@ -0,0 +1,10 @@
+{"positive": "The sky turned dark and a cold wind howled through the trees. Lily was scared and started to cry.", "negative": "The sun was shining and the birds sang sweetly. Lily laughed and clapped her hands."}
+{"positive": "The little dog was lost in the dark forest. He was cold, sad and all alone.", "negative": "The little dog played in the warm garden. He was happy and wagged his tail."}
+{"positive": "Tom looked at the broken toy and big tears ran down his face. Everything was ruined.", "negative": "Tom looked at his new toy and smiled. It was the best day ever."}
+{"positive": "The storm came and the thunder was very loud. Everyone hid and shivered with fear.", "negative": "The rain stopped and a pretty rainbow came out. Everyone danced and cheered."}
+{"positive": "The old house was dark and quiet. Something moved in the shadows and Anna froze.", "negative": "The little house was bright and cozy. The fire was warm and Anna felt safe."}
+{"positive": "Ben dropped his ice cream in the mud. He sat down and cried and cried.", "negative": "Ben got a big ice cream with a cherry on top. He jumped with joy."}
+{"positive": "The bird's wing was hurt and it could not fly. It sat in the cold rain, sad and weak.", "negative": "The bird spread its wings and flew high in the blue sky. It sang a happy song."}
+{"positive": "Nobody came to Mia's party. She sat alone in the empty room and felt very sad.", "negative": "All her friends came to Mia's party. They ate cake and played games and laughed."}
+{"positive": "The night was black and cold. A strange noise came closer and closer.", "negative": "The morning was bright and warm. A friendly puppy came running to say hello."}
+{"positive": "Sam lost his way home and the woods grew darker. He was afraid and wanted to cry.", "negative": "Sam found his way home and mom hugged him tight. He felt warm and loved."}

brainscope_adapter/ple_bin.py

@@ -0,0 +1,110 @@
+"""Parse artifacts/<model>/model.bin back into fp32 tensors.
+
+The binary is the exporter's int4 group-quantized format (research/tinystories/
+export.py). Dequantizing here reproduces exactly the weights the C runtime
+reconstructs on the device: codes are (q+8) nibbles, scales are fp16, groups of
+`group` along the last dim, ragged tail. So the model this yields is not "the
+checkpoint" - it is the model the ESP32 actually runs, bit-faithful.
+"""
+
+import struct
+from dataclasses import dataclass
+
+import numpy as np
+
+MAGIC = 0x00454C50  # "PLE\0"
+
+
+@dataclass
+class BinConfig:
+    vocab_size: int
+    out_vocab: int
+    d_model: int
+    n_layers: int
+    n_heads: int
+    ffn_hidden: int
+    ple_dim: int
+    seq_len: int
+    group: int
+    rope_theta: float
+    tied_head: bool
+
+
+def _dequant(buf, off, shape, group):
+    """Read one packed int4 tensor; returns (fp32 ndarray, new offset)."""
+    (tensor_group,) = struct.unpack_from("<i", buf, off)
+    off += 4
+    assert tensor_group == group, f"group {tensor_group} != header {group}"
+    cols = shape[-1]
+    rows = int(np.prod(shape)) // cols
+    row_bytes = (cols + 1) // 2
+    n_groups = (cols + group - 1) // group
+
+    packed = np.frombuffer(buf, np.uint8, rows * row_bytes, off).reshape(rows, row_bytes)
+    off += rows * row_bytes
+    scales = np.frombuffer(buf, np.float16, rows * n_groups, off).reshape(rows, n_groups)
+    off += rows * n_groups * 2
+
+    codes = np.empty((rows, cols), np.int8)
+    lo = (packed & 0x0F).astype(np.int8)
+    hi = (packed >> 4).astype(np.int8)
+    codes[:, 0::2] = lo[:, : (cols + 1) // 2]
+    codes[:, 1::2] = hi[:, : cols // 2]
+    q = codes.astype(np.float32) - 8.0
+
+    sc = np.repeat(scales.astype(np.float32), group, axis=1)[:, :cols]
+    return (q * sc).reshape(shape), off
+
+
+def _fp32(buf, off, shape):
+    n = int(np.prod(shape))
+    arr = np.frombuffer(buf, np.float32, n, off).reshape(shape).copy()
+    return arr, off + n * 4
+
+
+def load_model_bin(path):
+    """Returns (BinConfig, {name: fp32 ndarray}) with the exporter's names."""
+    buf = open(path, "rb").read()
+    magic, version, header_bytes, flags = struct.unpack_from("<IIII", buf, 0)
+    if magic != MAGIC:
+        raise ValueError(f"{path}: bad magic {magic:#x}")
+    if version != 1:
+        raise ValueError(f"{path}: unsupported format version {version}")
+    vocab, out_vocab = struct.unpack_from("<II", buf, 16)
+    d, L, H, F, P, S, G = struct.unpack_from("<7i", buf, 24)
+    (theta,) = struct.unpack_from("<f", buf, 52)
+    cfg = BinConfig(vocab, out_vocab, d, L, H, F, P, S, G, theta, bool(flags & 1))
+
+    # Same fixed order the exporter writes and the C reader hard-codes.
+    plan = [
+        ("tok_emb.weight", (vocab, d), True),
+        ("ple_model_proj.weight", (L * P, d), True),
+        ("ple_proj_norm.weight", (P,), False),
+        ("ple_table.weight", (vocab, L * P), True),
+    ]
+    for i in range(L):
+        p = f"blocks.{i}."
+        plan += [
+            (p + "attn_norm.weight", (d,), False),
+            (p + "attn.qkv.weight", (3 * d, d), True),
+            (p + "attn.proj.weight", (d, d), True),
+            (p + "ffn_norm.weight", (d,), False),
+            (p + "ffn.gate.weight", (F, d), True),
+            (p + "ffn.up.weight", (F, d), True),
+            (p + "ffn.down.weight", (d, F), True),
+            (p + "ple_gate.weight", (P, d), True),
+            (p + "ple_proj.weight", (d, P), True),
+            (p + "ple_norm.weight", (d,), False),
+        ]
+    plan.append(("out_norm.weight", (d,), False))
+
+    off = header_bytes
+    sd = {}
+    for name, shape, quant in plan:
+        if quant:
+            sd[name], off = _dequant(buf, off, shape, G)
+        else:
+            sd[name], off = _fp32(buf, off, shape)
+    if off != len(buf):
+        raise ValueError(f"{path}: {len(buf) - off} trailing bytes after last tensor")
+    return cfg, sd