@seanhogg/builderforce-memory-engine 2026.7.12 → 2026.9.27

package/src/training/cpu_ops.ts

@@ -0,0 +1,127 @@
+/**
+ * cpu_ops.ts — the small set of exact CPU primitives the reference forward and
+ * backward passes share (linear, RMSNorm, SiLU, softplus).
+ *
+ * Each one mirrors a WGSL kernel exactly, so a gradient computed here is the
+ * gradient of the function the GPU actually evaluates. They are deliberately
+ * plain loops over `Float32Array`: correctness first, and small enough that the
+ * finite-difference check in `tests/backprop.test.ts` is meaningful.
+ */
+
+/** Epsilon used by `rmsnorm_forward` in `ACTIVATIONS_WGSL`. */
+export const RMSNORM_EPS = 1e-6;
+
+export function sigmoid(x: number): number {
+    return 1 / (1 + Math.exp(-x));
+}
+
+/** SiLU(x) = x * sigmoid(x) — mirrors `silu_forward`. */
+export function silu(x: number): number {
+    return x / (1 + Math.exp(-x));
+}
+
+/** d/dx SiLU(x) = sigmoid(x) * (1 + x * (1 - sigmoid(x))). */
+export function siluGrad(x: number): number {
+    const s = sigmoid(x);
+    return s * (1 + x * (1 - s));
+}
+
+/** Numerically stable softplus — mirrors the kernel's `max(x,0) + log1p(exp(-|x|))`. */
+export function softplus(x: number): number {
+    return Math.max(x, 0) + Math.log1p(Math.exp(-Math.abs(x)));
+}
+
+/**
+ * Y = X @ W^T + b, mirroring `linear_forward`.
+ *
+ * @param X   (M, K)
+ * @param W   (N, K) row-major
+ * @param b   (N,)
+ * @param out (M, N) — written, not accumulated
+ */
+export function linearForward(
+    X: Float32Array, W: Float32Array, b: Float32Array,
+    M: number, K: number, N: number, out: Float32Array,
+): void {
+    for (let r = 0; r < M; r++) {
+        const xo = r * K;
+        const yo = r * N;
+        for (let n = 0; n < N; n++) {
+            let acc = b[n]!;
+            const wo = n * K;
+            for (let k = 0; k < K; k++) acc += X[xo + k]! * W[wo + k]!;
+            out[yo + n] = acc;
+        }
+    }
+}
+
+/**
+ * Backward of {@link linearForward}. `dX` is written; `dW` and `db` are
+ * ACCUMULATED into (so one parameter can collect gradient from several call
+ * sites without a temporary).
+ */
+export function linearBackward(
+    dY: Float32Array, X: Float32Array, W: Float32Array,
+    M: number, K: number, N: number,
+    dX: Float32Array, dW: Float32Array, db: Float32Array,
+): void {
+    dX.fill(0);
+    for (let r = 0; r < M; r++) {
+        const xo = r * K;
+        const yo = r * N;
+        for (let n = 0; n < N; n++) {
+            const g = dY[yo + n]!;
+            if (g === 0) continue;
+            db[n] = db[n]! + g;
+            const wo = n * K;
+            for (let k = 0; k < K; k++) {
+                dX[xo + k] = dX[xo + k]! + g * W[wo + k]!;
+                dW[wo + k] = dW[wo + k]! + g * X[xo + k]!;
+            }
+        }
+    }
+}
+
+/**
+ * RMSNorm forward — mirrors `rmsnorm_forward`:
+ *   inv[r] = 1 / sqrt(mean(x[r]²) + eps);  y[r,i] = x[r,i] * inv[r] * w[i]
+ */
+export function rmsNormForward(
+    x: Float32Array, w: Float32Array, M: number, D: number,
+    out: Float32Array, inv: Float32Array, eps: number = RMSNORM_EPS,
+): void {
+    for (let r = 0; r < M; r++) {
+        const base = r * D;
+        let sq = 0;
+        for (let i = 0; i < D; i++) { const v = x[base + i]!; sq += v * v; }
+        const iv = 1 / Math.sqrt(sq / D + eps);
+        inv[r] = iv;
+        for (let i = 0; i < D; i++) out[base + i] = x[base + i]! * iv * w[i]!;
+    }
+}
+
+/**
+ * RMSNorm backward. `dx` and `dw` are ACCUMULATED into — `dx` in particular is
+ * usually the residual path, which already carries a gradient.
+ */
+export function rmsNormBackward(
+    dY: Float32Array, x: Float32Array, w: Float32Array, inv: Float32Array,
+    M: number, D: number,
+    dx: Float32Array, dw: Float32Array, _eps: number = RMSNORM_EPS,
+): void {
+    for (let r = 0; r < M; r++) {
+        const base = r * D;
+        const iv = inv[r]!;
+        let dot = 0;   // Σ_i (dY[r,i] * w[i]) * x[r,i]
+        for (let i = 0; i < D; i++) {
+            const g = dY[base + i]! * w[i]!;
+            dot += g * x[base + i]!;
+            dw[i] = dw[i]! + dY[base + i]! * x[base + i]! * iv;
+        }
+        const scale = (iv * iv * iv) / D;
+        for (let i = 0; i < D; i++) {
+            const g = dY[base + i]! * w[i]!;
+            dx[base + i] = dx[base + i]! + g * iv - scale * dot * x[base + i]!;
+        }
+    }
+}

package/src/training/lora.ts

@@ -0,0 +1,365 @@
+/**
+ * lora.ts — Low-Rank Adaptation (LoRA) + QLoRA for the EvermindLM CPU reference.
+ *
+ * The Meta llama-cookbook's central PEFT idea: freeze the base weights and train
+ * a tiny low-rank delta ΔW = (α/r)·B·A on top. Three payoffs, all realised here:
+ *
+ *   • Cheap — you train r·(rows+cols) params instead of rows·cols. For the tied
+ *     embedding (vocab×dModel) at rank 8 that is orders of magnitude smaller.
+ *   • Composable — the adapter serialises to a few KB (see {@link LoRAAdapter.serialize}),
+ *     so a persona / tenant / project is an MB-scale artifact you swap at load
+ *     time, not a full checkpoint.
+ *   • Forgetting-safe — the base never moves, so an adapter cannot catastrophically
+ *     overwrite the pretrained model (the exact property WSLA approximates with a
+ *     trust region; LoRA gets it structurally).
+ *
+ * QLoRA adds one more: the frozen base is held QUANTIZED (fp16 or int8) and
+ * dequantized on the fly for the merged forward, while the small adapter trains
+ * in fp32. On a single constrained device (our WebGPU target) the frozen base is
+ * where most bytes live, so quantizing it is the biggest memory unlock.
+ *
+ * Pure CPU, exact gradients (finite-difference checked in tests). The WGSL path
+ * is a future acceleration with the same shapes — same contract as EvermindLM.
+ */
+
+import { AdamW, type AdamWOptions } from "../optim/adamw.js";
+import { SeededRng } from "../utils/rng.js";
+import {
+  quantizeFp16,
+  dequantizeFp16,
+  quantizeInt8,
+  dequantizeInt8,
+} from "../utils/quantization.js";
+
+/** How the frozen base matrix is stored. `none` = fp32 (plain LoRA); the others = QLoRA. */
+export type BaseQuant = "none" | "fp16" | "int8";
+
+export interface LoRAConfig {
+  /** Low-rank bottleneck. Higher = more capacity, larger adapter. Default 8. */
+  rank?: number;
+  /** LoRA scaling; the delta is scaled by alpha/rank. Default = rank (unit scale). */
+  alpha?: number;
+  /** Deterministic init seed for the A matrix. */
+  seed?: number;
+}
+
+const ADAPTER_MAGIC = 0x4c4f5241; // "LORA"
+
+/**
+ * A low-rank adapter over one base matrix of shape [rows × cols] (row-major).
+ *
+ *   ΔW[i,j] = (α/r) · Σ_k B[i,k]·A[k,j]          B: [rows×r]  A: [r×cols]
+ *
+ * Standard LoRA init: A ~ N(0, σ²), B = 0, so ΔW starts at exactly zero and the
+ * adapted model equals the base until training moves it. Only A and B are
+ * trainable ({@link parameters}/{@link gradients} expose them to {@link AdamW});
+ * the base is supplied by the caller and never mutated here.
+ */
+export class LoRAAdapter {
+  readonly rows: number;
+  readonly cols: number;
+  readonly rank: number;
+  readonly alpha: number;
+  readonly scale: number;
+  /** [rows × rank], zero-initialised. */
+  readonly B: Float32Array;
+  /** [rank × cols], gaussian-initialised. */
+  readonly A: Float32Array;
+  private readonly gB: Float32Array;
+  private readonly gA: Float32Array;
+
+  constructor(rows: number, cols: number, config: LoRAConfig = {}) {
+    if (rows <= 0 || cols <= 0) throw new Error("LoRAAdapter: rows and cols must be > 0");
+    const rank = config.rank ?? 8;
+    if (rank <= 0 || rank > Math.min(rows, cols)) {
+      throw new Error(`LoRAAdapter: rank must be in [1, min(rows,cols)=${Math.min(rows, cols)}]`);
+    }
+    this.rows = rows;
+    this.cols = cols;
+    this.rank = rank;
+    this.alpha = config.alpha ?? rank;
+    this.scale = this.alpha / rank;
+
+    this.B = new Float32Array(rows * rank); // zero-init ⇒ ΔW starts at 0
+    this.A = new Float32Array(rank * cols);
+    const rng = new SeededRng((config.seed ?? 0x4c6f5241) >>> 0 || 1);
+    // Kaiming-ish small init for A (std ~ 1/sqrt(cols)).
+    const std = 1 / Math.sqrt(cols);
+    for (let i = 0; i < this.A.length; i++) {
+      const u1 = Math.max(rng.next(), 1e-12);
+      const u2 = rng.next();
+      this.A[i] = std * Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
+    }
+    this.gB = new Float32Array(this.B.length);
+    this.gA = new Float32Array(this.A.length);
+  }
+
+  /** Trainable count — the LoRA saving vs a full rows·cols matrix. */
+  numParams(): number {
+    return this.B.length + this.A.length;
+  }
+
+  /** Materialise ΔW = scale·B·A as a flat [rows·cols] row-major array. */
+  delta(): Float32Array {
+    const { rows, cols, rank, scale } = this;
+    const out = new Float32Array(rows * cols);
+    for (let i = 0; i < rows; i++) {
+      const bOff = i * rank;
+      const oOff = i * cols;
+      for (let k = 0; k < rank; k++) {
+        const b = this.B[bOff + k]! * scale;
+        if (b === 0) continue;
+        const aOff = k * cols;
+        for (let j = 0; j < cols; j++) out[oOff + j] = out[oOff + j]! + b * this.A[aOff + j]!;
+      }
+    }
+    return out;
+  }
+
+  /** Base + ΔW (new array; `base` is not mutated). */
+  applyTo(base: Float32Array): Float32Array {
+    if (base.length !== this.rows * this.cols) {
+      throw new Error("LoRAAdapter.applyTo: base length mismatch");
+    }
+    const d = this.delta();
+    const out = new Float32Array(base.length);
+    for (let i = 0; i < out.length; i++) out[i] = base[i]! + d[i]!;
+    return out;
+  }
+
+  /**
+   * Project dL/dW (flat [rows·cols], the gradient the base matrix WOULD receive)
+   * onto the adapter, accumulating dL/dA and dL/dB. Exact chain rule for
+   * ΔW = scale·B·A:  gA = scale·Bᵀ·G,  gB = scale·G·Aᵀ.
+   */
+  accumulateGradient(gW: Float32Array): void {
+    const { rows, cols, rank, scale } = this;
+    if (gW.length !== rows * cols) throw new Error("LoRAAdapter.accumulateGradient: gW length mismatch");
+    for (let i = 0; i < rows; i++) {
+      const gOff = i * cols;
+      const bOff = i * rank;
+      for (let k = 0; k < rank; k++) {
+        const aOff = k * cols;
+        let gb = 0;
+        const bScaled = this.B[bOff + k]! * scale;
+        for (let j = 0; j < cols; j++) {
+          const g = gW[gOff + j]!;
+          gb += g * this.A[aOff + j]!; // gB[i,k] = scale·Σ_j G[i,j]·A[k,j]
+          this.gA[aOff + j] = this.gA[aOff + j]! + bScaled * g; // gA[k,j] += scale·B[i,k]·G[i,j]
+        }
+        this.gB[bOff + k] = this.gB[bOff + k]! + scale * gb;
+      }
+    }
+  }
+
+  // AdamW OptimTarget surface — trains ONLY the adapter.
+  parameters(): { data: Float32Array }[] {
+    return [{ data: this.B }, { data: this.A }];
+  }
+  gradients(): { data: Float32Array }[] {
+    return [{ data: this.gB }, { data: this.gA }];
+  }
+  zeroGrad(): void {
+    this.gB.fill(0);
+    this.gA.fill(0);
+  }
+
+  /** Compact self-describing adapter blob (magic, rows, cols, rank, alpha, B, A). */
+  serialize(): ArrayBuffer {
+    const headerEls = 5;
+    const headerBytes = headerEls * 4;
+    const buf = new ArrayBuffer(headerBytes + (this.B.length + this.A.length) * 4);
+    const head = new Uint32Array(buf, 0, headerEls);
+    head[0] = ADAPTER_MAGIC;
+    head[1] = this.rows;
+    head[2] = this.cols;
+    head[3] = this.rank;
+    new Float32Array(buf, 16, 1)[0] = this.alpha;
+    const body = new Float32Array(buf, headerBytes);
+    body.set(this.B, 0);
+    body.set(this.A, this.B.length);
+    return buf;
+  }
+
+  static deserialize(buffer: ArrayBuffer): LoRAAdapter {
+    const head = new Uint32Array(buffer, 0, 5);
+    if (head[0] !== ADAPTER_MAGIC) throw new Error("LoRAAdapter.deserialize: bad magic");
+    const rows = head[1]!;
+    const cols = head[2]!;
+    const rank = head[3]!;
+    const alpha = new Float32Array(buffer, 16, 1)[0]!;
+    const adapter = new LoRAAdapter(rows, cols, { rank, alpha });
+    const body = new Float32Array(buffer, 20);
+    adapter.B.set(body.subarray(0, adapter.B.length));
+    adapter.A.set(body.subarray(adapter.B.length, adapter.B.length + adapter.A.length));
+    return adapter;
+  }
+}
+
+/** Quantize a base matrix for QLoRA storage; returns a dequantized view + byte cost. */
+export function quantizeBase(base: Float32Array, mode: BaseQuant): { view: Float32Array; bytes: number } {
+  switch (mode) {
+    case "none":
+      return { view: base, bytes: base.length * 4 };
+    case "fp16": {
+      const q = quantizeFp16(base);
+      return { view: dequantizeFp16(q), bytes: q.length * 2 };
+    }
+    case "int8": {
+      const q = quantizeInt8(base);
+      return { view: dequantizeInt8(q.data, q.scale), bytes: q.data.length + 4 };
+    }
+  }
+}
+
+import { EvermindLM } from "../lm/evermind_lm.js";
+
+export interface LoRAFitOptions extends AdamWOptions {
+  epochs?: number;
+  /**
+   * Gradient accumulation: average the adapter gradient over this many sequences
+   * (micro-batches) before each optimiser step. Lets a memory-constrained device
+   * train at a larger *effective* batch. Default 1 (step per sequence).
+   */
+  accumSteps?: number;
+}
+
+/**
+ * LoRA / QLoRA fine-tuning of an {@link EvermindLM} through its tied token
+ * embedding — the dominant parameter (vocab×dModel) and the natural adapter
+ * target (input lookup and output head share it). The base model is frozen; only
+ * the {@link LoRAAdapter} trains.
+ *
+ * QLoRA: pass `baseQuant: "fp16" | "int8"` and the frozen base is stored
+ * quantized and dequantized on the fly for the merged forward, so the resident
+ * base costs half (fp16) or a quarter (int8) of the bytes while the adapter
+ * trains full-precision.
+ *
+ * The adapter is the shippable artifact: {@link serializeAdapter} emits a few KB
+ * you swap per persona/tenant/project, versus rewriting the whole checkpoint.
+ */
+export class EvermindLMLoRA {
+  readonly adapter: LoRAAdapter;
+  readonly baseQuant: BaseQuant;
+  /** Frozen base embedding as used in the merged forward (dequantized under QLoRA). */
+  private readonly frozenBase: Float32Array;
+  private readonly baseBytesStored: number;
+  private readonly rows: number;
+  private readonly cols: number;
+
+  constructor(
+    private readonly model: EvermindLM,
+    config: LoRAConfig & { baseQuant?: BaseQuant } = {},
+  ) {
+    this.rows = model.config.vocabSize;
+    this.cols = model.config.dModel;
+    this.baseQuant = config.baseQuant ?? "none";
+    const { view, bytes } = quantizeBase(Float32Array.from(model.emb), this.baseQuant);
+    this.frozenBase = view;
+    this.baseBytesStored = bytes;
+    this.adapter = new LoRAAdapter(this.rows, this.cols, config);
+  }
+
+  /** The frozen base model this adapter rides on. */
+  get baseModel(): EvermindLM {
+    return this.model;
+  }
+
+  /** Effective embedding used for training/generation: frozenBase + adapter delta. */
+  mergedEmb(): Float32Array {
+    return this.adapter.applyTo(this.frozenBase);
+  }
+
+  /**
+   * One forward+backward on the merged weights, projecting the base-embedding
+   * gradient onto the adapter. The underlying model's own weights are left
+   * exactly as they were (base frozen); adapter gradients ACCUMULATE (caller
+   * zeroes between optimiser windows).
+   */
+  private _accumulate(seq: number[]): number {
+    const merged = this.mergedEmb();
+    const saved = this.model.emb;
+    this.model.emb = merged; // swap in effective weights (tied lookup + head)
+    this.model.zeroGrad();
+    const loss = this.model.lossAndBackward(seq);
+    const gW = this.model.gradients()[0]!.data; // dL/dW_eff == dL/dΔW, shape [rows×cols]
+    this.adapter.accumulateGradient(gW);
+    this.model.emb = saved; // restore the frozen base
+    return loss;
+  }
+
+  /** Train the adapter (only) with AdamW + optional gradient accumulation. */
+  fit(sequences: number[][], opts: LoRAFitOptions = {}): number[] {
+    const epochs = opts.epochs ?? 1;
+    const accum = Math.max(1, opts.accumSteps ?? 1);
+    const adam = new AdamW(this.adapter, opts);
+    const grads = this.adapter.gradients();
+    const history: number[] = [];
+    for (let e = 0; e < epochs; e++) {
+      let total = 0;
+      let n = 0;
+      let pending = 0;
+      this.adapter.zeroGrad();
+      const flush = () => {
+        if (pending === 0) return;
+        if (pending !== 1) for (const g of grads) for (let i = 0; i < g.data.length; i++) g.data[i] = g.data[i]! / pending;
+        adam.step();
+        this.adapter.zeroGrad();
+        pending = 0;
+      };
+      for (const seq of sequences) {
+        if (seq.length < 2) continue;
+        total += this._accumulate(seq);
+        n++;
+        if (++pending >= accum) flush();
+      }
+      flush();
+      history.push(n > 0 ? total / n : 0);
+    }
+    return history;
+  }
+
+  generate(prompt: number[], opts: import("../lm/evermind_lm.js").LMGenerateOptions): number[] {
+    const saved = this.model.emb;
+    this.model.emb = this.mergedEmb();
+    try {
+      return this.model.generate(prompt, opts);
+    } finally {
+      this.model.emb = saved;
+    }
+  }
+
+  generateText(prompt: string, codec: import("../lm/evermind_lm.js").TextCodec, opts: import("../lm/evermind_lm.js").LMGenerateOptions): string {
+    return codec.decode(this.generate(codec.encode(prompt), opts));
+  }
+
+  /** The shippable adapter artifact (a few KB). */
+  serializeAdapter(): ArrayBuffer {
+    return this.adapter.serialize();
+  }
+
+  /** Reconstruct a fine-tuned model from a base model + a serialized adapter. */
+  static loadAdapter(model: EvermindLM, adapterBuffer: ArrayBuffer, baseQuant: BaseQuant = "none"): EvermindLMLoRA {
+    const loaded = LoRAAdapter.deserialize(adapterBuffer);
+    const wrap = new EvermindLMLoRA(model, { rank: loaded.rank, alpha: loaded.alpha, baseQuant });
+    wrap.adapter.B.set(loaded.B);
+    wrap.adapter.A.set(loaded.A);
+    return wrap;
+  }
+
+  /** Bake the adapter into the base and return a standalone EVL0 checkpoint. */
+  merge(opts: { fp16?: boolean } = {}): ArrayBuffer {
+    this.model.emb.set(this.mergedEmb());
+    return this.model.exportWeights(opts);
+  }
+
+  /** Byte cost of the trainable adapter vs the frozen base — the LoRA/QLoRA saving. */
+  footprint(): { adapterBytes: number; baseBytes: number; trainableParams: number; baseParams: number } {
+    return {
+      adapterBytes: this.adapter.serialize().byteLength,
+      baseBytes: this.baseBytesStored,
+      trainableParams: this.adapter.numParams(),
+      baseParams: this.rows * this.cols,
+    };
+  }
+}