@seanhogg/builderforce-memory-engine 2026.9.34 → 2026.10.2

package/src/ml/head_trainer.ts

@@ -0,0 +1,247 @@
+/**
+ * ml/head_trainer.ts — HeadTrainer: the ONE gradient trainer for a small head over
+ * an embedding (extracted from the limbic affect trainer, which is now a thin
+ * registration over it).
+ *
+ * Full-batch gradient descent with AdamW. When a GPUDevice is supplied the
+ * optimiser step runs on the GPU via the shared WEIGHT_UPDATE_WGSL kernel; with
+ * no device it runs the kernel's CPU mirror, `adamwUpdateInPlace` (`optim/adamw.ts`,
+ * NaN/Inf guard included), so training works everywhere (CI, Node without
+ * @webgpu/node, a laptop with no GPU). Neither path bounds the step (maxDelta 0).
+ *
+ * The head owns its forward/backward and its loss (MSE for regression, BCE /
+ * softmax cross-entropy for classification); this file owns only the loop:
+ * zero → accumulate per sample → average → clip → AdamW step.
+ */
+
+import {
+  createUniformBuffer,
+  createStorageBuffer,
+  createComputePipeline,
+  createBindGroup,
+  dispatchKernel,
+  readBuffer,
+  cdiv,
+} from "../utils/gpu_utils.js";
+import { WEIGHT_UPDATE_WGSL } from "../kernels/weight_update.js";
+import { adamwUpdateInPlace, type AdamWKernelStep } from "../optim/adamw.js";
+
+/** A named trainable tensor (flat row-major Float32Array). */
+export interface HeadParam {
+  name: string;
+  data: Float32Array;
+  numel: number;
+}
+
+/** What a head must expose to be trained by {@link HeadTrainer}. */
+export interface TrainableHead<S> {
+  /** Trainable parameters, in a stable order. */
+  parameters(): HeadParam[];
+  /** Gradient buffers, index-aligned with {@link parameters}. */
+  gradients(): HeadParam[];
+  zeroGrad(): void;
+  /** Accumulate gradients for one sample into {@link gradients}; returns its loss. */
+  accumulate(sample: S): number;
+  /** Forward-only loss of one sample. */
+  loss(sample: S): number;
+}
+
+export interface HeadTrainOptions {
+  learningRate?: number;
+  epochs?: number;
+  weightDecay?: number;
+  beta1?: number;
+  beta2?: number;
+  eps?: number;
+  /** Max global gradient L2 norm before the optimiser step. Default 1.0. */
+  maxGradNorm?: number;
+  /** Per-epoch hook. `gradNorm` is the pre-clip gradient L2 norm this epoch (the
+   *  instability early-warning) — computed CPU-side so it's always available. */
+  onEpochEnd?: ((epoch: number, loss: number, gradNorm?: number) => void) | null;
+}
+
+interface AdamMoment {
+  m: Float32Array;
+  v: Float32Array;
+}
+
+/** Byte size of WEIGHT_UPDATE_WGSL's `AdamParams` uniform: 1 u32 + 8 f32 + 3 f32 pad = 48. */
+const ADAM_PARAMS_BYTES = 48;
+
+/**
+ * Pack `AdamParams` exactly as WEIGHT_UPDATE_WGSL declares it:
+ * `num_elements: u32, lr, beta1, beta2, eps, weight_decay, beta1_t, beta2_t,
+ * max_delta, _pad0, _pad1, _pad2` (all f32 after the first) — 48 bytes.
+ */
+function packAdamParams(numElements: number, hp: AdamWKernelStep): ArrayBuffer {
+  const buf = new ArrayBuffer(ADAM_PARAMS_BYTES);
+  new Uint32Array(buf, 0, 1).set([numElements]);
+  // Pads stay 0 (fresh ArrayBuffer).
+  new Float32Array(buf, 4, 8).set([hp.lr, hp.beta1, hp.beta2, hp.eps, hp.weightDecay, hp.beta1_t, hp.beta2_t, hp.maxDelta]);
+  return buf;
+}
+
+export class HeadTrainer<S> {
+  readonly head: TrainableHead<S>;
+  readonly device: GPUDevice | null;
+  private _moments: AdamMoment[] | null = null;
+  private _step = 0;
+  private readonly _adamwPipeline: GPUComputePipeline | null;
+
+  constructor(
+    head: TrainableHead<S>,
+    device: GPUDevice | null = null,
+    /** Name used in error messages (the registration that owns this trainer). */
+    private readonly label: string = "HeadTrainer",
+  ) {
+    this.head = head;
+    this.device = device;
+    this._adamwPipeline = device
+      ? createComputePipeline(device, WEIGHT_UPDATE_WGSL, "adamw_update")
+      : null;
+  }
+
+  /** Whether the optimiser step runs on the GPU. */
+  get gpuTraining(): boolean {
+    return this.device != null && this._adamwPipeline != null;
+  }
+
+  private _initMoments(): void {
+    if (this._moments) return;
+    this._moments = this.head.parameters().map((p) => ({
+      m: new Float32Array(p.numel),
+      v: new Float32Array(p.numel),
+    }));
+  }
+
+  /**
+   * Train on `samples` for `epochs` passes; returns the per-epoch mean loss.
+   * Full-batch: grads accumulate over every sample, are averaged, clipped, then
+   * applied once per epoch.
+   */
+  async train(samples: S[], opts: HeadTrainOptions = {}): Promise<number[]> {
+    if (samples.length === 0) throw new Error(`${this.label}.train: no samples`);
+    const {
+      learningRate = 0.05,
+      epochs = 50,
+      weightDecay = 0.0,
+      beta1 = 0.9,
+      beta2 = 0.999,
+      eps = 1e-8,
+      maxGradNorm = 1.0,
+      onEpochEnd = null,
+    } = opts;
+
+    this._initMoments();
+    const losses: number[] = [];
+
+    for (let epoch = 0; epoch < epochs; epoch++) {
+      this.head.zeroGrad();
+      let epochLoss = 0;
+      for (const s of samples) epochLoss += this.head.accumulate(s);
+
+      const grads = this.head.gradients();
+      const invN = 1 / samples.length;
+      for (const g of grads) {
+        for (let i = 0; i < g.data.length; i++) g.data[i]! *= invN;
+      }
+
+      const gradNorm = clipGradients(grads, maxGradNorm);
+
+      this._step++;
+      const beta1_t = Math.pow(beta1, this._step);
+      const beta2_t = Math.pow(beta2, this._step);
+      const hp: AdamWKernelStep = { lr: learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t, maxDelta: 0 };
+      if (this.gpuTraining) {
+        await this._adamwStepGpu(grads, hp);
+      } else {
+        this._adamwStepCpu(grads, hp);
+      }
+
+      const avg = epochLoss / samples.length;
+      losses.push(avg);
+      if (onEpochEnd) onEpochEnd(epoch + 1, avg, gradNorm);
+    }
+    return losses;
+  }
+
+  /** Mean loss over samples (no weight update). 0 for an empty set. */
+  evaluate(samples: S[]): number {
+    if (samples.length === 0) return 0;
+    let total = 0;
+    for (const s of samples) total += this.head.loss(s);
+    return total / samples.length;
+  }
+
+  /** AdamW on the CPU via the kernel's shared mirror (`adamwUpdateInPlace`). */
+  private _adamwStepCpu(grads: HeadParam[], hp: AdamWKernelStep): void {
+    const params = this.head.parameters();
+    for (let pi = 0; pi < params.length; pi++) {
+      const mom = this._moments![pi]!;
+      adamwUpdateInPlace(params[pi]!.data, grads[pi]!.data, mom.m, mom.v, hp);
+    }
+  }
+
+  /** AdamW on the GPU via the shared WEIGHT_UPDATE_WGSL kernel. Awaited per step. */
+  private async _adamwStepGpu(grads: HeadParam[], hp: AdamWKernelStep): Promise<void> {
+    const device = this.device!;
+    const pipeline = this._adamwPipeline!;
+    const params = this.head.parameters();
+
+    for (let pi = 0; pi < params.length; pi++) {
+      const p = params[pi]!;
+      const mom = this._moments![pi]!;
+      const paramBuf = createStorageBuffer(device, p.data, true);
+      const gradBuf = createStorageBuffer(device, grads[pi]!.data, false);
+      const mBuf = createStorageBuffer(device, mom.m, true);
+      const vBuf = createStorageBuffer(device, mom.v, true);
+      const uni = createUniformBuffer(device, packAdamParams(p.numel, hp));
+      const bg = createBindGroup(device, pipeline, [uni, paramBuf, gradBuf, mBuf, vBuf]);
+      dispatchKernel(device, pipeline, bg, [cdiv(p.numel, 256), 1, 1]);
+
+      p.data.set((await readBuffer(device, paramBuf, p.numel * 4)).subarray(0, p.numel));
+      mom.m.set((await readBuffer(device, mBuf, p.numel * 4)).subarray(0, p.numel));
+      mom.v.set((await readBuffer(device, vBuf, p.numel * 4)).subarray(0, p.numel));
+
+      paramBuf.destroy();
+      gradBuf.destroy();
+      mBuf.destroy();
+      vBuf.destroy();
+      uni.destroy();
+    }
+  }
+}
+
+/** Clip grads to `maxNorm` in place; returns the PRE-clip L2 norm (the signal). */
+function clipGradients(grads: HeadParam[], maxNorm: number): number {
+  let normSq = 0;
+  for (const g of grads) for (let i = 0; i < g.data.length; i++) normSq += g.data[i]! * g.data[i]!;
+  const norm = Math.sqrt(normSq);
+  if (norm > maxNorm && norm > 0) {
+    const scale = maxNorm / norm;
+    for (const g of grads) for (let i = 0; i < g.data.length; i++) g.data[i]! *= scale;
+  }
+  return norm;
+}
+
+/** Concatenate a head's parameters into one flat vector (the federation unit). */
+export function flattenParams(params: readonly HeadParam[]): Float32Array {
+  const out = new Float32Array(params.reduce((n, p) => n + p.numel, 0));
+  let o = 0;
+  for (const p of params) {
+    out.set(p.data, o);
+    o += p.numel;
+  }
+  return out;
+}
+
+/** Write a flat vector back into a head's parameters (inverse of {@link flattenParams}). */
+export function assignParams(params: readonly HeadParam[], flat: ArrayLike<number>): void {
+  const total = params.reduce((n, p) => n + p.numel, 0);
+  if (flat.length !== total) throw new Error(`assignParams: expected ${total} values, got ${flat.length}`);
+  let o = 0;
+  for (const p of params) {
+    for (let i = 0; i < p.numel; i++) p.data[i] = flat[o + i]!;
+    o += p.numel;
+  }
+}

package/src/ml/index.ts

@@ -0,0 +1,57 @@
+/**
+ * ml — the ML task catalog (docs/ml-task-catalog.md): small, purpose-built models
+ * for bounded questions, answered on the person's own machine without a GPU or an
+ * LLM call. ONE registry, ONE runner, ONE trainer seam; trained heads federate on
+ * the existing delta wire. Pure and portable: no vscode, no DOM, no network.
+ */
+export type {
+  MlTask,
+  MlTaskKind,
+  MlEngineRef,
+  ClassicalAlgorithm,
+  InferenceMode,
+  MlPrediction,
+  MlPredictor,
+  FeatureSpec,
+  MlSample,
+  TrainSpec,
+  OutputSpec,
+} from "./types.js";
+export { engineKey, ML_TASK_ID_RE } from "./types.js";
+
+export { MlTaskRegistry } from "./registry.js";
+export {
+  runMlTask,
+  candidateEngines,
+  InMemoryMlModels,
+  MlDecisionLog,
+  MlEngineUnavailableError,
+} from "./run.js";
+export type { MlRunContext, MlResult, MlModelStore, MlCloudPort, MlDecision, MlDecisionSink, MlSkipReason } from "./run.js";
+
+export { trainMlTask, mlEvalLoss } from "./train.js";
+export type { MlTrainExample, MlTrainOptions, MlTrainResult } from "./train.js";
+
+export { HeadTrainer, flattenParams, assignParams } from "./head_trainer.js";
+export type { HeadParam, TrainableHead, HeadTrainOptions } from "./head_trainer.js";
+export { DenseHead } from "./dense_head.js";
+export type { DenseHeadConfig } from "./dense_head.js";
+
+export { extractFeatures, extractFeatureMatrix, logScale, clampRatio } from "./features.js";
+
+export * from "./classical/index.js";
+
+export {
+  headCheckpoint,
+  headParams,
+  mlHeadEntry,
+  mlHeadParams,
+  mlHeadUnusableReason,
+  packageMlHead,
+  withPackageMlHead,
+  buildMlHeadDelta,
+  mergeMlHeadDeltas,
+} from "./federate.js";
+export type { EvermindMlHead, MlHeadMergeOutcome } from "./federate.js";
+
+export * from "./tasks/index.js";

package/src/ml/registry.ts

@@ -0,0 +1,44 @@
+/**
+ * ml/registry.ts — the ONE registry of ML tasks, keyed by task id.
+ *
+ * A new task is a new row, not a new module. The registry only stores and looks
+ * up rows; running lives in `run.ts`, training in `train.ts`.
+ */
+import { ML_TASK_ID_RE, type MlTask } from "./types.js";
+
+export class MlTaskRegistry {
+  private readonly _tasks = new Map<string, MlTask<any, any>>();
+
+  /** Register a row. Throws on a malformed id, a duplicate id, or a row with no engine. */
+  register<I, O>(task: MlTask<I, O>): MlTask<I, O> {
+    if (!ML_TASK_ID_RE.test(task.id)) throw new Error(`MlTaskRegistry: malformed task id '${task.id}'`);
+    if (this._tasks.has(task.id)) throw new Error(`MlTaskRegistry: task '${task.id}' is already registered`);
+    if (task.engines.length === 0) throw new Error(`MlTaskRegistry: task '${task.id}' lists no engine`);
+    if (task.kind === "classify" && !(task.output.labels && task.output.labels.length >= 2)) {
+      throw new Error(`MlTaskRegistry: classify task '${task.id}' needs at least two output labels`);
+    }
+    this._tasks.set(task.id, task);
+    return task;
+  }
+
+  /** The row for `id`, or undefined. */
+  get<I = unknown, O = unknown>(id: string): MlTask<I, O> | undefined {
+    return this._tasks.get(id) as MlTask<I, O> | undefined;
+  }
+
+  /** The row for `id`; throws when it is not registered. */
+  require<I = unknown, O = unknown>(id: string): MlTask<I, O> {
+    const t = this.get<I, O>(id);
+    if (!t) throw new Error(`MlTaskRegistry: unknown task '${id}'`);
+    return t;
+  }
+
+  has(id: string): boolean {
+    return this._tasks.has(id);
+  }
+
+  /** Every row, in registration order. */
+  list(): MlTask<any, any>[] {
+    return [...this._tasks.values()];
+  }
+}

package/src/ml/run.ts

@@ -0,0 +1,222 @@
+/**
+ * ml/run.ts — the ONE runner: `runMlTask(id, input, ctx)`.
+ *
+ * Picks the first engine the inference mode allows that can answer (Local only:
+ * local tiers; Hybrid: local first, then the host's cloud port; Specific: exactly
+ * the named engine), returns `{ output, confidence, engine }`, and records the
+ * decision so routing can be evaluated later. The cloud is a PORT the host
+ * supplies — nothing here touches the network.
+ */
+import { extractFeatures, type TextEmbedder } from "./features.js";
+import type { MlTaskRegistry } from "./registry.js";
+import { engineKey, type InferenceMode, type MlEngineRef, type MlPrediction, type MlPredictor, type MlTask } from "./types.js";
+
+/** Fitted local models, per (task id, engine key). */
+export interface MlModelStore {
+  get(taskId: string, engine: string): MlPredictor | undefined;
+  set(taskId: string, engine: string, predictor: MlPredictor): void;
+}
+
+/** The plain in-process store. */
+export class InMemoryMlModels implements MlModelStore {
+  private readonly _m = new Map<string, MlPredictor>();
+  get(taskId: string, engine: string): MlPredictor | undefined {
+    return this._m.get(`${taskId}|${engine}`);
+  }
+  set(taskId: string, engine: string, predictor: MlPredictor): void {
+    this._m.set(`${taskId}|${engine}`, predictor);
+  }
+}
+
+/** A cloud engine, supplied by the host (Builderforce.ai passes its gateway client). */
+export interface MlCloudPort {
+  run(request: { taskId: string; engine: MlEngineRef; input: unknown }): Promise<MlPrediction>;
+}
+
+/** Why the runner moved past an engine. */
+export type MlSkipReason = "unavailable" | "low-confidence" | "error";
+
+/** One routing decision, recorded for eval. */
+export interface MlDecision {
+  taskId: string;
+  mode: InferenceMode;
+  /** The engine that answered (key), or null when none could. */
+  engine: string | null;
+  tier: "local" | "cloud" | null;
+  confidence: number | null;
+  /** Engines passed over before the answer, with why. */
+  skipped: { engine: string; reason: MlSkipReason; detail?: string }[];
+}
+
+export interface MlDecisionSink {
+  record(decision: MlDecision): void;
+}
+
+/** A bounded in-memory decision log with a per-task summary. */
+export class MlDecisionLog implements MlDecisionSink {
+  private readonly _d: MlDecision[] = [];
+  constructor(readonly capacity = 500) {}
+
+  record(decision: MlDecision): void {
+    this._d.push(decision);
+    if (this._d.length > this.capacity) this._d.splice(0, this._d.length - this.capacity);
+  }
+
+  list(): readonly MlDecision[] {
+    return this._d;
+  }
+
+  /** Per task: how many runs, which engines answered, mean confidence, unanswered count. */
+  summary(): Record<string, { runs: number; byEngine: Record<string, number>; unanswered: number; meanConfidence: number | null }> {
+    const out: Record<string, { runs: number; byEngine: Record<string, number>; unanswered: number; meanConfidence: number | null }> = {};
+    const sums: Record<string, { s: number; n: number }> = {};
+    for (const d of this._d) {
+      const t = (out[d.taskId] ??= { runs: 0, byEngine: {}, unanswered: 0, meanConfidence: null });
+      t.runs++;
+      if (d.engine == null) {
+        t.unanswered++;
+        continue;
+      }
+      t.byEngine[d.engine] = (t.byEngine[d.engine] ?? 0) + 1;
+      const c = (sums[d.taskId] ??= { s: 0, n: 0 });
+      c.s += d.confidence ?? 0;
+      c.n++;
+      t.meanConfidence = c.s / c.n;
+    }
+    return out;
+  }
+}
+
+export interface MlRunContext {
+  registry: MlTaskRegistry;
+  mode: InferenceMode;
+  /** For `specific`: the engine (ref or key) to use. */
+  engine?: MlEngineRef | string;
+  models?: MlModelStore;
+  embedder?: TextEmbedder;
+  cloud?: MlCloudPort;
+  /** Hybrid only: a local answer below this confidence defers to the cloud when it can answer. */
+  minConfidence?: number;
+  decisions?: MlDecisionSink;
+}
+
+export interface MlResult<O> {
+  output: O;
+  confidence: number;
+  engine: MlEngineRef;
+}
+
+/** No engine allowed by the mode could answer. */
+export class MlEngineUnavailableError extends Error {
+  constructor(
+    readonly taskId: string,
+    readonly decision: MlDecision,
+  ) {
+    super(
+      `runMlTask(${taskId}): no engine could answer in '${decision.mode}' mode (${
+        decision.skipped.map((s) => `${s.engine}: ${s.reason}${s.detail ? ` (${s.detail})` : ""}`).join(", ") || "none allowed"
+      })`,
+    );
+    this.name = "MlEngineUnavailableError";
+  }
+}
+
+/**
+ * The task's engines the mode allows, in the order the runner tries them. Exhaustive:
+ * a mode string outside {@link InferenceMode} (a host passing a stale or mistyped
+ * setting) throws instead of being read as `specific`.
+ */
+export function candidateEngines(task: MlTask<any, any>, mode: InferenceMode, specific?: MlEngineRef | string): MlEngineRef[] {
+  switch (mode) {
+    case "local-only":
+      return task.engines.filter((e) => e.tier === "local");
+    case "hybrid":
+      return [...task.engines.filter((e) => e.tier === "local"), ...task.engines.filter((e) => e.tier === "cloud")];
+    case "specific": {
+      if (specific == null) throw new Error(`runMlTask(${task.id}): 'specific' mode needs ctx.engine`);
+      const key = typeof specific === "string" ? specific : engineKey(specific);
+      const match = task.engines.filter((e) => engineKey(e) === key);
+      if (match.length === 0) throw new Error(`runMlTask(${task.id}): engine '${key}' is not one of this task's engines`);
+      return match;
+    }
+    default: {
+      const unknown: never = mode;
+      throw new Error(`runMlTask(${task.id}): unknown inference mode '${String(unknown)}'`);
+    }
+  }
+}
+
+const errorDetail = (error: unknown): string => (error instanceof Error ? error.message : String(error));
+
+/**
+ * Answer task `id` for `input` with the first engine the mode allows that can. An
+ * engine that throws (feature extraction or its own predict) is recorded as an
+ * `error` skip and the next allowed engine is tried — in Hybrid, the cloud. Throws
+ * {@link MlEngineUnavailableError}, carrying every engine's reason, only when none answers.
+ */
+export async function runMlTask<O = unknown, I = unknown>(id: string, input: I, ctx: MlRunContext): Promise<MlResult<O>> {
+  const task = ctx.registry.require<I, O>(id);
+  const engines = candidateEngines(task, ctx.mode, ctx.engine);
+  const skipped: MlDecision["skipped"] = [];
+  let features: Float32Array | null = null;
+  let featureError: unknown = null;
+  let deferred: { ref: MlEngineRef; p: MlPrediction } | null = null;
+
+  const finish = (ref: MlEngineRef, p: MlPrediction): MlResult<O> => {
+    ctx.decisions?.record({ taskId: id, mode: ctx.mode, engine: engineKey(ref), tier: ref.tier, confidence: p.confidence, skipped });
+    return { output: task.output.decode(p), confidence: p.confidence, engine: ref };
+  };
+
+  for (let i = 0; i < engines.length; i++) {
+    const ref = engines[i]!;
+    const key = engineKey(ref);
+    if (ref.tier === "local") {
+      const predictor = ctx.models?.get(id, key);
+      if (!predictor) {
+        skipped.push({ engine: key, reason: "unavailable" });
+        continue;
+      }
+      // Extract once; a failure (e.g. no embedder) fails every local engine alike, so
+      // it is remembered rather than retried per engine.
+      if (!features && featureError == null) {
+        try {
+          features = await extractFeatures(task.features, input, ctx.embedder);
+        } catch (error) {
+          featureError = error;
+        }
+      }
+      if (!features) {
+        skipped.push({ engine: key, reason: "error", detail: `features: ${errorDetail(featureError)}` });
+        continue;
+      }
+      let p: MlPrediction;
+      try {
+        p = predictor.predict(features);
+      } catch (error) {
+        skipped.push({ engine: key, reason: "error", detail: errorDetail(error) });
+        continue;
+      }
+      const cloudLater = ctx.mode === "hybrid" && !!ctx.cloud && engines.slice(i + 1).some((e) => e.tier === "cloud");
+      if (ctx.minConfidence != null && p.confidence < ctx.minConfidence && cloudLater) {
+        skipped.push({ engine: key, reason: "low-confidence", detail: p.confidence.toFixed(3) });
+        deferred ??= { ref, p };
+        continue;
+      }
+      return finish(ref, p);
+    }
+    if (!ctx.cloud) {
+      skipped.push({ engine: key, reason: "unavailable" });
+      continue;
+    }
+    try {
+      return finish(ref, await ctx.cloud.run({ taskId: id, engine: ref, input }));
+    } catch (error) {
+      skipped.push({ engine: key, reason: "error", detail: errorDetail(error) });
+    }
+  }
+  // The cloud could not answer after a low-confidence local one: the local answer stands.
+  if (deferred) return finish(deferred.ref, deferred.p);
+  const decision: MlDecision = { taskId: id, mode: ctx.mode, engine: null, tier: null, confidence: null, skipped };
+  ctx.decisions?.record(decision);
+  throw new MlEngineUnavailableError(id, decision);
+}