From 751cc098e21532d9c1b52e54e408179c1fed7b23 Mon Sep 17 00:00:00 2001 From: Katerina Fajmanova Date: Thu, 6 Aug 2026 16:40:50 +0200 Subject: [PATCH] feat: brainscope twin of the deployed model, gated against the C runtime --- brainscope_adapter/README.md | 38 ++++ brainscope_adapter/build_hf.py | 101 ++++++++++ brainscope_adapter/dump_logits | Bin 0 -> 37352 bytes brainscope_adapter/dump_logits.c | 52 +++++ brainscope_adapter/extract_direction.py | 40 ++++ brainscope_adapter/mood_pairs.jsonl | 10 + brainscope_adapter/ple_bin.py | 110 +++++++++++ brainscope_adapter/ple_hf.py | 240 ++++++++++++++++++++++++ brainscope_adapter/serve.py | 38 ++++ brainscope_adapter/tell_box.py | 65 +++++++ brainscope_adapter/verify_vs_c.py | 75 ++++++++ 11 files changed, 769 insertions(+) create mode 100644 brainscope_adapter/README.md create mode 100644 brainscope_adapter/build_hf.py create mode 100755 brainscope_adapter/dump_logits create mode 100644 brainscope_adapter/dump_logits.c create mode 100644 brainscope_adapter/extract_direction.py create mode 100644 brainscope_adapter/mood_pairs.jsonl create mode 100644 brainscope_adapter/ple_bin.py create mode 100644 brainscope_adapter/ple_hf.py create mode 100644 brainscope_adapter/serve.py create mode 100644 brainscope_adapter/tell_box.py create mode 100644 brainscope_adapter/verify_vs_c.py 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 0000000000000000000000000000000000000000..627b62ab99b8d55205ecd0494834cfb06bf64cb0 GIT binary patch literal 37352 zcmeHw4SZD9nfIM!fPj%Zk*HX)j_%kFR%3!_8AZ*&1n%er(V|VZ)DR{Kq=qCWlM1e4 z=p?{xn1HX?VvD}mcek}&x3XR8YS-FHkPvXINvsg1Y6L%KKqZJpK+Np_f9^dqlOa(S z+xP32%WpXMJm)#z&-Xc#xxc_SE6rk2M;IBSES0^rj6?Gfa3oQr4Tumm$(g23V1)93fNI8}7+J!zhudVl(orH;S)j zJ~e%2o`EM*Ssw4N)+kp}(ug@QnI_lc+Su+h>e)hrQ^lOO-ArXYE*IU zJ!LA_lWW4yG_$?bG#6!jWSUxUu~{!Q^%&I@nX(T)-+Cudb~^pdHS29M*D*DnZZxmR zRQ7iZ>Tz98^|6_Vx0?M8w-5d1ddk$?PUVWK%7s@?y`rk@@~X<}hNYJ;op$x*S5M8U zugRGrTIRUy_>iV%&$*4;5hhG-XIi6?m+i4E+bk#k*!BbU-IJram;UtJpFRGoO;7Le zw(TYk$IUv#;luXimNU=y^Y9^#SP#nzyeVmfeFdNH@)wu5K5CI{!?HooEac7_1Tiwd z9R|*R>;}r9|H&>xwX=E{_`YG_-x~(Le;D}d!@w^9+>TG`^B@32jq}hj@HxZ4e?1I5 zHVph9hJhCk1OLe|aQ`syJBER`4g)_r41Cit@XG+V*f{Ctg5LlpI5rDs$3~4S+u04x}-i(S{EoOF;Is@tgo#0 z*VI=i^>@_;DwOi2wG~Qj&E0sbtH4`%U0qEb8bdq&J4^hF?<}b(t*io|wyv@|P(g@N zQCD7ChMJWDr9My>sHsvaYHG`?fmD%%wWPGFs>Tl>xZiJ%Ye{KkHR@mr?8Cpbw4|c4 zy0ohDo^qukP+ouPXmzUqEptU6w9arB=nR`B*pxvyAywqTH(7GN^u}s ze0&_LDC?5oY7%@y5}dRspG`^dED1#VcoIBWu5V6)Pe{UlItfnwE1z%@{A>wC+L;7D zCkft_1V1+kKJq%U0wXIhvI76FuE1gStH-tA`x#oxs8{b-6sEdKoLDEhYce zl>9mIyrv(FJdU0E18Few9QQ0oy|4+H1iiYy?Y`Z8n|t1E-_lmRPm5OzW;&G5d0!2J z)1P(F#MJaHjsDwc}2N7ucn~`#)c|CbY@UH3GB0m@4gOOH|zN+c_pW}Q`!;8|R zpP`a#o7(*H&4j%f`2+xOOM&8Txh>cWnK zmisO%XeoDUZ3jg8F?UNrZiQOY>27J5$_g*@zmGM2zQYOp-Kc()>xJ2RLxHnF(p5W& zvD>4292x55HI94~5z=X(FMLto=*T4un0RHh(1Vq4#N>%S7y0HLYHPEGbdTD4AJT&6 zBWi0vyoA-(JH^Y}YU@Jr5>s1?@q(3ay#=W|G&2@FdZD_?BTCWIwRrL9{SUi8Io8~x zHeZHpu+O4nPqE-?^CaZ;qkydTIBH$5-j~r7R-4BF=nKuUwp{JjcNXZqH-#phRiPf- z!HIn6^}p{9d9IFO?z?>Y8pmRDhVMRHpg#G$ngs}`g{Xi2WPo#82LysMAQibTxy)$D zxI^E!bPs6$0W9HaAZhw>yfFggRiEtfS$1mrdPgmr$1-ahH1)}?VinQ!0M@Tv&B83q z-ou)u9fjN3?@>?7IA7= z>=KPnZgTV+C2g7|EZZJ1`Z=hDtiNJqhI2hF8NSd`(VZPxbeECn&K;U-nH{oz-|O$} zKJE({6rt%phaK~I2`%KKistVTQ})kuWN=Nx*c;+Z)AVEJD4jsVI$)Gs2e1DP4ivAT z+p2a>u3yE7a|PDIPkbTk`2_CLEL}dm#}}IZ2q;XnnU6Mu?H2537=7+6wDhnio8H}h zO!If+RnzotcW9o&!q!2?rcm~)5^b@D*%d+6Y5F|J(iSn5d5(z~ zG&b^C$7=MCF~5ZTais5Yc-Y2;Z}Q~?UUqx+6%#RWp9LFe)3BoYU3N`(VLi`@RTQYz z*l@QKTXQqj4mO39wTDP4}!-?ph@ zzrC(wYO|M`3c|DqA`djsc(0=t9T>!}HgBA*D3lGF6>=>MVM5;CgcOSki3EW}5!MO7 z!=D_}T6_*Cc|z57^@C<>TT#ZEOF=`qZGCAzf1gi(FY*MyNQs^g z=8kMkKZa0q#mT2 z{}B52iMKV5CZ?c^Mzn&yoj{rg5Vo%tZJz)Xkhz3_yf7G#9)9?fPqh|9C89B(A*`cLlvTR&^@(`>tniG$3GC*kr*ateZ`9hAmybnAW&b>wpxf-^tP=}27 zfUQH;JbY#_Q^f7mKh&cI`Wu#QZuQ@-4v*_~yC=Tf)JEmLtzw1BdmaZGGpC?ege7ZF|R~ zQs4bJB?9&D#HUB?*wN>49l7uFBW;#|Q}3s~3f*Pth9RInWlax9tm;$YU?k(PdTIEu zTHJovt!5P{p;;EUt7Ca-LFjfE7B9Ldf24(ONppAapBb8dmS)*~z^Vn?)AUX@<58cQ zp6=0)y2A&p!9Epjq2I7ueQLY$iY^y+sNV^z#cghNf>npQ(_F_Jo|_R0+@9uk9a(;* z-fs*iJ6-QW3;9PnEP>JC{TNxhXyRa4^B+kbp!HuhEazNr@cA^aWt-1c<|tB|e@Yt| zds#~(n9A@Wwfq(|tv=azSiPgIpcz)=dTvU{}rdosU1DAwEL`@{+ed#Rc+(2gf_K!r~uG1y4R9jIM&bBY82^xfes_L&O;7 z>>9`Pbo2uHrQHfiB%~6p<0Pr{D{d#M@aP^ob1D*F=rYJ9U&|Vy3$TwuE(tkAp;V~9 zK)Ix2T7)f-vdOZIYz2>T&NjE+0U=eOP(nf42%$s~1-WF1sbuN&F-|m2M>@SGrPF1U zPGA2*>9m#7DILPe^4hlo*|R*`Y^wXb9Nkn-rE+r(9+v|l#~f+2EDzRlZImDdo55jbVH&J zeri!$e}h6S>XCxHw>>RSh#-O%+US@Lg$$_|K*OOY9Mh01&^v^HpN4hN07f&tm(49sT} zw>($C#4rR5+6r`A=2&q5Lj`(2RFIPa2@4LFy#8LgU~us;3w;|{^AXGm76z1UkQMzjY2fC3>Bs-CN%?6y_khOz=49w0Q%8%idmjYl%Q6UJ0hC*WN9cUn(xEqiK0@YcQ z<9!YbL&;}s3;9Oo&E4#0Jt}Ra*iHR)Y|3G%xU}B^fEL#{BE$%3Hoo`a9cI@8bq;Ly zqim}mSb8Ulrm=yysH2dX&pPJFG`u6419E|5ZEd!SW-&ch1f|7i&F9%*X-B%f$#E|dKE9TJ6hw=rSLoe7>}sxt;;<;hdZ{;qQkxuIK#$X& z81E)v#K*fC#q?X2%0aIq0KE#8jHcPfQGrAk!_v>Zg8ksqEhZfKMp51*M(AsKdLps~ z%^#?(-^5GsS+QqOzvHiBEqF0hDjQHlvQO2+U1~=c^>FtfJ-pMDQ`PjDceRp4u4PZ9Zj)8ntA0l1l$rv^Jg9h zmQO!QwG4YjZ8{&v=kDNV(9AaqzuLzb&{0jO=M1TymtDx|&*SvB2nOOn)4=h?)-hmi zp{p&_)talLVV{qb33IZL=43Y7XVr-`AJH_Fw>R{R@(Q$ftV?3jD{um-7D2*RD~OE) zYW>Yw=9W>OZ3jajLEfNikN$*X6U!|S1{=lLG~g5o98m9EA9Op}MALh<;8ySh z{ajRz2-rqP7cfzA8nQkebQs*1jwT-m6!L@AKu>rgh=yiCzo>+uB8u~N&{4CY#H)v8 z44gEAnQjz9g&rgu!tNqoJ^c)Ka3Ec6&BO5pXTy%U{}^GF!|4lU3+bhWF0&MdCXMmB z_G_VHoBBw&(6u#ReY9h2f1&HmeD#rzLf3w``p7oV*bgu2@VedvhDM9Xtv(7^L^O6$ zzt{BwD|ub70k_>Vb{9bhjp`8-l;6~^HvKOfM&bIS`e@ke>ScL;Mu*xI0u-YiqqbIw z(QXtHLrhLEt1mR^j2lB!(>YTOYceO4t>j@5s87?}sx31e|7Bb8-=xv1vrlzK)#3v* zy^Hr!!_I#PcB7i*04i)eU|G{pCyogmAa5nUzyy{f4Va^8^=huI4G&}ajIF_jk;v*W z2EZu5^bQYLl-mjS%MDjRZEC|xbJ#^(?unq<<=I>2geF}*Cv@3QeVD;!loMK@6j}~@ z!Fdv4K9EfF-=UyDrL2n~Xn|=P zjZcO#IMz2&w+=#Av%UswB?w-$6k1*a&F2@=gn8D{C>FqA5T9j-&q8jq+>52he{)?P zN3Gi2dxF6ggKZ6J!q$sc{O&|m8{5Lu8qPK*h1!L1*pXs}}{n*K-14<~1*|D%vp)O{o? zN-ks8PdGY}r zGKDFal4y&FeB*=*leyprZU`$2TyJ=@AEI8xI=EsE+e0rCqQ9seR(iuuY@U((zomC)JE`B=R0 zbJ!gvl7*Fn4YCw-V4ftzgU2%M8&+`=sW$(7ijXqI&{jOZ8in8au6w@wyY4&OB{&;0 z_RAnwO|Nnkp~F4g)pcmD5XNt4-YC)RYUD9chGO(rJU9)~-}LGGy!!Jv-{(rf4WEy5 z51eD<0}X7iI@=<*zC}CQh8*-{J}e9@Lod}O4S91LpF@O@`>uh1^LD5mZ-Zo}KaTBG z(d=|pE&h{Qv>Uvhts?4d^SWMa*k|sVTkmvr2fnB8_DtT6Q_Fyx%@nE5x6nx6oKU7g zh=rC2-dBq*AbFNs8`a|HG34bJsYTC2YFYA+Y_|j)!S@C%IL+C?B_ua`L(|u47D}DW z)flykNvYo8^H$CBy4UhL<51JoW}g3h^*kOi<@OlGt_BELjr1lxwTr0<3!DJ|*Ez)H5HZ_J>wrX_tpZlKUH8N!l& z5|>p@-I5lRF%#{^=|Si+7}8?(KzuNwNgOn{1!ag3?khqL;23nU&3ynYMkk~Qt4`vL zSE35j(HQls@Z3O$OWEaI#5HjwkWcI`)biUj0vQeLL(u@IxVF?=~y>8DdKP z_NZT8E2J9~f8k<46Hg?0;L>>K_=yhZLF#50t67WZw}9-(W*eifftgs>WSF)UIbVgEd0{|c~wuQ&V~ z_YWk3-p&0h68pD&lDU6A{iV2MTN-=Wpi?X~B^q7ssiwpxXS+5intALd(D9!+T8iHc$nVchpW!&YC zN5o$dMQaR?R(}n9mFjR)e#7BzYuINH{b-ZxAz8ne=ppt0w-7z3l_ZUm?rD}a1_1|+ zoyWf7h#6Zhm74~HFex4IUc!-YG}F~gak|N9n%eC9FcvF-Xab!VLGU5;q_$24MnPyn z%%g9^K?7{sekd9Q)3_KmkWM1(h9Fp{YCFp~N5T=$7Z|tV2*rIQoF2H@h|P&}u)!3R z*=huBcAzzwww)qq<2%PZNim^e5w!80107l{K%kr^K{=#sa@f0YI>b@RjLCSRiFXU& zwo$XL#_aW8;J1mGO$Nd-?%+{N0HFrVj-d_2Yv6$ekK!!lnTxRAeq-W!J%$5l%gmzr zH39#j)vU}c%pSDU)$le~pI4@Y;us_1JnO}Iw=qowPd8!Bc}LZVfFe|GylAj{rdhEk zJhN8U&VhZJw?_+I{zFWHICBB!x+CoS2ebfg_8P+@2)j2e~pHre6__%CxME$8ee{y%A@ z;l7eQXg^}krnf8HS=er-*NdSkG2vXxbACHRGouLF39yo2>Du=N(GYwv+l6@jDZ z2mQQzh~bAY{Iwnmetk?g{BqVfBX{F{JDQtN-LjVx81pu4yHhoUvV_`vQ9w5|k1Dscyt~Z4; zzf#!pjDi~U!*S6{7|W&`P8XL3ua_tLAS7(QrhlU8Z-6|A$Fbzt2xI`5a2cknG-7Pwgjw4ixUAlLm2Zwrtf zQP0xj4n{3UI;??@+-+~2>u%dO-aYmhTH=iXkN%Qp?3*~b{{%;C<3!M=Igjmtc5%4s z`bZ6}LQ}X8;PZD*Zuf?krFHN3gtE`_T8{hR;5iSB6?Cixx2F~8xO#AKkGuPLp}%Lg z`c(G01^UZw9O%1)`_uyUDU9PlhC8@F14Mr59{SfA2m%#P!nt-4+|Pr12!D;B`7Ks2 zuI>Q?23OHmpX<$rM?;>uX}Fs(O>gf$jthUdXfX9${iEPs^>IreBR~A1HD7&d7td2N zv<>J(w8!Sq0Z6_GTvOcP#mTVGvJcs=LRUm>7WZcvy!X*374a~(0JFr{7m5KA4p`;$&EA&E0%QgYmt4}jx1>y zNiq#TkRbqrX;_dZ7G){Im#G8{bd!L^ctF4$`0xa5#D;lER;8j&Cg6Jv0u}=>37AAH z2v`8;pG?3qp7CYqHxF{=wCOjNpx+5>=l@6g-9CtZaj-m?eo49r#$$Jq>6b)9`Yp7) zP7+R`UzZW#SDO!fG4Xm}-WP$`MfmO3F%Y(UqCbRRW0-;$o(nP}PYY>}Ya~+O-^Tlq zpJxC{NiVFxwa2MD!+AZ?TwFFYu3Th_87M~E%Aw;+B zjl~ik?oJW%TVsUgoZVn@UUs8J2P<~!7${s^^hgBiIvd_b--Fy)wdo9Wl?c{N5d;f6 zV0reB8M0%vj%_%*c?XVfM#xTW+HnXYkzqTid?N&kh#j1J+QK3hkQV(J(c-7zW5s|G zS2P;LIX#Do@y3TLjtn$9YFNmIH%}h#rMM1)@xE%rRi<*hHekP13d45>{Feu@+J^kNmp$Fd_AAb zy%o#NOJ_Ep3(JjLb86EMUd3^h-shPdfybpbm7;2azTczo=NAofagK;ZH!e!z{U;=a zTuumLb@{x9&}!1^BSfD5ING=ZO(MF?Q@zB(Lkc(3S78t)Et&}wP*S8fdHt8+!k7@c zr@#IS`c8NKzlNHpw*DvHjKyFSHQukcNB{Xx@OSo}Irg}|d&NtH!Pq_qr>OW->Xx=z zn$$3&$L@~K2BaIe%HT0>fOTpelX}TE-5r|2v!myuQ-BJ@o~RRncG=ULSSs3wUy0n0tx=IYsXA09A0$OoT3m;>`~ba~a(F9zag9^$wtl*8hNz z&QPttgfPkA)_)JkDYpJNP(|yHK{1iral$}umk|wQ$?Zszx&46*0T`CsPc|Wee>x^) zBCPSGuslr@(hkNqO-L~i(<_XJHP{;&|0ITohCTgw#_lh}?@9jPGfDjZx3wSbpUdp0 zV?Wx54Q8C=494xpiZr@bT?v~JB4QK7 ztnll~kspe$W}hSzt~zy@V3#uC9LNOn6%M4xSN{P|h*}g`@)eR~zIr%A0EXr(L|IPK zqCo%eXicfypT3komc&oV{@C9(K4|+~ zemc$gAkH_3i4P8DslR%B@I!DFENvbqeLEp%aGqjaNA`n0y=%#sm2l#il7$N3jL)MMTpR#TiSy8 z?v~1IYBHe&=)f#BrN5wD8QK0a^!nJBk^f2YgJ6>U2fh9c;|EJWmtIdt{x2PtSO*8Q zVdfd~9}*^6{-1U*`}pUm05g~!JqnL%C`o?bNs^rQ-YHs0w@XRxgCvLEAiMAz6a6RP zQ0|HT7QWd~eg1X)v2UQ0xaL}ElLykN?yBM6%kYNqx8pH`*G`TZ9D#ilj~S> zqh4|#3b7L$#v0+M;>5Vm1DXi4MGr?f1BlKY(J`;lmB2apdFEmDO^9)S6px&5LX6$( z?=sHF+xDU{TP6Ys#hDCwi(_!}I}0ht;LXRtn3{jR_-Z?^9x6DrrJKehZLG=D5b0DR zntzSsS-kuF8^!nQ@%_vn@uZUXiE*_DUpJWtXWsuZuirfmJe-1tCOrgxRB$XGnm(2X z0$Tt#u2|_Wis@o@Q7ga4 z?ZcOu)}eK3kOi9lQVS#0>c5}k)?an&dvG_$@`Aog?{@1w9@nSpoC9rdnctl@Pej{3 z9Jk^v)W?^n-^F=o7{rWt0lqx_9FId$&C(P6P@IGwZhIF|o1U>R#=knP9~#?@W6*B5 zE6Q`wt@yJzKWBZ!3EZJuEUx`(=w3A934M3K>)+wwZSI1}_&ViX_BZ!Mgq69Ow|K+w|Z+ULOEd)1*@yfet zVp~8jUC0}Zy_+}j3N8OB7~8l?sy2^BLkt6x)T00SoMCqD)Vo*g5Ml7(0HWYO0#2}_ zRzHL&_}K2~Dpt{VQep8gEQAX}H(8{#%FY&l+M$QSPrbAqylBM2zlS=&-4nelW^!FP z`uDFA#`{q;^-}u0{rI3hXCrI&$v=5;N_}36I*e4TEy8_X-Zz=@di=Qj-zKIDZ3%q== zniP1MQsA96w7~n6M=YrVPY63oJwxE_22S$GpTmiDay|zK+MtN|%&K`&spYzy90-Os z@%0j%x`)pG4*t?%VV6#0;NTFAItr-inNznJn2{|@5ef5PTFDHbq==r5ig+|S)u zQp&~)uyRB7E=&c&Vx&;-0h<~=EH@EX?(Jrp{T2LKX(y8{;sqYIy9B*&m~ zd_#pa$_rXd@I_%!0i^wpg8U=jMor9)NU&Il1eIWv#W{twATx8L1HZ?GJ!iYo=LLyA zyP`LS*(YWf84a+w1swg}HglZQK>;kp?Q(pRBJ%dH4GcJA$VTtqX;d{!I|i55bjgPQ z#Vl|}n{e*LF*qX+1ID2_IqzTZ;(DF=&}C36XJ8^ysr(9eULrAI8h_aE9lROKx*vXA z;(FOS{{79o54?K!9h$yZ3%=bycOK&Oc{{Ywbrwu2cJ^W1`~5gk0r2zCmlhgz7886K zM?7kc|fZ z$!s7a)@Sg!tk3Z^UiW(xD-GYcO$%Q41HNkdJAq8l<4R;A(I-G3Wf*J6sE3fZd>3`& z{j+}bZWN*x-^Jf54a&bx3gE~yvH~M3FtP$8D=@MGBP%en0wXIhvH~M3aLOyd-|5-; zk@wK=zq#@%R0`vB11mSwmo6&5#(qh?efpA`vhu2&g_YIU+ouO=?kumiSC(CG|4wC@ zJtrsUyW^C)g-ZS6n!13UKe04USy)|BS_Ik#=p0#rU%XDT*4>e6Gj{W-^+uyKgPzyfzMQY zGSC;J>&kfBtnm|XRL{6OqfxoxE3QkgbX*Lms6QWRF2*1#5{QTG2JmqLp3nN`!||-a zr}nW}?3KBAnKjvz|vXhGwaiz;e(q?;=z}g#Isy7u!VsS+g8Bt3_VU z!)*Sqmc)N|J*CB5)9K4M)4684z)Wk+w9!o0ndv4o-E5{|Gwm|dJ~Qn%Q~ap1cx+~B zH`81*%{SAzX1c&kQ-=G6S^4|9vuDn{#_qgrVMBGG!Jc;|ex5k@@~ay}cFOW8Q*&~s z=HyK_3qG$=l8;q6#ACZb59xuwSiHvEFuPgGpU4{u|02NC6q^zjj&0(%ym3%CBtHHy zjll2s<0a{thIh7OH~CT)-DJWgze{*2UM%<}U*|dl!ujx7f{z71HYE8~0P){?eAvE} z9}>O^FBWi}M*izbQ7zUqz zkLj$mEJey)*>CGG_`e(m{w(0^H+B3whQZ$hxLtOZlx416x&Hqdi!%6hhoae}{0ei5-S^l{rT?JP3ek-v8>$6qh- zZAzax2@&H$;nHE+d0-g$W5dAzZ5VjRF!0^Oz&`7s$)0!B5VYF03pGlrB=r8kW?S zRMjl14AkfN75*%Jj=!d=OaWqjprN7yK>SF3`J&4DKzUtBU`Yvnm%h5Z9xasBlq{;M zSy)|o_Ia6}3LglK;h5nNIKxti|gwv0g zUO81M@y;tONlbrU;Y?0`Ug^TBa;1D}ZBk3$#5@qFj#m~H;;X z7-jV}C5ua|%P@L$;ICa4?;wwJy&1L3DyvHx>dVWN`n&1^N%O->Eh(+4s`1CG7&ELc zQ!47q%kLZvE?HP#Z;q5pika5dRaOToIMq1jl{r(D+M2tQI>92=)Rrf9fL+v<*VWb3 z#VZk1URRM+iGZb*gV(C!?z+lA%9590_ax<%R17m7vEh>12W@vriGOJ+skgMM@}88Y zNvK$H5X#_zDmnGbmIO){BE@!ybg_KHa+KHADmm3Pf%2S1)eSkdb(niyU|AAqVMAqA z+2xgGCdfU*dpTKFpf4`1U##SmEvrTsMjEIyO1@uSS6^9EJt$LxvbyrBQdThEYO4ZD zju?3kUUL@JAj_Bfa=(%jC|`=KpwgVW8bQW6<%>=3TwI3MWX@>AU{j-pd`DNMODfTr zQ5__qK(s7b0=6G4txi8jzXz%YHj?r1heG+tRQf#z)WDVRO#coan*6ff3aC#rabpaex#n2FfU$`>r4Aprgn*Frm6M64LI&d*2wWoe@>=Ov!3+l zQpaD4a{6+zey=%xnXWbqI7Y6stS|k(O2Ak})|dXHOr^gF6!EZ)I`d>2T z`tp(Emv}N=jxzdoGB5o*nM$)nuCW}yXcYg2@Zp-u`qJN&sq{Bx|FXWE|Buc3Gm!F; z{;N#;%n7qT>!j9y3^1a|_Tzfo)aTN#2a0%7_kSax$@T4KJ(-HLMzbyxvpAmqBBg$< z>F3E*;>#vvUZ$H<>NlG8Wtwl!f4KV3q|^_aMv+XTx5%1CD*KgC!>>uKpS06uD&r;+ zKXv`Kp@2h_`!DTbnQoF5%`~-t!%s@omwu5E3U?SBu-ll$+hH7dxIUrS4VB>hO( zJ9y!BYLji_{-EtX10xk*=4E;aXlhb@>3>c<#rmJ2E@@ZxFa1TkSzj(J`Hzq6SFYnf z1I8)F>#rr2qS)^@(mo@J|1ayybTld_*N^*GjYbvOk8D8Jlle?kCM@3nI=1;{d%A;&M}591zk>_!da f{|(JXT9Z;J`<1w4*Ax^{G)6ZtzLHWQC9C`wAJ}<* literal 0 HcmV?d00001 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())