@seanhogg/builderforce-memory-engine 2026.6.39 → 2026.7.11

package/src/moe/moe_package.ts

@@ -13,6 +13,7 @@
 
 import { SharedExpertMoE, type MoEConfig } from "./moe_model.js";
 import { EvermindLM, type EvermindLMConfig } from "../lm/evermind_lm.js";
+import { VideoRVQCodec } from "../codec/video_rvq.js";
 
 /** First 4 bytes of a serialised package: "EVM1". */
 const PKG_MAGIC = 0x45564d31;
@@ -21,6 +22,9 @@ const PKG_VERSION = 1;
 /** The kinds of model an `.evermind` package can carry. */
 export type EvermindModelType = "shared-expert-moe" | "evermind-lm";
 
+/** Output modality. `text` (default) needs no codec; `video`/`image` bundle a VRQ codec. */
+export type EvermindModality = "text" | "video" | "image";
+
 /** Human-facing description published with the model (the "model card"). */
 export interface EvermindModelCard {
   description: string;
@@ -38,14 +42,26 @@ export interface EvermindModelManifest {
   name: string;
   version: string;
   modelType: EvermindModelType;
+  /** Output modality; absent ⇒ "text". Video/image also carry a `codec` section. */
+  modality?: EvermindModality;
   /** Flat numeric model config (the constructor args), serialised verbatim. */
   config: Record<string, number>;
   /** Total trainable scalar parameters (for sizing / pricing / display). */
   paramCount: number;
   checkpointFormat: "MoE0" | "EVL0";
   checkpointFp16: boolean;
+  /**
+   * Byte length of the checkpoint section — REQUIRED when a codec section follows
+   * (so the reader knows where the checkpoint ends). Absent ⇒ checkpoint runs to
+   * end-of-blob (the original text-only layout, still read verbatim).
+   */
+  checkpointBytes?: number;
   /** 32-bit FNV-1a over the checkpoint bytes — integrity for download/purchase. */
   checksum: number;
+  /** Media codec section format (present for video/image packages). */
+  codecFormat?: "VRQ0";
+  /** 32-bit FNV-1a over the codec bytes — integrity for the bundled codec. */
+  codecChecksum?: number;
   card: EvermindModelCard;
   /** ISO timestamp, caller-supplied (the engine avoids Date for determinism). */
   createdAt?: string;
@@ -84,6 +100,8 @@ export class EvermindModelPackage {
   constructor(
     readonly manifest: EvermindModelManifest,
     readonly checkpoint: ArrayBuffer,
+    /** Serialized media codec ("VRQ0" blob), present for video/image packages. */
+    readonly codec?: ArrayBuffer,
   ) {}
 
   /** Package a trained model with its publishing metadata. */
@@ -126,17 +144,61 @@ export class EvermindModelPackage {
     return new EvermindModelPackage(manifest, checkpoint);
   }
 
-  /** Serialise to a single `.evermind` blob: magic, version, manifest, checkpoint. */
+  /**
+   * Package a video/image generative model: its {@link EvermindLM} weights AND the
+   * {@link VideoRVQCodec} needed to turn generated tokens back into frames. Without
+   * the bundled codec a served media model would emit undecodable token ids — this
+   * is what makes an `.evermind` media artifact self-contained and servable.
+   */
+  static fromMediaLM(
+    lm: EvermindLM,
+    codec: VideoRVQCodec,
+    meta: PackageMeta & { modality: "video" | "image" },
+  ): EvermindModelPackage {
+    if (lm.config.vocabSize !== codec.vocabSize) {
+      throw new Error(
+        `fromMediaLM: LM vocabSize (${lm.config.vocabSize}) must equal codec.vocabSize (${codec.vocabSize})`,
+      );
+    }
+    const fp16 = meta.fp16 ?? false;
+    const checkpoint = lm.exportWeights({ fp16 });
+    const codecBlob = codec.serialize();
+    const manifest: EvermindModelManifest = {
+      schema: "evermind.model/1",
+      name: meta.name,
+      version: meta.version,
+      modelType: "evermind-lm",
+      modality: meta.modality,
+      config: lm.config as unknown as Record<string, number>,
+      paramCount: lm.parameters().reduce((n, p) => n + p.data.length, 0),
+      checkpointFormat: "EVL0",
+      checkpointFp16: fp16,
+      checkpointBytes: checkpoint.byteLength,
+      checksum: fnv1a(new Uint8Array(checkpoint)),
+      codecFormat: "VRQ0",
+      codecChecksum: fnv1a(new Uint8Array(codecBlob)),
+      card: meta.card,
+      ...(meta.createdAt ? { createdAt: meta.createdAt } : {}),
+    };
+    return new EvermindModelPackage(manifest, checkpoint, codecBlob);
+  }
+
+  /** Serialise to a single `.evermind` blob: magic, version, manifest, checkpoint[, codec]. */
   toBlob(): ArrayBuffer {
     const manifestBytes = new TextEncoder().encode(JSON.stringify(this.manifest));
     const headerBytes = 12; // magic, version, manifestLen
-    const out = new ArrayBuffer(headerBytes + manifestBytes.byteLength + this.checkpoint.byteLength);
+    const codecBytes = this.codec?.byteLength ?? 0;
+    const out = new ArrayBuffer(headerBytes + manifestBytes.byteLength + this.checkpoint.byteLength + codecBytes);
     const head = new Uint32Array(out, 0, 3);
     head[0] = PKG_MAGIC;
     head[1] = PKG_VERSION;
     head[2] = manifestBytes.byteLength;
-    new Uint8Array(out, headerBytes, manifestBytes.byteLength).set(manifestBytes);
-    new Uint8Array(out, headerBytes + manifestBytes.byteLength).set(new Uint8Array(this.checkpoint));
+    let o = headerBytes;
+    new Uint8Array(out, o, manifestBytes.byteLength).set(manifestBytes);
+    o += manifestBytes.byteLength;
+    new Uint8Array(out, o, this.checkpoint.byteLength).set(new Uint8Array(this.checkpoint));
+    o += this.checkpoint.byteLength;
+    if (this.codec) new Uint8Array(out, o, codecBytes).set(new Uint8Array(this.codec));
     return out;
   }
 
@@ -152,8 +214,19 @@ export class EvermindModelPackage {
     }
     const manifestBytes = new Uint8Array(buffer, headerBytes, manifestLen);
     const manifest = JSON.parse(new TextDecoder().decode(manifestBytes)) as EvermindModelManifest;
-    const checkpoint = buffer.slice(headerBytes + manifestLen);
-    return new EvermindModelPackage(manifest, checkpoint);
+    const bodyStart = headerBytes + manifestLen;
+    // A codec section follows the checkpoint only when the manifest fixed the
+    // checkpoint length; otherwise the checkpoint runs to end-of-blob (text layout).
+    const cpBytes = manifest.checkpointBytes;
+    if (cpBytes != null) {
+      if (bodyStart + cpBytes > buffer.byteLength) {
+        throw new Error("EvermindModelPackage.fromBlob: truncated (checkpointBytes exceeds blob)");
+      }
+      const checkpoint = buffer.slice(bodyStart, bodyStart + cpBytes);
+      const codec = bodyStart + cpBytes < buffer.byteLength ? buffer.slice(bodyStart + cpBytes) : undefined;
+      return new EvermindModelPackage(manifest, checkpoint, codec);
+    }
+    return new EvermindModelPackage(manifest, buffer.slice(bodyStart));
   }
 
   /** Verify integrity + structural sanity before trusting a downloaded package. */
@@ -176,6 +249,19 @@ export class EvermindModelPackage {
     if (actual !== this.manifest.checksum) {
       errors.push(`checksum mismatch (manifest ${this.manifest.checksum}, actual ${actual}) — corrupt or tampered checkpoint`);
     }
+    // Media packages must carry an integral codec — a served media model can't
+    // decode tokens → frames without it.
+    const modality = this.manifest.modality ?? "text";
+    if (modality === "video" || modality === "image") {
+      if (!this.codec) {
+        errors.push(`modality '${modality}' package is missing its codec section`);
+      } else {
+        const codecActual = fnv1a(new Uint8Array(this.codec));
+        if (codecActual !== this.manifest.codecChecksum) {
+          errors.push(`codec checksum mismatch (manifest ${this.manifest.codecChecksum}, actual ${codecActual}) — corrupt codec`);
+        }
+      }
+    }
     return { ok: errors.length === 0, errors };
   }
 
@@ -202,4 +288,19 @@ export class EvermindModelPackage {
     lm.loadWeights(this.checkpoint);
     return lm;
   }
+
+  /**
+   * Reconstruct a video/image model: the generative {@link EvermindLM} plus its
+   * {@link VideoRVQCodec}. Pair with `generateVideo` / `generateImage` to produce
+   * frames. This is what a media-serving path (gateway) loads and runs.
+   */
+  loadMediaLM(): { lm: EvermindLM; codec: VideoRVQCodec; modality: EvermindModality } {
+    const v = this.validate();
+    if (!v.ok) throw new Error(`EvermindModelPackage.loadMediaLM: ${v.errors.join("; ")}`);
+    const modality = this.manifest.modality ?? "text";
+    if (modality !== "video" && modality !== "image") {
+      throw new Error(`loadMediaLM: package modality is '${modality}', use loadLM()`);
+    }
+    return { lm: this.loadLM(), codec: VideoRVQCodec.deserialize(this.codec!), modality };
+  }
 }

package/src/training/trainer.ts

@@ -9,6 +9,7 @@ import {
     createComputePipeline,
     createBindGroup,
     dispatchKernel,
+    readBuffer,
     cdiv,
 } from '../utils/gpu_utils.js';
 
@@ -29,9 +30,26 @@ export interface TrainOptions {
   beta2?: number;
   eps?: number;
   wsla?: boolean;
-  onEpochEnd?: ((epoch: number, loss: number) => void) | null;
+  /** Trust region: max |Δθ| per optimizer step. 0 disables. Defaults to
+   *  {@link WSLA_MAX_DELTA} in WSLA (write-through) mode, else 0 (full training). */
+  maxDelta?: number;
+  /** Read back the pre-clip gradient L2 norm each step (the `grad_norm_reduce`
+   *  kernel already computes it — this just maps it to the CPU). OFF by default: the
+   *  readback is a GPU sync point that would slow a long finetune. `adapt()` turns it
+   *  ON so the norm reaches {@link AdaptResult} as an instability signal. */
+  trackGradNorm?: boolean;
+  onEpochEnd?: ((epoch: number, loss: number, gradNorm?: number) => void) | null;
 }
 
+/**
+ * Default per-step trust region for WSLA / write-through adaptation. Small
+ * enough that a single `adapt()` nudges the narrow params without lurching, so
+ * repeated adapts stay stable (and any that regress are cheaply rolled back by
+ * the session). Full-training callers pass `maxDelta: 0` (or omit it in
+ * non-WSLA mode) for unbounded steps.
+ */
+export const WSLA_MAX_DELTA = 0.05;
+
 interface AdamMoments {
   m: GPUBuffer;
   v: GPUBuffer;
@@ -45,13 +63,22 @@ interface AdamHyperparams {
   eps: number;
   beta1_t: number;
   beta2_t: number;
+  maxDelta: number;
 }
 
 export class MambaTrainer {
     model: HybridMambaModel;
     tokenizer: BPETokenizer | null;
     device: GPUDevice;
-    private _moments: AdamMoments[] | null;
+    /**
+     * Adam moments keyed by parameter NAME (not array index). Name-keying is what
+     * lets WSLA toggle safely: a narrow write-through step updates only the
+     * `layer{i}.wXProj/bXProj` subset, a full fine-tune updates everything, and
+     * both reuse the SAME `m`/`v` buffers per parameter. Index-keyed moments (the
+     * old scheme) silently misaligned the moment with the wrong parameter the
+     * moment the trainable set changed shape — corrupting the update.
+     */
+    private _moments: Map<string, AdamMoments>;
     private _step: number;
     private _adamwPipeline: GPUComputePipeline;
     private _clipReducePipeline: GPUComputePipeline;
@@ -62,7 +89,7 @@ export class MambaTrainer {
         this.tokenizer = tokenizer;
         this.device    = model.device;
 
-        this._moments = null;
+        this._moments = new Map();
         this._step = 0;
 
         this._adamwPipeline   = createComputePipeline(this.device, WEIGHT_UPDATE_WGSL, 'adamw_update');
@@ -70,12 +97,17 @@ export class MambaTrainer {
         this._clipScalePipeline  = createComputePipeline(this.device, GRAD_CLIP_WGSL, 'grad_clip_scale');
     }
 
-    private _initMoments(): void {
-        if (this._moments) return;
-        this._moments = this.model.parameters().map(p => ({
-            m: createEmptyStorageBuffer(this.device, p.numel * 4, false),
-            v: createEmptyStorageBuffer(this.device, p.numel * 4, false),
-        }));
+    /** Get-or-create the Adam first/second moments for a parameter, by name. */
+    private _momentFor(p: BlockParam): AdamMoments {
+        let mom = this._moments.get(p.name);
+        if (!mom) {
+            mom = {
+                m: createEmptyStorageBuffer(this.device, p.numel * 4, false),
+                v: createEmptyStorageBuffer(this.device, p.numel * 4, false),
+            };
+            this._moments.set(p.name, mom);
+        }
+        return mom;
     }
 
     async train(input: string | number[], opts: TrainOptions = {}): Promise<number[]> {
@@ -90,8 +122,11 @@ export class MambaTrainer {
             beta2        = 0.999,
             eps          = 1e-8,
             wsla         = false,
+            trackGradNorm = false,
             onEpochEnd   = null,
         } = opts;
+        // Trust region defaults ON for write-through (WSLA) adapts, OFF otherwise.
+        const maxDelta = opts.maxDelta ?? (wsla ? WSLA_MAX_DELTA : 0);
 
         if (wsla) this.model.setWSLAMode(true);
 
@@ -117,27 +152,33 @@ export class MambaTrainer {
             throw new Error('Input is too short to form any training chunk.');
         }
 
-        this._initMoments();
+        // Moments are created lazily per-parameter (name-keyed) on first touch —
+        // no upfront allocation needed, and it stays correct as WSLA narrows the
+        // trainable set.
 
         const epochLosses: number[] = [];
 
         for (let epoch = 0; epoch < epochs; epoch++) {
             let epochLoss = 0;
+            let epochGradNorm = 0;
             let numSteps  = 0;
 
             for (const { inputs, targets } of chunks) {
-                const loss = await this._trainStep(
+                const { loss, gradNorm } = await this._trainStep(
                     inputs, targets, batchSize,
-                    { learningRate, maxGradNorm, weightDecay, beta1, beta2, eps, wsla }
+                    { learningRate, maxGradNorm, weightDecay, beta1, beta2, eps, wsla, maxDelta, trackGradNorm }
                 );
                 epochLoss += loss;
+                if (gradNorm != null) epochGradNorm += gradNorm;
                 numSteps++;
             }
 
             const avgLoss = epochLoss / numSteps;
             epochLosses.push(avgLoss);
 
-            if (onEpochEnd) onEpochEnd(epoch + 1, avgLoss);
+            // Mean pre-clip grad norm this epoch (only measured when trackGradNorm is on).
+            const avgGradNorm = trackGradNorm ? epochGradNorm / numSteps : undefined;
+            if (onEpochEnd) onEpochEnd(epoch + 1, avgLoss, avgGradNorm);
         }
 
         if (wsla) this.model.setWSLAMode(false);
@@ -148,9 +189,9 @@ export class MambaTrainer {
         inputs: number[],
         targets: number[],
         batch: number,
-        hyperparams: TrainOptions & { learningRate: number; maxGradNorm: number; weightDecay: number; beta1: number; beta2: number; eps: number }
-    ): Promise<number> {
-        const { learningRate, maxGradNorm, weightDecay, beta1, beta2, eps } = hyperparams;
+        hyperparams: TrainOptions & { learningRate: number; maxGradNorm: number; weightDecay: number; beta1: number; beta2: number; eps: number; maxDelta: number; trackGradNorm?: boolean }
+    ): Promise<{ loss: number; gradNorm: number | null }> {
+        const { learningRate, maxGradNorm, weightDecay, beta1, beta2, eps, maxDelta, trackGradNorm } = hyperparams;
 
         this._step++;
         const seqLen    = inputs.length;
@@ -177,21 +218,24 @@ export class MambaTrainer {
 
         const dLogitsBuf = createStorageBuffer(this.device, dLogits, false);
 
-        await this._clipGradients(dLogitsBuf, dLogits.length, maxGradNorm);
+        const gradNorm = await this._clipGradients(dLogitsBuf, dLogits.length, maxGradNorm, !!trackGradNorm);
 
-        const params  = this.model.parameters();
+        // Only the trainable set is updated. Under WSLA (write-through) that is
+        // the narrow per-layer subset and the backbone (incl. every A_log) stays
+        // frozen — the guarantee that repeated adapts can't destabilise the SSM.
+        const params  = this.model.getTrainableParams();
         const beta1_t = Math.pow(beta1, this._step);
         const beta2_t = Math.pow(beta2, this._step);
 
         await this._adamwStep(
             params, [dLogitsBuf],
-            { learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t }
+            { learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t, maxDelta }
         );
 
         dLogitsBuf.destroy();
         gpuLogits.destroy();
 
-        return loss;
+        return { loss, gradNorm };
     }
 
     private async _adamwStep(
@@ -199,7 +243,7 @@ export class MambaTrainer {
         gradBufs: GPUBuffer[],
         hp: AdamHyperparams
     ): Promise<void> {
-        const { learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t } = hp;
+        const { learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t, maxDelta } = hp;
 
         for (let i = 0; i < params.length; i++) {
             const p       = params[i]!;
@@ -207,16 +251,17 @@ export class MambaTrainer {
 
             if (!gradBuf || gradBuf.size < p.numel * 4) continue;
 
+            const mom = this._momentFor(p);
             const paramsBuf = createUniformBuffer(this.device, packAdamParams(
-                p.numel, learningRate, beta1, beta2, eps, weightDecay, beta1_t, beta2_t
+                p.numel, learningRate, beta1, beta2, eps, weightDecay, beta1_t, beta2_t, maxDelta
             ));
 
             const bg = createBindGroup(this.device, this._adamwPipeline, [
                 paramsBuf,
                 p.buf,
                 gradBuf,
-                this._moments![i]!.m,
-                this._moments![i]!.v,
+                mom.m,
+                mom.v,
             ]);
 
             dispatchKernel(this.device, this._adamwPipeline, bg,
@@ -226,7 +271,14 @@ export class MambaTrainer {
         }
     }
 
-    private async _clipGradients(gradBuf: GPUBuffer, numel: number, maxNorm: number): Promise<void> {
+    /**
+     * Clip gradients to `maxNorm` in place. The reduce kernel writes Σg² into
+     * `normSqBuf`; the scale kernel then rescales the grads. When `trackNorm` is set
+     * we read Σg² back and return the true pre-clip L2 norm (√Σg²) — the kernel
+     * already computed it, so this is only a small buffer readback (a GPU sync,
+     * hence opt-in). Returns null when not tracking.
+     */
+    private async _clipGradients(gradBuf: GPUBuffer, numel: number, maxNorm: number, trackNorm = false): Promise<number | null> {
         const normSqBuf = createEmptyStorageBuffer(this.device, 4, true);
         this.device.queue.writeBuffer(normSqBuf, 0, new Float32Array([0.0]));
 
@@ -240,6 +292,14 @@ export class MambaTrainer {
         dispatchKernel(this.device, this._clipReducePipeline, bg1,
             [cdiv(numel, 256), 1, 1]);
 
+        // Read Σg² AFTER the reduce (the scale kernel only reads it, never overwrites),
+        // so the returned norm is the pre-clip magnitude — the instability signal.
+        let gradNorm: number | null = null;
+        if (trackNorm) {
+            const out = await readBuffer(this.device, normSqBuf, 4);
+            gradNorm = Math.sqrt(Math.max(0, out[0] ?? 0));
+        }
+
         const bg2 = createBindGroup(this.device, this._clipScalePipeline,
             [pBuf, gradBuf, normSqBuf]);
         dispatchKernel(this.device, this._clipScalePipeline, bg2,
@@ -247,6 +307,7 @@ export class MambaTrainer {
 
         pBuf.destroy();
         normSqBuf.destroy();
+        return gradNorm;
     }
 
     async evaluate(input: string | number[]): Promise<number> {
@@ -300,10 +361,13 @@ function buildChunks(ids: number[], seqLen: number): Array<{inputs: number[], ta
 
 function packAdamParams(
     numElements: number, lr: number, beta1: number, beta2: number,
-    eps: number, weightDecay: number, beta1_t: number, beta2_t: number
+    eps: number, weightDecay: number, beta1_t: number, beta2_t: number,
+    maxDelta: number
 ): ArrayBuffer {
-    const buf = new ArrayBuffer(32);
+    // 48 bytes: u32 + 8×f32 = 36, padded to a 16-byte multiple for the uniform
+    // layout (matches AdamParams' _pad0.._pad2 in WEIGHT_UPDATE_WGSL).
+    const buf = new ArrayBuffer(48);
     new Uint32Array(buf, 0, 1).set([numElements]);
-    new Float32Array(buf, 4, 7).set([lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t]);
+    new Float32Array(buf, 4, 8).set([lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t, maxDelta]);
     return buf;
 }