@seanhogg/builderforce-memory-engine 2026.6.35 → 2026.7.11

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

package/src/utils/quantization.ts

@@ -54,7 +54,10 @@ export function fp16ToFloat(val: number): number {
     const mantissa =  val         & 0x3FF;
 
     if (exponent === 0) {
-        const f = mantissa / 1024.0;
+        // Subnormal: value = mantissa/1024 × 2^-14 (no implicit leading 1, min
+        // exponent -14). The 2^-14 scale was missing, so subnormals (|x| < ~6.1e-5)
+        // decoded ~16384× too large — corrupting every fp16 round-trip of tiny weights.
+        const f = (mantissa / 1024.0) * Math.pow(2, -14);
         return sign ? -f : f;
     }