@seanhogg/builderforce-memory-engine 2026.7.12 → 2026.7.14

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,
+    };
+  }
+}

package/src/training/mixed_precision.ts

@@ -0,0 +1,119 @@
+/**
+ * mixed_precision.ts — mixed-precision training discipline (AMP) for the CPU
+ * reference, faithful to the cookbook's recipe:
+ *
+ *   • fp16 compute, fp32 MASTER weights. The model's Float32Array params ARE the
+ *     master copy; the forward runs on an fp16-rounded view, so tiny weight
+ *     updates still accumulate in fp32 instead of vanishing.
+ *   • Dynamic LOSS SCALING. fp16 gradients underflow to zero below ~6e-5; scaling
+ *     the loss by a large factor S lifts them into fp16 range, then we unscale by
+ *     S before the optimiser step. On overflow (inf/nan) we skip the step and back
+ *     the scale off; after a clean streak we grow it. This is the mechanism that
+ *     makes fp16 training numerically stable.
+ *
+ * The payoff is memory/throughput: activations and gradients at half precision.
+ * On CPU we can't reclaim the bytes, so this module models the NUMERICS exactly
+ * (fp16 rounding + loss scaling) — the property the WGSL path will inherit — and
+ * its tests prove a small gradient survives WITH scaling and underflows WITHOUT.
+ */
+
+import { floatToFp16, fp16ToFloat } from "../utils/quantization.js";
+
+/** Round a value to what fp16 storage would preserve (half-precision simulation). */
+export function roundFp16(x: number): number {
+  return fp16ToFloat(floatToFp16(x));
+}
+
+/** Round every element of a copy to fp16 precision (leaves the input untouched). */
+export function fp16View(a: Float32Array): Float32Array {
+  const out = new Float32Array(a.length);
+  for (let i = 0; i < a.length; i++) out[i] = roundFp16(a[i]!);
+  return out;
+}
+
+export interface LossScalerOptions {
+  /** Initial scale. Default 2^14 = 16384. */
+  initScale?: number;
+  /** Multiply the scale by this after `growthInterval` clean steps. Default 2. */
+  growthFactor?: number;
+  /** Multiply the scale by this on overflow. Default 0.5. */
+  backoffFactor?: number;
+  /** Clean steps between growth attempts. Default 200. */
+  growthInterval?: number;
+  /** Floor for the scale so it can't collapse to zero. Default 1. */
+  minScale?: number;
+}
+
+/**
+ * Dynamic loss scaler — the standard AMP controller. Scale the loss/grads up
+ * before backward, {@link check} them for overflow, {@link unscale} the good ones,
+ * and {@link update} the scale based on whether this step overflowed.
+ */
+export class DynamicLossScaler {
+  private _scale: number;
+  private readonly growthFactor: number;
+  private readonly backoffFactor: number;
+  private readonly growthInterval: number;
+  private readonly minScale: number;
+  private cleanStreak = 0;
+  private _overflows = 0;
+  private _skipped = 0;
+
+  constructor(opts: LossScalerOptions = {}) {
+    this._scale = opts.initScale ?? 16384;
+    this.growthFactor = opts.growthFactor ?? 2;
+    this.backoffFactor = opts.backoffFactor ?? 0.5;
+    this.growthInterval = opts.growthInterval ?? 200;
+    this.minScale = opts.minScale ?? 1;
+  }
+
+  get scale(): number {
+    return this._scale;
+  }
+  get overflows(): number {
+    return this._overflows;
+  }
+  get skipped(): number {
+    return this._skipped;
+  }
+
+  /** True if any gradient is inf/nan (an fp16 overflow) — such a step must be skipped. */
+  check(grads: { data: Float32Array }[]): boolean {
+    for (const g of grads) {
+      const a = g.data;
+      for (let i = 0; i < a.length; i++) {
+        const x = a[i]!;
+        if (!Number.isFinite(x)) return true;
+      }
+    }
+    return false;
+  }
+
+  /** Divide gradients back down by the current scale, in place. */
+  unscale(grads: { data: Float32Array }[]): void {
+    const inv = 1 / this._scale;
+    for (const g of grads) {
+      const a = g.data;
+      for (let i = 0; i < a.length; i++) a[i] = a[i]! * inv;
+    }
+  }
+
+  /**
+   * Advance the scale controller. Pass whether this step overflowed. Returns
+   * `true` if the optimiser step should PROCEED, `false` if it must be skipped.
+   */
+  update(overflow: boolean): boolean {
+    if (overflow) {
+      this._overflows++;
+      this._skipped++;
+      this._scale = Math.max(this.minScale, this._scale * this.backoffFactor);
+      this.cleanStreak = 0;
+      return false;
+    }
+    if (++this.cleanStreak >= this.growthInterval) {
+      this._scale *= this.growthFactor;
+      this.cleanStreak = 0;
+    }
+    return true;
+  }
+}