Web demo: mini transformer LM + scaling sliders

web/public/app.js

@@ -2,8 +2,8 @@
 // train a shared model together — averaging gradients over the data channels.
 "use strict";
 
-const DIN = 16, DOUT = 4, NPER = 128;
-let D = { n: NPER, din: DIN, h: 16, dout: DOUT };   // h is set by the training-settings slider
+// model + task now live in transformer.js — a mini transformer LM trained
+// through the verified INT8 units. Settings come from the sliders.
 const STUN = [{ urls: "stun:stun.l.google.com:19302" }];
 
 const ui = {
@@ -24,7 +24,9 @@ const ui = {
   roomInfo: document.getElementById("roomInfo"),
   roomCode: document.getElementById("roomCode"),
   copyLink: document.getElementById("copyLink"),
-  cfgH: document.getElementById("cfgH"),
+  cfgC: document.getElementById("cfgC"),
+  cfgT: document.getElementById("cfgT"),
+  cfgB: document.getElementById("cfgB"),
   cfgSteps: document.getElementById("cfgSteps"),
   cfgLr: document.getElementById("cfgLr"),
 };
@@ -48,7 +50,7 @@ let myId = null, compute = null, ws = null, L = null, wasDenied = false;
 const pcs = new Map(), chans = new Map();      // peerId -> RTCPeerConnection / DataChannel
 const names = new Map();                       // peerId -> device name
 const incoming = new Map();                    // step -> Map(peerId -> Float32Array)
-let W1, W2, Xdata, Ydata, training = false;    // 2-layer verified model
+let model = null, training = false;            // the mini transformer (transformer.js)
 let trainedSteps = 0;                          // steps baked into the current weights
 function nmeOf(id) { return names.get(id) || id; }
 
@@ -154,32 +156,35 @@ function cleanupPeer(id) { const pc = pcs.get(id); if (pc) pc.close(); pcs.delet
 
 // ---- checkpoints ------------------------------------------------------------
 // File layout (also the broadcast payload after the sentinel):
-//   8 bytes magic "DAISYPT1" | int32 din,h,dout,steps | f32 W1 | f32 W2
+//   8 bytes magic "DAISYPT2" | int32 c,t,vocab,steps | f32 flat params
 // DaisyChain's own format (not torch-pickle) — .pt extension for familiarity.
-const CKPT_MAGIC = "DAISYPT1";
+const CKPT_MAGIC = "DAISYPT2";
 const CKPT_SENTINEL = -2;                      // wire: [int32 -2][checkpoint bytes]
 
 function packCheckpoint() {
-  const buf = new ArrayBuffer(8 + 16 + (W1.length + W2.length) * 4);
+  const flat = Transformer.getFlatParams(model);
+  const buf = new ArrayBuffer(8 + 16 + flat.length * 4);
   new Uint8Array(buf, 0, 8).set([...CKPT_MAGIC].map(c => c.charCodeAt(0)));
-  new Int32Array(buf, 8, 4).set([DIN, D.h, DOUT, trainedSteps]);
-  new Float32Array(buf, 24, W1.length).set(W1);
-  new Float32Array(buf, 24 + W1.length * 4, W2.length).set(W2);
+  new Int32Array(buf, 8, 4).set([model.cfg.c, model.cfg.t, model.cfg.vocab, trainedSteps]);
+  new Float32Array(buf, 24).set(flat);
   return buf;
 }
 function parseCheckpoint(buf) {
   const magic = String.fromCharCode(...new Uint8Array(buf, 0, 8));
-  if (magic !== CKPT_MAGIC) throw new Error("not a DaisyChain checkpoint");
-  const [din, h, dout, steps] = new Int32Array(buf, 8, 4);
-  if (din !== DIN || dout !== DOUT || h < 8 || h > 64)
-    throw new Error(`shape mismatch: file is ${din}×${h}×${dout}, this build is ${DIN}×(8–64)×${DOUT}`);
-  if (buf.byteLength !== 24 + (din * h + h * dout) * 4) throw new Error("truncated checkpoint");
-  return { h, steps, w1: new Float32Array(buf.slice(24, 24 + din * h * 4)),
-           w2: new Float32Array(buf.slice(24 + din * h * 4)) };
+  if (magic !== CKPT_MAGIC) throw new Error("not a DaisyChain v2 checkpoint");
+  const [c, t, vocab, steps] = new Int32Array(buf, 8, 4);
+  if (vocab !== Transformer.VOCAB) throw new Error(`vocab mismatch (file ${vocab}, this build ${Transformer.VOCAB})`);
+  if (c < 16 || c > 64 || t < 16 || t > 64) throw new Error(`bad dims in checkpoint (width ${c}, seq ${t})`);
+  return { c, t, steps, flat: new Float32Array(buf.slice(24)) };
 }
 function applyCheckpoint(ck, from) {
-  if (ck.h !== D.h) { buildModel(ck.h); showCfgInUI({ h: ck.h, steps: +ui.cfgSteps.value, lr: +ui.cfgLr.value / 100 }); }
-  W1.set(ck.w1); W2.set(ck.w2); trainedSteps = ck.steps;
+  if (!model || model.cfg.c !== ck.c || model.cfg.t !== ck.t) {
+    buildModel({ ...readCfgFromUI(), c: ck.c, t: ck.t });
+    showCfgInUI({ ...readCfgFromUI(), c: ck.c, t: ck.t });
+  }
+  if (ck.flat.length !== model.nParams) throw new Error("truncated checkpoint");
+  Transformer.setFlatParams(model, ck.flat);
+  trainedSteps = ck.steps;
   ui.save.disabled = false;
   ui.step.textContent = `${ck.steps} baked in`;
   log(`checkpoint loaded (${ck.steps} steps) ${from ? "from " + from : "from file"} — all set to resume`);
@@ -197,38 +202,41 @@ function saveCheckpoint() {
   const blob = new Blob([packCheckpoint()], { type: "application/octet-stream" });
   const a = document.createElement("a");
   a.href = URL.createObjectURL(blob);
-  a.download = `daisychain-${DIN}x${D.h}x${DOUT}-step${trainedSteps}.pt`;
+  a.download = `daisychain-lm-w${model.cfg.c}-step${trainedSteps}.pt`;
   a.click();
   URL.revokeObjectURL(a.href);
 }
 
 // ---- training config: whoever presses Start sets it for the whole group ----
-const CFG_SENTINEL = -3;                       // wire: [int32 -3][int32 h,steps][f32 lr]
+const CFG_SENTINEL = -3;                       // wire: [int32 -3][int32 c,t,b,steps][f32 lr]
 function readCfgFromUI() {
-  return { h: +ui.cfgH.value, steps: +ui.cfgSteps.value, lr: +ui.cfgLr.value / 100 };
+  return { c: +ui.cfgC.value, t: +ui.cfgT.value, b: +ui.cfgB.value,
+           steps: +ui.cfgSteps.value, lr: +ui.cfgLr.value / 1000 };
 }
 function showCfgInUI(cfg) {
-  ui.cfgH.value = cfg.h; document.getElementById("vcfgH").textContent = cfg.h;
-  ui.cfgSteps.value = cfg.steps; document.getElementById("vcfgSteps").textContent = cfg.steps;
-  ui.cfgLr.value = Math.round(cfg.lr * 100); document.getElementById("vcfgLr").textContent = Math.round(cfg.lr * 100);
+  const set = (el, vid, val) => { el.value = val; document.getElementById(vid).textContent = val; };
+  set(ui.cfgC, "vcfgC", cfg.c); set(ui.cfgT, "vcfgT", cfg.t); set(ui.cfgB, "vcfgB", cfg.b);
+  set(ui.cfgSteps, "vcfgSteps", cfg.steps);
+  set(ui.cfgLr, "vcfgLr", Math.round(cfg.lr * 1000));
 }
 function broadcastConfig(cfg) {
-  const buf = new ArrayBuffer(16);
-  new Int32Array(buf, 0, 3).set([CFG_SENTINEL, cfg.h, cfg.steps]);
-  new Float32Array(buf, 12, 1)[0] = cfg.lr;
+  const buf = new ArrayBuffer(24);
+  new Int32Array(buf, 0, 5).set([CFG_SENTINEL, cfg.c, cfg.t, cfg.b, cfg.steps]);
+  new Float32Array(buf, 20, 1)[0] = cfg.lr;
   for (const dc of chans.values()) if (dc.readyState === "open") dc.send(buf);
 }
 function onConfig(peerId, buf) {
   if (training) { log(`ignored settings from ${nmeOf(peerId)} (already training)`); return; }
-  const [, h, steps] = new Int32Array(buf, 0, 3);
-  const lr = new Float32Array(buf, 12, 1)[0];
-  if (!(h >= 8 && h <= 64 && steps >= 1 && steps <= 10000 && lr > 0 && lr <= 1)) {
+  const [, c, t, b, steps] = new Int32Array(buf, 0, 5);
+  const lr = new Float32Array(buf, 20, 1)[0];
+  if (!(c >= 16 && c <= 64 && c % 2 === 0 && t >= 16 && t <= 64 && b >= 1 && b <= 32 &&
+        steps >= 1 && steps <= 10000 && lr > 0 && lr <= 0.2)) {
     log(`rejected bad settings from ${nmeOf(peerId)}`); return;
   }
-  const cfg = { h, steps, lr };
+  const cfg = { c, t, b, steps, lr };
   showCfgInUI(cfg);
-  log(`${nmeOf(peerId)} started the group: hidden=${h}, ${steps} steps, lr=${lr}`);
-  buildModel(cfg.h);
+  log(`${nmeOf(peerId)} started the group: width=${c}, seq=${t}, batch=${b}, ${steps} steps, lr=${lr}`);
+  buildModel(cfg);
   train(cfg);                                  // follow automatically
 }
 
@@ -281,20 +289,19 @@ function waitForGrads(step, cohort, timeoutMs = 8000) {
 
 // ---- compute: one async training step THROUGH the verified units -----------
 async function localStep() {
-  // forward runs through the verified INT8 multiply (WebGPU or CPU); STE backward
-  const fwd = await Verified.forward(Xdata, Ydata, W1, W2, D, L, compute.matmulInt8);
-  const grad = Verified.backward(Xdata, W1, W2, fwd, D);          // flat [gW1, gW2]
-  return { loss: fwd.loss, grad };
+  // transformer forward runs through the verified INT8 multiply; STE backward
+  return Transformer.trainStep(model);
 }
 
 // ---- the training loop -----------------------------------------------------
 async function train(cfg) {
   if (training) return; training = true; ui.start.disabled = true;
-  [ui.cfgH, ui.cfgSteps, ui.cfgLr].forEach(el => el.disabled = true);
+  cfgSliders().forEach(el => el.disabled = true);
   const cohort = [...chans.keys()];                    // lock the cohort (departed peers are pruned per-step)
   const steps = cfg.steps;
-  const opt = TrainCore.makeAdam(W1.length + W2.length, { lr: cfg.lr });
-  log(`training started — cohort ${cohort.length} peer(s), world ${cohort.length + 1}, hidden=${cfg.h}, optimizer ${opt.name}`);
+  const opt = TrainCore.makeAdam(model.nParams, { lr: cfg.lr });
+  log(`training started — cohort ${cohort.length} peer(s), world ${cohort.length + 1}, ` +
+      `width=${cfg.c} seq=${cfg.t} batch=${cfg.b}×${cohort.length + 1}, optimizer ${opt.name}`);
   for (let s = 0; s < steps; s++) {
     const { loss, grad } = await localStep();
     broadcastGrad(s, grad);
@@ -302,7 +309,7 @@ async function train(cfg) {
     const all = [grad, ...remote];
     const avg = TrainCore.averageGrads(all);
     const upd = opt.step(avg);                      // DaisyAdam on the cluster-avg grad
-    Verified.splitApply(W1, W2, upd, 1);            // W -= 1 * upd (lr folded into upd)
+    Transformer.applyUpdate(model, upd);            // W -= upd (lr folded into upd)
     incoming.delete(s);
     trainedSteps++;
     if (s % 10 === 0 || s === steps - 1) {
@@ -312,30 +319,24 @@ async function train(cfg) {
       await new Promise(r => setTimeout(r, 0));         // yield to UI
     }
   }
-  ui.diff.textContent = `done — trained through the verified units; all peers share one model.`;
   log(`training done — final loss ${ui.loss.textContent}`);
+  try {
+    const sample = await Transformer.generate(model, "the ", 70);
+    ui.diff.textContent = `the model speaks: “${sample.trim()}”`;
+  } catch (e) {
+    ui.diff.textContent = "done — trained through the verified units; all peers share one model.";
+  }
   training = false;
   ui.save.disabled = false;
   ui.start.disabled = false;
-  [ui.cfgH, ui.cfgSteps, ui.cfgLr].forEach(el => el.disabled = false);
-}
-
-// 2-layer float target (matches the model shape) for a learnable task
-function target(X, h) {
-  const Wt1 = randn(DIN * h, mulberry32(42)), Wt2 = randn(h * DOUT, mulberry32(43));
-  const hpre = TrainCore.matmul(X, Wt1, NPER, DIN, h);
-  for (let i = 0; i < hpre.length; i++) hpre[i] = Math.max(0, hpre[i]);
-  return TrainCore.matmul(hpre, Wt2, NPER, h, DOUT);
+  cfgSliders().forEach(el => el.disabled = false);
 }
+function cfgSliders() { return [ui.cfgC, ui.cfgT, ui.cfgB, ui.cfgSteps, ui.cfgLr]; }
 
-// (re)build the model for a hidden size — deterministic shared init (same seeds
-// on every peer), fresh per-peer data shard
-function buildModel(h) {
-  D = { n: NPER, din: DIN, h, dout: DOUT };
-  W1 = randn(DIN * h, mulberry32(7));
-  W2 = randn(h * DOUT, mulberry32(8));
-  Xdata = randn(NPER * DIN);
-  Ydata = target(Xdata, h);
+// (re)build the mini transformer — deterministic shared init (same seeds on
+// every peer); each device samples its own batch windows from the shared corpus
+function buildModel(cfg) {
+  model = Transformer.init(cfg, L, compute.matmulInt8);
   trainedSteps = 0;
 }
 
@@ -387,9 +388,10 @@ function buildModel(h) {
   }
   ui.backend.textContent = `${compute.backend.toUpperCase()} — ${compute.label} · through verified INT8 units`;
   // slider readouts
-  for (const [el, v] of [[ui.cfgH, "vcfgH"], [ui.cfgSteps, "vcfgSteps"], [ui.cfgLr, "vcfgLr"]])
+  for (const [el, v] of [[ui.cfgC, "vcfgC"], [ui.cfgT, "vcfgT"], [ui.cfgB, "vcfgB"],
+                         [ui.cfgSteps, "vcfgSteps"], [ui.cfgLr, "vcfgLr"]])
     el.oninput = () => document.getElementById(v).textContent = el.value;
-  buildModel(+ui.cfgH.value);
+  buildModel(readCfgFromUI());
   ui.me.textContent = deviceName;
   if (room()) {
     ui.roomInfo.style.display = "";
@@ -404,7 +406,7 @@ function buildModel(h) {
   connectSignaling();
   ui.start.onclick = () => {
     const cfg = readCfgFromUI();
-    buildModel(cfg.h);
+    buildModel(cfg);
     broadcastConfig(cfg);                                // everyone follows these settings
     train(cfg);
   };

web/public/index.html

@@ -105,13 +105,17 @@
 
   <div class="card">
     <div class="lbl">🎛 Training settings</div>
-    <label class="slbl">Hidden units <span class="sval" id="vcfgH">16</span></label>
-    <input type="range" id="cfgH" min="8" max="64" step="8" value="16">
+    <label class="slbl">Model width <span class="sval" id="vcfgC">32</span></label>
+    <input type="range" id="cfgC" min="16" max="64" step="16" value="32">
+    <label class="slbl">Sequence length <span class="sval" id="vcfgT">32</span></label>
+    <input type="range" id="cfgT" min="16" max="64" step="16" value="32">
+    <label class="slbl">Batch per device <span class="sval" id="vcfgB">8</span></label>
+    <input type="range" id="cfgB" min="2" max="32" step="2" value="8">
     <label class="slbl">Steps <span class="sval" id="vcfgSteps">300</span></label>
-    <input type="range" id="cfgSteps" min="100" max="1000" step="100" value="300">
-    <label class="slbl">Learning rate ×100 <span class="sval" id="vcfgLr">20</span></label>
+    <input type="range" id="cfgSteps" min="100" max="10000" step="100" value="300">
+    <label class="slbl">Learning rate ×1000 <span class="sval" id="vcfgLr">20</span></label>
     <input type="range" id="cfgLr" min="5" max="50" step="5" value="20">
-    <p class="note" style="margin:.5rem 0 0">Whoever presses Start sets the settings for the whole group — every device follows automatically.</p>
+    <p class="note" style="margin:.5rem 0 0">A mini transformer language model — attention, MLP blocks, next-character prediction — every multiply through the verified INT8 units. Whoever presses Start sets the settings for the whole group; total batch scales with every device that joins.</p>
   </div>
 
   <div class="card" style="text-align:center">
@@ -147,6 +151,7 @@
 
   <script src="traincore.js"></script>
   <script src="verified_core.js"></script>
+  <script src="transformer.js"></script>
   <script src="webgpu.js"></script>
   <script src="app.js"></script>
 </body>

web/public/transformer.js

@@ -0,0 +1,311 @@
+// A miniature transformer language model that trains THROUGH the verified INT8
+// units: every matrix product in the forward pass — QKV projections, attention
+// scores, attention·values, output projection, the MLP, and the unembedding —
+// runs through the verified multiply LUT (an emulated INT8 tensor core).
+// Backward is a straight-through estimator in float (the integer path has no
+// gradient), exactly like the Python VerifiedLinear.
+//
+// Task: next-character prediction on a deterministic, self-generated corpus
+// (every peer builds the same text from the same seed — nothing to download).
+(function (root) {
+  "use strict";
+
+  let TC, V;   // TrainCore / Verified — resolved per environment at the end
+
+  function mulberry32(a) { return function () { a |= 0; a = a + 0x6D2B79F5 | 0; let t = Math.imul(a ^ a >>> 15, 1 | a); t = t + Math.imul(t ^ t >>> 7, 61 | t) ^ t; return ((t ^ t >>> 14) >>> 0) / 4294967296; }; }
+  function randn(n, rng) { const o = new Float32Array(n); for (let i = 0; i < n; i += 2) { let u = 0, v = 0; while (u === 0) u = rng(); while (v === 0) v = rng(); const m = Math.sqrt(-2 * Math.log(u)); o[i] = m * Math.cos(2 * Math.PI * v); if (i + 1 < n) o[i + 1] = m * Math.sin(2 * Math.PI * v); } return o; }
+
+  // ---- corpus: deterministic cottagecore prose, identical on every peer -----
+  const W_ADJ = ["mossy", "golden", "amber", "quiet", "little", "misty", "sunny", "wild", "cozy", "dusty", "merry", "brave"];
+  const W_NOUN = ["fox", "hare", "owl", "badger", "toad", "sparrow", "otter", "deer", "mushroom", "acorn", "willow", "robin", "river", "meadow", "garden", "lantern"];
+  const W_VERB = ["naps", "sings", "wanders", "hides", "dreams", "waits", "dances", "listens", "rests", "grows"];
+  const W_PREP = ["by", "under", "near", "beside", "beyond", "inside"];
+  function buildCorpus() {
+    const rng = mulberry32(20260712);
+    const pick = (a) => a[Math.floor(rng() * a.length)];
+    let s = "";
+    while (s.length < 60000)
+      s += `the ${pick(W_ADJ)} ${pick(W_NOUN)} ${pick(W_VERB)} ${pick(W_PREP)} the ${pick(W_ADJ)} ${pick(W_NOUN)}. `;
+    return s;
+  }
+  const CORPUS = buildCorpus();
+  const CHARS = [...new Set(CORPUS)].sort();
+  const VOCAB = CHARS.length;
+  const STOI = Object.fromEntries(CHARS.map((c, i) => [c, i]));
+  const IDS = new Int32Array(CORPUS.length);
+  for (let i = 0; i < CORPUS.length; i++) IDS[i] = STOI[CORPUS[i]];
+
+  // ---- verified matmul: float in -> quantize -> LUT multiply -> dequant -----
+  // weights/projections may go through WebGPU; attention (per-head, many small
+  // matmuls) uses the CPU LUT path — same verified units, no dispatch overhead.
+  async function vmm(Xf, Wf, m, k, n, ctx) {
+    const xq = V.quantize(Xf), wq = V.quantize(Wf);
+    const acc = await ctx.matmulInt8(xq.q, wq.q, m, k, n, ctx.L);
+    const dq = xq.scale * wq.scale, out = new Float32Array(m * n);
+    for (let i = 0; i < out.length; i++) out[i] = acc[i] * dq;
+    return out;
+  }
+  function vmmCPU(Xf, Wf, m, k, n, ctx) {
+    const xq = V.quantize(Xf), wq = V.quantize(Wf);
+    const acc = V.lutMatmulJS(xq.q, wq.q, m, k, n, ctx.L);
+    const dq = xq.scale * wq.scale, out = new Float32Array(m * n);
+    for (let i = 0; i < out.length; i++) out[i] = acc[i] * dq;
+    return out;
+  }
+
+  // ---- layernorm (no affine) -------------------------------------------------
+  function lnFwd(x, rows, C) {
+    const y = new Float32Array(rows * C), sig = new Float32Array(rows);
+    for (let r = 0; r < rows; r++) {
+      let mu = 0; for (let j = 0; j < C; j++) mu += x[r * C + j]; mu /= C;
+      let v = 0; for (let j = 0; j < C; j++) { const d = x[r * C + j] - mu; v += d * d; }
+      const s = Math.sqrt(v / C + 1e-5); sig[r] = s;
+      for (let j = 0; j < C; j++) y[r * C + j] = (x[r * C + j] - mu) / s;
+    }
+    return { y, sig };
+  }
+  function lnBwd(dy, y, sig, rows, C) {
+    const dx = new Float32Array(rows * C);
+    for (let r = 0; r < rows; r++) {
+      let mdy = 0, mdyy = 0;
+      for (let j = 0; j < C; j++) { mdy += dy[r * C + j]; mdyy += dy[r * C + j] * y[r * C + j]; }
+      mdy /= C; mdyy /= C;
+      for (let j = 0; j < C; j++) dx[r * C + j] = (dy[r * C + j] - mdy - y[r * C + j] * mdyy) / sig[r];
+    }
+    return dx;
+  }
+
+  // ---- model -----------------------------------------------------------------
+  // cfg: { c: width, t: seq len, b: batch/device, layers, heads, steps, lr }
+  function init(cfg, L, matmulInt8) {
+    const c = cfg.c, layers = cfg.layers || 2, heads = cfg.heads || 2, hidden = 2 * c;
+    let seed = 100;
+    const mk = (nEl, scale) => { const w = randn(nEl, mulberry32(seed++)); for (let i = 0; i < nEl; i++) w[i] *= scale; return w; };
+    const params = [], names = [];
+    const add = (name, w) => { params.push(w); names.push(name); return w; };
+    const m = {
+      cfg: { ...cfg, layers, heads, hidden, vocab: VOCAB },
+      ctx: { L, matmulInt8 },
+      emb: add("emb", mk(VOCAB * c, 0.08)),
+      pos: add("pos", mk(cfg.t * c, 0.02)),
+      blocks: [], params, names,
+    };
+    for (let l = 0; l < layers; l++)
+      m.blocks.push({
+        Wq: add(`b${l}.Wq`, mk(c * c, 0.08)), Wk: add(`b${l}.Wk`, mk(c * c, 0.08)),
+        Wv: add(`b${l}.Wv`, mk(c * c, 0.08)), Wo: add(`b${l}.Wo`, mk(c * c, 0.08)),
+        W1: add(`b${l}.W1`, mk(c * hidden, 0.08)), W2: add(`b${l}.W2`, mk(hidden * c, 0.08)),
+      });
+    m.Wu = add("Wu", mk(c * VOCAB, 0.08));
+    m.nParams = params.reduce((a, p) => a + p.length, 0);
+    return m;
+  }
+
+  function sampleBatch(cfg) {
+    const { b, t } = cfg;
+    const X = new Int32Array(b * t), Y = new Int32Array(b * t);
+    for (let i = 0; i < b; i++) {
+      const off = Math.floor(Math.random() * (IDS.length - t - 1));
+      for (let j = 0; j < t; j++) { X[i * t + j] = IDS[off + j]; Y[i * t + j] = IDS[off + j + 1]; }
+    }
+    return { X, Y };
+  }
+
+  // per-head slice helpers (q: BT×C row-major; head h occupies cols [h*hd, (h+1)*hd))
+  function headSlice(q, bIdx, h, T, C, hd) {
+    const out = new Float32Array(T * hd);
+    for (let ti = 0; ti < T; ti++)
+      for (let j = 0; j < hd; j++) out[ti * hd + j] = q[(bIdx * T + ti) * C + h * hd + j];
+    return out;
+  }
+  function headUnslice(dst, src, bIdx, h, T, C, hd, accumulate) {
+    for (let ti = 0; ti < T; ti++)
+      for (let j = 0; j < hd; j++) {
+        const di = (bIdx * T + ti) * C + h * hd + j;
+        dst[di] = (accumulate ? dst[di] : 0) + src[ti * hd + j];
+      }
+  }
+
+  // ---- forward THROUGH the verified units (caches kept for STE backward) -----
+  async function forward(m, X, Y) {
+    const { c: C, t: T, b: B, layers, heads, hidden, vocab } = m.cfg;
+    const BT = B * T, hd = C / heads, ctx = m.ctx;
+    const cache = { X, Y, blocks: [] };
+    let x = new Float32Array(BT * C);
+    for (let i = 0; i < BT; i++) {
+      const id = X[i], tpos = i % T;
+      for (let j = 0; j < C; j++) x[i * C + j] = m.emb[id * C + j] + m.pos[tpos * C + j];
+    }
+    for (let l = 0; l < layers; l++) {
+      const bl = m.blocks[l], cb = { xin: x };
+      const l1 = lnFwd(x, BT, C); cb.ln1 = l1;
+      const q = await vmm(l1.y, bl.Wq, BT, C, C, ctx);
+      const k = await vmm(l1.y, bl.Wk, BT, C, C, ctx);
+      const v = await vmm(l1.y, bl.Wv, BT, C, C, ctx);
+      cb.q = q; cb.k = k; cb.v = v;
+      const ctxOut = new Float32Array(BT * C);
+      cb.att = [];                                     // per (b,h): softmax probs
+      const scale = 1 / Math.sqrt(hd);
+      for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
+        const qh = headSlice(q, bi, h, T, C, hd), kh = headSlice(k, bi, h, T, C, hd), vh = headSlice(v, bi, h, T, C, hd);
+        const kT = TC.transpose(kh, T, hd);            // hd×T
+        const s = vmmCPU(qh, kT, T, hd, T, ctx);       // scores through the units
+        for (let i = 0; i < T * T; i++) s[i] *= scale;
+        const a = new Float32Array(T * T);             // causal softmax
+        for (let ti = 0; ti < T; ti++) {
+          let mx = -1e30;
+          for (let tj = 0; tj <= ti; tj++) mx = Math.max(mx, s[ti * T + tj]);
+          let z = 0;
+          for (let tj = 0; tj <= ti; tj++) { const e = Math.exp(s[ti * T + tj] - mx); a[ti * T + tj] = e; z += e; }
+          for (let tj = 0; tj <= ti; tj++) a[ti * T + tj] /= z;
+        }
+        const ch = vmmCPU(a, vh, T, T, hd, ctx);       // attn·V through the units
+        headUnslice(ctxOut, ch, bi, h, T, C, hd, false);
+        cb.att.push({ a, qh, kh, vh });
+      }
+      cb.ctxOut = ctxOut;
+      const attnOut = await vmm(ctxOut, bl.Wo, BT, C, C, ctx);
+      const x2 = new Float32Array(BT * C);
+      for (let i = 0; i < x2.length; i++) x2[i] = x[i] + attnOut[i];
+      cb.x2 = x2;
+      const l2 = lnFwd(x2, BT, C); cb.ln2 = l2;
+      const h1 = await vmm(l2.y, bl.W1, BT, C, hidden, ctx);
+      const mask = new Uint8Array(h1.length);
+      for (let i = 0; i < h1.length; i++) { if (h1[i] > 0) mask[i] = 1; else h1[i] = 0; }
+      cb.h1 = h1; cb.mask = mask;
+      const mlpOut = await vmm(h1, bl.W2, BT, hidden, C, ctx);
+      x = new Float32Array(BT * C);
+      for (let i = 0; i < x.length; i++) x[i] = x2[i] + mlpOut[i];
+      cache.blocks.push(cb);
+    }
+    const lf = lnFwd(x, BT, C); cache.lnf = lf; cache.xf = x;
+    const logits = await vmm(lf.y, m.Wu, BT, C, vocab, ctx);
+    // cross-entropy + dlogits
+    let loss = 0;
+    const dlogits = new Float32Array(BT * vocab);
+    for (let i = 0; i < BT; i++) {
+      let mx = -1e30;
+      for (let j = 0; j < vocab; j++) mx = Math.max(mx, logits[i * vocab + j]);
+      let z = 0;
+      for (let j = 0; j < vocab; j++) z += Math.exp(logits[i * vocab + j] - mx);
+      const lz = Math.log(z) + mx;
+      loss += lz - logits[i * vocab + Y[i]];
+      for (let j = 0; j < vocab; j++)
+        dlogits[i * vocab + j] = (Math.exp(logits[i * vocab + j] - lz) - (j === Y[i] ? 1 : 0)) / BT;
+    }
+    loss /= BT;
+    cache.dlogits = dlogits;
+    return { loss, cache, logits };
+  }
+
+  // ---- STE backward (float), mirrors forward exactly --------------------------
+  function backward(m, cache) {
+    const { c: C, t: T, b: B, layers, heads, hidden, vocab } = m.cfg;
+    const BT = B * T, hd = C / heads, mm = TC.matmul, tr = TC.transpose;
+    const g = m.params.map(p => new Float32Array(p.length));
+    const gi = Object.fromEntries(m.names.map((n, i) => [n, i]));
+    // unembed
+    g[gi.Wu] = mm(tr(cache.lnf.y, BT, C), cache.dlogits, C, BT, vocab);
+    let dx = lnBwd(mm(cache.dlogits, tr(m.Wu, C, vocab), BT, vocab, C), cache.lnf.y, cache.lnf.sig, BT, C);
+    const scale = 1 / Math.sqrt(hd);
+    for (let l = layers - 1; l >= 0; l--) {
+      const bl = m.blocks[l], cb = cache.blocks[l];
+      // mlp: x3 = x2 + relu(ln2 @ W1) @ W2
+      const dmlpOut = dx;                              // residual passthrough handled below
+      g[gi[`b${l}.W2`]] = mm(tr(cb.h1, BT, hidden), dmlpOut, hidden, BT, C);
+      const dh1 = mm(dmlpOut, tr(bl.W2, hidden, C), BT, C, hidden);
+      for (let i = 0; i < dh1.length; i++) if (!cb.mask[i]) dh1[i] = 0;
+      g[gi[`b${l}.W1`]] = mm(tr(cb.ln2.y, BT, C), dh1, C, BT, hidden);
+      const dln2in = lnBwd(mm(dh1, tr(bl.W1, C, hidden), BT, hidden, C), cb.ln2.y, cb.ln2.sig, BT, C);
+      const dx2 = new Float32Array(BT * C);
+      for (let i = 0; i < dx2.length; i++) dx2[i] = dx[i] + dln2in[i];
+      // attention: x2 = xin + (ctxOut @ Wo)
+      g[gi[`b${l}.Wo`]] = mm(tr(cb.ctxOut, BT, C), dx2, C, BT, C);
+      const dctx = mm(dx2, tr(bl.Wo, C, C), BT, C, C);
+      const dq = new Float32Array(BT * C), dk = new Float32Array(BT * C), dv = new Float32Array(BT * C);
+      let ai = 0;
+      for (let bi = 0; bi < B; bi++) for (let h = 0; h < heads; h++) {
+        const { a, qh, kh, vh } = cb.att[ai++];
+        const dch = headSlice(dctx, bi, h, T, C, hd);        // T×hd
+        const dvh = mm(tr(a, T, T), dch, T, T, hd);          // aᵀ @ dctx
+        const da = mm(dch, tr(vh, T, hd), T, hd, T);         // dctx @ vᵀ
+        const ds = new Float32Array(T * T);                  // softmax bwd (causal)
+        for (let ti = 0; ti < T; ti++) {
+          let dot = 0;
+          for (let tj = 0; tj <= ti; tj++) dot += da[ti * T + tj] * a[ti * T + tj];
+          for (let tj = 0; tj <= ti; tj++) ds[ti * T + tj] = a[ti * T + tj] * (da[ti * T + tj] - dot) * scale;
+        }
+        const dqh = mm(ds, kh, T, T, hd);                    // ds @ k
+        const dkh = mm(tr(ds, T, T), qh, T, T, hd);          // dsᵀ @ q
+        headUnslice(dq, dqh, bi, h, T, C, hd, true);
+        headUnslice(dk, dkh, bi, h, T, C, hd, true);
+        headUnslice(dv, dvh, bi, h, T, C, hd, true);
+      }
+      g[gi[`b${l}.Wq`]] = mm(tr(cb.ln1.y, BT, C), dq, C, BT, C);
+      g[gi[`b${l}.Wk`]] = mm(tr(cb.ln1.y, BT, C), dk, C, BT, C);
+      g[gi[`b${l}.Wv`]] = mm(tr(cb.ln1.y, BT, C), dv, C, BT, C);
+      const dln1in = new Float32Array(BT * C);
+      const addIn = (dsrc, W) => { const t2 = mm(dsrc, tr(W, C, C), BT, C, C); for (let i = 0; i < dln1in.length; i++) dln1in[i] += t2[i]; };
+      addIn(dq, bl.Wq); addIn(dk, bl.Wk); addIn(dv, bl.Wv);
+      const dxin = lnBwd(dln1in, cb.ln1.y, cb.ln1.sig, BT, C);
+      dx = new Float32Array(BT * C);
+      for (let i = 0; i < dx.length; i++) dx[i] = dx2[i] + dxin[i];
+    }
+    // embedding + positional
+    const ge = g[gi.emb], gp = g[gi.pos];
+    for (let i = 0; i < BT; i++) {
+      const id = cache.X[i], tpos = i % T;
+      for (let j = 0; j < C; j++) { ge[id * C + j] += dx[i * C + j]; gp[tpos * C + j] += dx[i * C + j]; }
+    }
+    // flatten
+    const flat = new Float32Array(m.nParams);
+    let off = 0;
+    for (const t of g) { flat.set(t, off); off += t.length; }
+    return flat;
+  }
+
+  async function trainStep(m) {
+    const { X, Y } = sampleBatch(m.cfg);
+    const { loss, cache } = await forward(m, X, Y);
+    const grad = backward(m, cache);
+    return { loss, grad };
+  }
+
+  function applyUpdate(m, upd) {          // W -= upd (lr folded in by the optimizer)
+    let off = 0;
+    for (const p of m.params) { for (let i = 0; i < p.length; i++) p[i] -= upd[off + i]; off += p.length; }
+  }
+  function getFlatParams(m) {
+    const flat = new Float32Array(m.nParams);
+    let off = 0;
+    for (const p of m.params) { flat.set(p, off); off += p.length; }
+    return flat;
+  }
+  function setFlatParams(m, flat) {
+    let off = 0;
+    for (const p of m.params) { p.set(flat.subarray(off, off + p.length)); off += p.length; }
+  }
+
+  // greedy sampling — watch the model actually speak
+  async function generate(m, prompt, nChars) {
+    const { t: T } = m.cfg;
+    let ids = [...prompt].map(ch => STOI[ch] ?? 0);
+    for (let step = 0; step < nChars; step++) {
+      const win = ids.slice(-T);
+      const X = new Int32Array(T), Y = new Int32Array(T);
+      for (let i = 0; i < win.length; i++) X[T - win.length + i] = win[i];
+      const save = m.cfg.b; m.cfg.b = 1;
+      const { logits } = await forward(m, X, Y);
+      m.cfg.b = save;
+      const row = (T - 1) * m.cfg.vocab;
+      let best = 0, bv = -1e30;
+      for (let j = 0; j < m.cfg.vocab; j++) if (logits[row + j] > bv) { bv = logits[row + j]; best = j; }
+      ids.push(best);
+    }
+    return ids.map(i => CHARS[i]).join("");
+  }
+
+  const api = { init, trainStep, applyUpdate, getFlatParams, setFlatParams, generate, VOCAB, CORPUS };
+  if (typeof module !== "undefined" && module.exports) { TC = require("./traincore.js"); V = require("./verified_core.js"); module.exports = api; }
+  else { TC = root.TrainCore; V = root.Verified; root.Transformer = api; }
+})(typeof self !== "undefined" ? self : this);

web/test_transformer.js

@@ -0,0 +1,46 @@
+// Mini transformer through the verified units: (a) loss drops well below the
+// uniform baseline ln(V), (b) two replicas fed the same averaged gradients stay
+// bit-identical, (c) generation produces corpus-like text.
+const fs = require("fs");
+const path = require("path");
+const T = require("./public/traincore.js");
+const V = require("./public/verified_core.js");
+const X = require("./public/transformer.js");
+
+function loadLUTs() {
+  const p = (f) => path.join(__dirname, "public", f);
+  return { mul: new Int16Array(fs.readFileSync(p("mul_lut.bin")).buffer.slice(0)),
+           requant: new Int8Array(fs.readFileSync(p("requant_lut.bin")).buffer.slice(0)),
+           relu: new Int8Array(fs.readFileSync(p("relu_lut.bin")).buffer.slice(0)) };
+}
+const L = loadLUTs();
+const matmulInt8 = (Xq, Wq, m, k, n, LL) => V.lutMatmulJS(Xq, Wq, m, k, n, LL);
+const cfg = { c: 32, t: 32, b: 8, layers: 2, heads: 2, steps: 120, lr: 0.02 };
+
+(async function () {
+  const A = X.init(cfg, L, matmulInt8);
+  const B = X.init(cfg, L, matmulInt8);
+  console.log(`vocab=${X.VOCAB}, params=${A.nParams}, baseline loss=${Math.log(X.VOCAB).toFixed(3)}`);
+  const oa = T.makeAdam(A.nParams, { lr: cfg.lr });
+  const ob = T.makeAdam(B.nParams, { lr: cfg.lr });
+  let first = 0, loss = 0;
+  for (let s = 0; s < cfg.steps; s++) {
+    const ra = await X.trainStep(A);
+    const rb = await X.trainStep(B);
+    const avg = T.averageGrads([ra.grad, rb.grad]);
+    X.applyUpdate(A, oa.step(avg));
+    X.applyUpdate(B, ob.step(avg));
+    loss = (ra.loss + rb.loss) / 2;
+    if (s === 0) first = loss;
+    if (s % 30 === 0 || s === cfg.steps - 1) console.log(`  step ${s}  loss ${loss.toFixed(4)}`);
+  }
+  const pa = X.getFlatParams(A), pb = X.getFlatParams(B);
+  let diff = 0;
+  for (let i = 0; i < pa.length; i++) diff = Math.max(diff, Math.abs(pa[i] - pb[i]));
+  const sample = await X.generate(A, "the ", 60);
+  console.log(`sample: "${sample}"`);
+  console.log(`replica max param diff: ${diff.toExponential(3)}`);
+  const ok = loss < first * 0.75 && loss < Math.log(X.VOCAB) && diff === 0;
+  console.log(ok ? "TRANSFORMER TEST PASSED" : "TRANSFORMER TEST FAILED");
+  process.exit(ok ? 0 : 1);
+})();