diff --git a/brainscope_adapter/README.md b/brainscope_adapter/README.md new file mode 100644 index 0000000..b9a3ad6 --- /dev/null +++ b/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 | diff --git a/brainscope_adapter/build_hf.py b/brainscope_adapter/build_hf.py new file mode 100644 index 0000000..6231226 --- /dev/null +++ b/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: /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) diff --git a/brainscope_adapter/dump_logits b/brainscope_adapter/dump_logits new file mode 100755 index 0000000..627b62a Binary files /dev/null and b/brainscope_adapter/dump_logits differ diff --git a/brainscope_adapter/dump_logits.c b/brainscope_adapter/dump_logits.c new file mode 100644 index 0000000..4e47c96 --- /dev/null +++ b/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 +#include +#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 [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; +} diff --git a/brainscope_adapter/extract_direction.py b/brainscope_adapter/extract_direction.py new file mode 100644 index 0000000..b1a2110 --- /dev/null +++ b/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() diff --git a/brainscope_adapter/mood_pairs.jsonl b/brainscope_adapter/mood_pairs.jsonl new file mode 100644 index 0000000..25d7fee --- /dev/null +++ b/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."} diff --git a/brainscope_adapter/ple_bin.py b/brainscope_adapter/ple_bin.py new file mode 100644 index 0000000..d81f8b9 --- /dev/null +++ b/brainscope_adapter/ple_bin.py @@ -0,0 +1,110 @@ +"""Parse artifacts//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("> 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(" q_pos[:, None] # future positions + scores = scores.masked_fill(mask[None, None], float("-inf")) + weights = torch.softmax(scores, dim=-1) + out = (weights @ v).transpose(1, 2).reshape(B, T, C) + return self.o_proj(out), (k, v), (weights if output_attentions else None) + + +class PLEMLP(nn.Module): + def __init__(self, cfg: PLEConfig): + super().__init__() + self.gate_proj = nn.Linear(cfg.hidden_size, cfg.ffn_hidden, bias=False) + self.up_proj = nn.Linear(cfg.hidden_size, cfg.ffn_hidden, bias=False) + self.down_proj = nn.Linear(cfg.ffn_hidden, cfg.hidden_size, bias=False) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class PLEDecoderLayer(nn.Module): + def __init__(self, cfg: PLEConfig): + super().__init__() + self.input_layernorm = RMSNorm(cfg.hidden_size) + self.self_attn = PLEAttention(cfg) + self.post_attention_layernorm = RMSNorm(cfg.hidden_size) + self.mlp = PLEMLP(cfg) + self.ple_gate = nn.Linear(cfg.hidden_size, cfg.ple_dim, bias=False) + self.ple_proj = nn.Linear(cfg.ple_dim, cfg.hidden_size, bias=False) + self.ple_norm = RMSNorm(cfg.hidden_size) + + def forward(self, x, cos, sin, ple, past=None, output_attentions=False): + a, kv, w = self.self_attn(self.input_layernorm(x), cos, sin, past, output_attentions) + x = x + a + x = x + self.mlp(self.post_attention_layernorm(x)) + g = F.gelu(self.ple_gate(x)) + x = x + self.ple_norm(self.ple_proj(g * ple)) + return x, kv, w + + +class PLEModel(nn.Module): + def __init__(self, cfg: PLEConfig): + super().__init__() + self.cfg = cfg + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size) + self.ple_model_proj = nn.Linear(cfg.hidden_size, cfg.num_hidden_layers * cfg.ple_dim, bias=False) + self.ple_proj_norm = RMSNorm(cfg.ple_dim) + self.ple_table = nn.Embedding(cfg.vocab_size, cfg.num_hidden_layers * cfg.ple_dim) + self.layers = nn.ModuleList(PLEDecoderLayer(cfg) for _ in range(cfg.num_hidden_layers)) + self.norm = RMSNorm(cfg.hidden_size) + + inv = 1.0 / (cfg.rope_theta ** (torch.arange(0, cfg.head_dim, 2).float() / cfg.head_dim)) + t = torch.arange(cfg.max_position_embeddings).float() + freqs = torch.outer(t, inv) + # persistent=True on purpose: from_pretrained fast-inits on the meta + # device, so a non-persistent buffer would come back as uninitialized + # memory. Shipping the tables in the checkpoint sidesteps that. + self.register_buffer("cos", freqs.cos(), persistent=True) + self.register_buffer("sin", freqs.sin(), persistent=True) + + +class PLETinyLMForCausalLM(PreTrainedModel): + config_class = PLEConfig + base_model_prefix = "model" + _no_split_modules = ["PLEDecoderLayer"] + + def __init__(self, config: PLEConfig): + super().__init__(config) + self.model = PLEModel(config) + # On the device the head IS the first out_vocab rows of the embedding + # (tied, scanned once per token from PSRAM). Held as its own tensor here + # so get_output_embeddings/logit lens see a plain Linear. + self.lm_head = nn.Linear(config.hidden_size, config.out_vocab, bias=False) + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def get_output_embeddings(self): + return self.lm_head + + def forward(self, input_ids=None, past_key_values=None, use_cache=False, + output_hidden_states=False, output_attentions=False, + attention_mask=None, **kwargs): + cfg = self.config + core = self.model + B, T = input_ids.shape + past_len = past_key_values[0][0].shape[2] if past_key_values else 0 + cos = core.cos[past_len:past_len + T] + sin = core.sin[past_len:past_len + T] + + x = core.embed_tokens(input_ids) + L, P = cfg.num_hidden_layers, cfg.ple_dim + ple = core.ple_model_proj(x) * (cfg.hidden_size**-0.5) + ple = core.ple_proj_norm(ple.view(B, T, L, P)) + table = core.ple_table(input_ids).view(B, T, L, P) + # embed_scale sqrt(P) on the table, then average the two halves - the + # undocumented-but-load-bearing scaling from Gemma, kept from training. + ple = (ple + table * (P**0.5)) * (2**-0.5) + + hidden = [x] if output_hidden_states else None + new_past = [] if use_cache else None + attns = [] if output_attentions else None + for i, layer in enumerate(core.layers): + past = past_key_values[i] if past_key_values else None + x, kv, w = layer(x, cos, sin, ple[:, :, i], past, output_attentions) + if output_hidden_states: + hidden.append(x) + if use_cache: + new_past.append(kv) + if output_attentions: + attns.append(w) + + logits = self.lm_head(core.norm(x)) + return CausalLMOutputWithPast( + logits=logits, + past_key_values=tuple(new_past) if use_cache else None, + hidden_states=tuple(hidden) if output_hidden_states else None, + attentions=tuple(attns) if output_attentions else None, + ) + + +AutoConfig.register("ple-tinylm", PLEConfig) +AutoModelForCausalLM.register(PLEConfig, PLETinyLMForCausalLM) + + +# exporter tensor name -> wrapper tensor name +def bin_to_hf_key(name: str) -> str: + out = (name + .replace("tok_emb.weight", "embed_tokens.weight") + .replace("out_norm.weight", "norm.weight")) + if out.startswith("blocks."): + out = (out.replace("blocks.", "layers.") + .replace(".attn_norm.", ".input_layernorm.") + .replace(".attn.qkv.", ".self_attn.qkv.") + .replace(".attn.proj.", ".self_attn.o_proj.") + .replace(".ffn_norm.", ".post_attention_layernorm.") + .replace(".ffn.gate.", ".mlp.gate_proj.") + .replace(".ffn.up.", ".mlp.up_proj.") + .replace(".ffn.down.", ".mlp.down_proj.")) + return "model." + out diff --git a/brainscope_adapter/serve.py b/brainscope_adapter/serve.py new file mode 100644 index 0000000..2830a63 --- /dev/null +++ b/brainscope_adapter/serve.py @@ -0,0 +1,38 @@ +"""Serve the ESP32's deployed model in brainscope, untouched brainscope. + +Registers the PLE architecture with transformers in-process (import ple_hf) and +then hands control to brainscope's own CLI, pointed at the twin built by +build_hf.py. Any extra arguments pass straight through to brainscope +(--port, --lens, --no-browser, ...). + +Run with brainscope's environment: + ~/projekty/brainscope/.venv/bin/python brainscope_adapter/serve.py +""" + +import os +import sys +from pathlib import Path + +HERE = Path(__file__).resolve().parent +ROOT = HERE.parent +BRAINSCOPE_REPO = Path.home() / "projekty" / "brainscope" +# PLE_HF_DIR selects which twin to serve, e.g. the --zero-ple-table build. +HF_DIR = Path(os.environ.get("PLE_HF_DIR", ROOT / "artifacts" / "tinystories" / "hf")) + +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) + +if not HF_DIR.is_dir(): + raise SystemExit(f"{HF_DIR} missing - run: python brainscope_adapter/build_hf.py") + +from brainscope import server # noqa: E402 + +extra = sys.argv[1:] +argv = ["brainscope", "--model", str(HF_DIR)] +if "--lens" not in extra: + argv += ["--lens", "on"] # 6 layers x 96 dims: the lens is free, keep it on +sys.argv = argv + extra +server.main() diff --git a/brainscope_adapter/tell_box.py b/brainscope_adapter/tell_box.py new file mode 100644 index 0000000..e416c0b --- /dev/null +++ b/brainscope_adapter/tell_box.py @@ -0,0 +1,65 @@ +"""Send one prompt to the matchbox AND to brainscope, in lockstep. + +The board encodes the prompt itself and decodes greedily on-chip; brainscope's +twin is bit-faithful to those weights, so fed the same prompt at temperature 0 +it writes the SAME story - the OLED shows the words, the open brainscope tab +shows the layers producing them. + + sg dialout -c 'python brainscope_adapter/tell_box.py "One day, a little cat"' + +The board listens between stories; if it is mid-story, the prompt waits in its +serial buffer until the current one finishes (up to ~45 s). +""" + +import argparse +import subprocess +import sys +import urllib.request +import json + +PORT = "/dev/ttyACM0" +BRAINSCOPE = "http://localhost:8010/v1/chat/completions" +UNPLUGGED = "http://localhost:8011/v1/chat/completions" + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("prompt", help="ASCII story opening, e.g. 'Once upon a time'") + ap.add_argument("--port", default=PORT) + ap.add_argument("--max-tokens", type=int, default=200, + help="brainscope-side cap; the board itself writes 200") + ap.add_argument("--also-unplugged", action="store_true", + help="feed the same prompt to the flash-unplugged twin on " + ":8011 too - three minds, one prompt") + args = ap.parse_args() + if not args.prompt.isascii(): + sys.exit("the device tokenizer is ASCII-only (no diacritics)") + + subprocess.run(["stty", "-F", args.port, "115200", "raw", "-echo"], check=True) + with open(args.port, "wb", buffering=0) as ser: + ser.write(args.prompt.encode("ascii") + b"\n") + print(f"box <- {args.prompt!r}") + + def ask(url): + req = urllib.request.Request( + url, + data=json.dumps({ + "messages": [{"role": "user", "content": args.prompt}], + "max_tokens": args.max_tokens, + "temperature": 0, + }).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=600) as r: + return json.load(r)["choices"][0]["message"]["content"] + + print("brainscope: generating (watch the open tab)...") + text = ask(BRAINSCOPE) + print(f"twin story: {text[:160]}{'...' if len(text) > 160 else ''}") + if args.also_unplugged: + broken = ask(UNPLUGGED) + print(f"unplugged : {broken[:80]!r}...") + print("the OLED should be writing the same words as the twin, ~10 tok/s.") + + +if __name__ == "__main__": + main() diff --git a/brainscope_adapter/verify_vs_c.py b/brainscope_adapter/verify_vs_c.py new file mode 100644 index 0000000..e4c4c4d --- /dev/null +++ b/brainscope_adapter/verify_vs_c.py @@ -0,0 +1,75 @@ +"""Gate: the brainscope twin must match the C runtime that ships to the board. + +Two checks, both on the deployed artifact: + 1. C-vs-Python logits on a fixed prompt (same prompt family as the exporter's + golden). The C side is runtime/llm.h - the code verify.c gates before + flashing - so agreement here means brainscope shows the device's model, + not an approximation of it. + 2. KV-cache parity: decoding token-by-token (what brainscope does) must equal + one full forward. + +Run from the repo root: + python brainscope_adapter/verify_vs_c.py +""" + +import subprocess +import sys +from pathlib import Path + +import numpy as np +import torch + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from build_hf import build # noqa: E402 + +PROMPT = [1, 500, 1000, 200, 42, 777, 13, 99] +TOLERANCE = 0.02 # same bar verify.c holds the C port to against PyTorch + + +def c_logits(artifacts: Path) -> np.ndarray: + exe = ROOT / "brainscope_adapter" / "dump_logits" + subprocess.run( + ["cc", "-O3", "-Wall", "-Wextra", "-I", str(ROOT / "runtime"), + "-o", str(exe), str(ROOT / "brainscope_adapter" / "dump_logits.c"), "-lm"], + check=True) + out = subprocess.run( + [str(exe), str(artifacts / "model.bin"), *map(str, PROMPT)], + check=True, capture_output=True, text=True).stdout + return np.array([float(v) for v in out.split()], dtype=np.float32) + + +def main(): + artifacts = ROOT / "artifacts" / "tinystories" + model = build(artifacts) + ids = torch.tensor([PROMPT]) + + with torch.no_grad(): + full = model(input_ids=ids).logits[0, -1].numpy() + + ref = c_logits(artifacts) + assert full.shape == ref.shape, f"vocab mismatch {full.shape} vs {ref.shape}" + diff = np.abs(full - ref) + print(f"C vs Python: max abs diff {diff.max():.6f} rms {np.sqrt((diff**2).mean()):.6f}") + print(f"top token: C={int(ref.argmax())} Python={int(full.argmax())}") + ok_c = diff.max() < TOLERANCE and ref.argmax() == full.argmax() + + with torch.no_grad(): + past = None + for t in PROMPT: + out = model(input_ids=torch.tensor([[t]]), past_key_values=past, use_cache=True) + past = out.past_key_values + step = out.logits[0, -1].numpy() + cache_diff = np.abs(step - full).max() + print(f"KV-cache vs full forward: max abs diff {cache_diff:.8f}") + ok_cache = cache_diff < 1e-4 + + if ok_c and ok_cache: + print("PASS: brainscope twin matches the device runtime") + return 0 + print("FAIL: twin diverges from the device runtime") + return 2 + + +if __name__ == "__main__": + sys.exit(main())