@seanhogg/builderforce-memory-engine 2026.7.14 → 2026.9.28

package/dist/training/model_cpu.d.ts

@@ -0,0 +1,100 @@
+/**
+ * model_cpu.ts — the exact CPU forward + backward for a whole
+ * {@link HybridMambaModel}: token embedding → layer stack → final RMSNorm →
+ * tied LM head → cross-entropy.
+ *
+ * This is the gradient engine {@link MambaTrainer} uses. Before it existed the
+ * trainer handed the SAME `dLogits` buffer to AdamW for every parameter, so no
+ * parameter ever received its own gradient — `learn()` was safe but inert. Now
+ * the backward pass is real, and `tests/backprop.test.ts` proves it against
+ * finite differences without needing a GPU.
+ *
+ * Parameter names match {@link HybridMambaModel.parameters} exactly
+ * (`embedding`, `layer{i}.{name}`, `final_norm`), which is how the trainer
+ * pairs each parameter with its OWN gradient.
+ */
+import { type Mamba1CpuWeights, type Mamba1CpuGrads } from './mamba1_cpu.js';
+/** Layer types the CPU gradient engine can differentiate today. */
+export declare const CPU_GRADIENT_LAYER_TYPES: readonly ["mamba1"];
+export type CpuGradientLayerType = (typeof CPU_GRADIENT_LAYER_TYPES)[number];
+export interface CpuModelDims {
+    vocabSize: number;
+    dModel: number;
+    dState: number;
+    dConv: number;
+    dInner: number;
+    dtRank: number;
+    numLayers: number;
+}
+/** Every tensor the CPU model reads, mirrored out of the GPU buffers. */
+export interface CpuModelWeights {
+    embedding: Float32Array;
+    finalNorm: Float32Array;
+    lmHeadBias: Float32Array;
+    layers: Mamba1CpuWeights[];
+}
+export interface CpuModelGrads {
+    embedding: Float32Array;
+    finalNorm: Float32Array;
+    lmHeadBias: Float32Array;
+    layers: Mamba1CpuGrads[];
+}
+/** Knobs on one backward pass. */
+export interface CpuBackwardOptions {
+    /**
+     * Recompute each layer's activations during the backward pass instead of
+     * retaining every layer's cache from the forward. Trades one extra forward per
+     * layer for O(1) instead of O(layers) retained activations; the gradients are
+     * numerically identical (`tests/peft_trainer.test.ts` asserts bit equality).
+     */
+    activationCheckpointing?: boolean;
+}
+/** Loss + the gradient of every parameter, keyed by `parameters()` name. */
+export interface CpuBackwardResult {
+    loss: number;
+    grads: Map<string, Float32Array>;
+}
+export declare function zeroCpuModelGrads(dims: CpuModelDims): CpuModelGrads;
+/** Flatten a {@link CpuModelGrads} into the `parameters()` naming the optimizer uses. */
+export declare function toNamedGrads(g: CpuModelGrads): Map<string, Float32Array>;
+/** The model's tensors keyed the same way — the inverse of {@link cpuWeightsFromNamed}. */
+export declare function namedWeights(w: CpuModelWeights): Map<string, Float32Array>;
+/**
+ * Assemble a {@link CpuModelWeights} from tensors keyed by `parameters()` name —
+ * how {@link MambaTrainer} mirrors the GPU buffers into the gradient engine.
+ * Throws on a missing or wrongly-sized tensor rather than training on garbage.
+ */
+export declare function cpuWeightsFromNamed(named: Map<string, Float32Array>, dims: CpuModelDims): CpuModelWeights;
+/** Logits only — the CPU mirror of `HybridMambaModel.forward`. */
+export declare function cpuModelForward(tokenIds: ArrayLike<number>, w: CpuModelWeights, dims: CpuModelDims, batch: number, seqLen: number): Float32Array;
+/** Mean next-token cross-entropy over the sequence — the trainer's loss. */
+export declare function cpuModelLoss(tokenIds: ArrayLike<number>, targets: ArrayLike<number>, w: CpuModelWeights, dims: CpuModelDims, batch: number, seqLen: number): number;
+/**
+ * Forward + backward. Returns the mean cross-entropy loss and the gradient of
+ * EVERY parameter, keyed by its `parameters()` name.
+ *
+ * The loss is averaged over the `batch * seqLen` predicted positions, so the
+ * gradients are the gradients of that mean — the quantity the trainer reports.
+ */
+export declare function cpuModelBackward(tokenIds: ArrayLike<number>, targets: ArrayLike<number>, w: CpuModelWeights, dims: CpuModelDims, batch: number, seqLen: number, opts?: CpuBackwardOptions): CpuBackwardResult;
+/**
+ * The 2-D shape of a parameter, or `null` when it is a vector (bias / norm gain
+ * / `D_vec`).
+ *
+ * Low-rank adaptation only means anything over a matrix, so this is the single
+ * place that says which parameters an adapter can cover — and it derives the
+ * shapes from `dims` rather than restating the block's layout, so it cannot
+ * drift from {@link cpuWeightsFromNamed}.
+ */
+export declare function paramMatrixShape(name: string, dims: CpuModelDims): [rows: number, cols: number] | null;
+/**
+ * Parameters a low-rank adapter covers by default: the four projections inside
+ * each block.
+ *
+ * Deliberately NOT the token embedding (by far the largest tensor, and tied to
+ * the output head — adapting it low-rank changes the head too) and NOT `A_log`
+ * (the state-transition decay, which WSLA freezes precisely because drifting it
+ * destabilises the SSM).
+ */
+export declare const DEFAULT_LORA_TARGETS: readonly ["wInProj", "wXProj", "wDtProj", "wOutProj"];
+//# sourceMappingURL=model_cpu.d.ts.map

package/dist/training/model_cpu.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"model_cpu.d.ts","sourceRoot":"","sources":["../../src/training/model_cpu.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;GAcG;AAQH,OAAO,EAKH,KAAK,gBAAgB,EAGrB,KAAK,cAAc,EACtB,MAAM,iBAAiB,CAAC;AAGzB,mEAAmE;AACnE,eAAO,MAAM,wBAAwB,qBAAsB,CAAC;AAC5D,MAAM,MAAM,oBAAoB,GAAG,CAAC,OAAO,wBAAwB,CAAC,CAAC,MAAM,CAAC,CAAC;AAE7E,MAAM,WAAW,YAAY;IACzB,SAAS,EAAE,MAAM,CAAC;IAClB,MAAM,EAAK,MAAM,CAAC;IAClB,MAAM,EAAK,MAAM,CAAC;IAClB,KAAK,EAAM,MAAM,CAAC;IAClB,MAAM,EAAK,MAAM,CAAC;IAClB,MAAM,EAAK,MAAM,CAAC;IAClB,SAAS,EAAE,MAAM,CAAC;CACrB;AAED,yEAAyE;AACzE,MAAM,WAAW,eAAe;IAC5B,SAAS,EAAI,YAAY,CAAC;IAC1B,SAAS,EAAI,YAAY,CAAC;IAC1B,UAAU,EAAG,YAAY,CAAC;IAC1B,MAAM,EAAO,gBAAgB,EAAE,CAAC;CACnC;AAED,MAAM,WAAW,aAAa;IAC1B,SAAS,EAAI,YAAY,CAAC;IAC1B,SAAS,EAAI,YAAY,CAAC;IAC1B,UAAU,EAAG,YAAY,CAAC;IAC1B,MAAM,EAAO,cAAc,EAAE,CAAC;CACjC;AAED,kCAAkC;AAClC,MAAM,WAAW,kBAAkB;IAC/B;;;;;OAKG;IACH,uBAAuB,CAAC,EAAE,OAAO,CAAC;CACrC;AAED,4EAA4E;AAC5E,MAAM,WAAW,iBAAiB;IAC9B,IAAI,EAAG,MAAM,CAAC;IACd,KAAK,EAAE,GAAG,CAAC,MAAM,EAAE,YAAY,CAAC,CAAC;CACpC;AAED,wBAAgB,iBAAiB,CAAC,IAAI,EAAE,YAAY,GAAG,aAAa,CAWnE;AAED,yFAAyF;AACzF,wBAAgB,YAAY,CAAC,CAAC,EAAE,aAAa,GAAG,GAAG,CAAC,MAAM,EAAE,YAAY,CAAC,CAUxE;AAED,2FAA2F;AAC3F,wBAAgB,YAAY,CAAC,CAAC,EAAE,eAAe,GAAG,GAAG,CAAC,MAAM,EAAE,YAAY,CAAC,CAU1E;AAED;;;;GAIG;AACH,wBAAgB,mBAAmB,CAAC,KAAK,EAAE,GAAG,CAAC,MAAM,EAAE,YAAY,CAAC,EAAE,IAAI,EAAE,YAAY,GAAG,eAAe,CA2BzG;AA6DD,kEAAkE;AAClE,wBAAgB,eAAe,CAC3B,QAAQ,EAAE,SAAS,CAAC,MAAM,CAAC,EAC3B,CAAC,EAAE,eAAe,EAClB,IAAI,EAAE,YAAY,EAClB,KAAK,EAAE,MAAM,EACb,MAAM,EAAE,MAAM,GACf,YAAY,CAEd;AAED,4EAA4E;AAC5E,wBAAgB,YAAY,CACxB,QAAQ,EAAE,SAAS,CAAC,MAAM,CAAC,EAC3B,OAAO,EAAE,SAAS,CAAC,MAAM,CAAC,EAC1B,CAAC,EAAE,eAAe,EAClB,IAAI,EAAE,YAAY,EAClB,KAAK,EAAE,MAAM,EACb,MAAM,EAAE,MAAM,GACf,MAAM,CASR;AAED;;;;;;GAMG;AACH,wBAAgB,gBAAgB,CAC5B,QAAQ,EAAE,SAAS,CAAC,MAAM,CAAC,EAC3B,OAAO,EAAE,SAAS,CAAC,MAAM,CAAC,EAC1B,CAAC,EAAE,eAAe,EAClB,IAAI,EAAE,YAAY,EAClB,KAAK,EAAE,MAAM,EACb,MAAM,EAAE,MAAM,EACd,IAAI,GAAE,kBAAuB,GAC9B,iBAAiB,CAiEnB;AAID;;;;;;;;GAQG;AACH,wBAAgB,gBAAgB,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,YAAY,GAAG,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,CAAC,GAAG,IAAI,CAatG;AAED;;;;;;;;GAQG;AACH,eAAO,MAAM,oBAAoB,uDAAwD,CAAC"}

package/dist/training/model_cpu.js

@@ -0,0 +1,252 @@
+/**
+ * model_cpu.ts — the exact CPU forward + backward for a whole
+ * {@link HybridMambaModel}: token embedding → layer stack → final RMSNorm →
+ * tied LM head → cross-entropy.
+ *
+ * This is the gradient engine {@link MambaTrainer} uses. Before it existed the
+ * trainer handed the SAME `dLogits` buffer to AdamW for every parameter, so no
+ * parameter ever received its own gradient — `learn()` was safe but inert. Now
+ * the backward pass is real, and `tests/backprop.test.ts` proves it against
+ * finite differences without needing a GPU.
+ *
+ * Parameter names match {@link HybridMambaModel.parameters} exactly
+ * (`embedding`, `layer{i}.{name}`, `final_norm`), which is how the trainer
+ * pairs each parameter with its OWN gradient.
+ */
+import { linearForward, rmsNormForward, rmsNormBackward, RMSNORM_EPS, } from './cpu_ops.js';
+import { mamba1CpuForward, mamba1CpuBackward, zeroMamba1Grads, MAMBA1_PARAM_NAMES, } from './mamba1_cpu.js';
+import { crossEntropyLoss, crossEntropyGrad } from './autograd.js';
+/** Layer types the CPU gradient engine can differentiate today. */
+export const CPU_GRADIENT_LAYER_TYPES = ['mamba1'];
+export function zeroCpuModelGrads(dims) {
+    const blockDims = {
+        dModel: dims.dModel, dState: dims.dState, dConv: dims.dConv,
+        dInner: dims.dInner, dtRank: dims.dtRank, batch: 1, seqLen: 1,
+    };
+    return {
+        embedding: new Float32Array(dims.vocabSize * dims.dModel),
+        finalNorm: new Float32Array(dims.dModel),
+        lmHeadBias: new Float32Array(dims.vocabSize),
+        layers: Array.from({ length: dims.numLayers }, () => zeroMamba1Grads(blockDims)),
+    };
+}
+/** Flatten a {@link CpuModelGrads} into the `parameters()` naming the optimizer uses. */
+export function toNamedGrads(g) {
+    const out = new Map();
+    out.set('embedding', g.embedding);
+    for (let i = 0; i < g.layers.length; i++) {
+        const lg = g.layers[i];
+        for (const name of MAMBA1_PARAM_NAMES)
+            out.set(`layer${i}.${name}`, lg[name]);
+    }
+    out.set('final_norm', g.finalNorm);
+    out.set('lm_head_bias', g.lmHeadBias);
+    return out;
+}
+/** The model's tensors keyed the same way — the inverse of {@link cpuWeightsFromNamed}. */
+export function namedWeights(w) {
+    const out = new Map();
+    out.set('embedding', w.embedding);
+    for (let i = 0; i < w.layers.length; i++) {
+        const lw = w.layers[i];
+        for (const name of MAMBA1_PARAM_NAMES)
+            out.set(`layer${i}.${name}`, lw[name]);
+    }
+    out.set('final_norm', w.finalNorm);
+    out.set('lm_head_bias', w.lmHeadBias);
+    return out;
+}
+/**
+ * Assemble a {@link CpuModelWeights} from tensors keyed by `parameters()` name —
+ * how {@link MambaTrainer} mirrors the GPU buffers into the gradient engine.
+ * Throws on a missing or wrongly-sized tensor rather than training on garbage.
+ */
+export function cpuWeightsFromNamed(named, dims) {
+    const take = (key, expected) => {
+        const t = named.get(key);
+        if (!t)
+            throw new Error(`cpuWeightsFromNamed: missing tensor "${key}"`);
+        if (t.length !== expected) {
+            throw new Error(`cpuWeightsFromNamed: "${key}" has ${t.length} elements, expected ${expected}`);
+        }
+        return t;
+    };
+    const { dModel, dState: N, dConv: K, dInner: D, dtRank: R } = dims;
+    const sizes = {
+        wInProj: 2 * D * dModel, bInProj: 2 * D, wConv: D * K, bConv: D,
+        wXProj: (R + 2 * N) * D, bXProj: R + 2 * N, wDtProj: D * R, bDtProj: D,
+        A_log: D * N, D_vec: D, wOutProj: dModel * D, bOutProj: dModel, normWeight: dModel,
+    };
+    const layers = [];
+    for (let i = 0; i < dims.numLayers; i++) {
+        const lw = {};
+        for (const name of MAMBA1_PARAM_NAMES)
+            lw[name] = take(`layer${i}.${name}`, sizes[name]);
+        layers.push(lw);
+    }
+    return {
+        embedding: take('embedding', dims.vocabSize * dModel),
+        finalNorm: take('final_norm', dModel),
+        lmHeadBias: take('lm_head_bias', dims.vocabSize),
+        layers,
+    };
+}
+function forwardInternal(tokenIds, w, dims, batch, seqLen, 
+/** Drop each layer's activation cache after use — see {@link CpuBackwardOptions}. */
+checkpointed = false) {
+    const { vocabSize, dModel } = dims;
+    const M = batch * seqLen;
+    // Token embedding lookup.
+    let hidden = new Float32Array(M * dModel);
+    for (let r = 0; r < M; r++) {
+        const id = tokenIds[r];
+        if (id < 0 || id >= vocabSize)
+            throw new Error(`token id ${id} out of range [0, ${vocabSize})`);
+        hidden.set(w.embedding.subarray(id * dModel, (id + 1) * dModel), r * dModel);
+    }
+    const blockDims = {
+        dModel, dState: dims.dState, dConv: dims.dConv,
+        dInner: dims.dInner, dtRank: dims.dtRank, batch, seqLen,
+    };
+    const hiddenIn = [];
+    const layerCaches = [];
+    for (let i = 0; i < dims.numLayers; i++) {
+        hiddenIn.push(hidden);
+        const { output, cache } = mamba1CpuForward(hidden, w.layers[i], blockDims);
+        // Under activation checkpointing only the layer INPUTS are retained; each
+        // layer's cache is recomputed from its input when its backward runs. Memory
+        // drops from O(layers) caches to one, at the cost of a second forward per
+        // layer — gradients are numerically identical either way.
+        if (!checkpointed)
+            layerCaches.push(cache);
+        hidden = output;
+    }
+    // Final RMSNorm.
+    const normOut = new Float32Array(M * dModel);
+    const normInv = new Float32Array(M);
+    rmsNormForward(hidden, w.finalNorm, M, dModel, normOut, normInv, RMSNORM_EPS);
+    // Tied LM head: logits = normOut @ embedding^T + lmHeadBias.
+    const logits = new Float32Array(M * vocabSize);
+    linearForward(normOut, w.embedding, w.lmHeadBias, M, dModel, vocabSize, logits);
+    return { hiddenIn, layer: layerCaches, lastHidden: hidden, normOut, normInv, logits };
+}
+/** Logits only — the CPU mirror of `HybridMambaModel.forward`. */
+export function cpuModelForward(tokenIds, w, dims, batch, seqLen) {
+    return forwardInternal(tokenIds, w, dims, batch, seqLen).logits;
+}
+/** Mean next-token cross-entropy over the sequence — the trainer's loss. */
+export function cpuModelLoss(tokenIds, targets, w, dims, batch, seqLen) {
+    const logits = cpuModelForward(tokenIds, w, dims, batch, seqLen);
+    const M = batch * seqLen;
+    const V = dims.vocabSize;
+    let total = 0;
+    for (let r = 0; r < M; r++) {
+        total += crossEntropyLoss(logits.subarray(r * V, (r + 1) * V), targets[r]);
+    }
+    return total / M;
+}
+/**
+ * Forward + backward. Returns the mean cross-entropy loss and the gradient of
+ * EVERY parameter, keyed by its `parameters()` name.
+ *
+ * The loss is averaged over the `batch * seqLen` predicted positions, so the
+ * gradients are the gradients of that mean — the quantity the trainer reports.
+ */
+export function cpuModelBackward(tokenIds, targets, w, dims, batch, seqLen, opts = {}) {
+    const { vocabSize: V, dModel } = dims;
+    const M = batch * seqLen;
+    const checkpointed = opts.activationCheckpointing ?? false;
+    const cache = forwardInternal(tokenIds, w, dims, batch, seqLen, checkpointed);
+    const g = zeroCpuModelGrads(dims);
+    // Cross-entropy over every position, averaged.
+    let total = 0;
+    const dLogits = new Float32Array(M * V);
+    for (let r = 0; r < M; r++) {
+        const slice = cache.logits.subarray(r * V, (r + 1) * V);
+        const target = targets[r];
+        total += crossEntropyLoss(slice, target);
+        const gr = crossEntropyGrad(slice, target);
+        for (let v = 0; v < V; v++)
+            dLogits[r * V + v] = gr[v] / M;
+    }
+    const loss = total / M;
+    // LM head (tied): logits = normOut @ embedding^T + bias.
+    // dNormOut = dLogits @ embedding ; dEmbedding += dLogits^T @ normOut.
+    const dNormOut = new Float32Array(M * dModel);
+    for (let r = 0; r < M; r++) {
+        const lo = r * V;
+        const no = r * dModel;
+        for (let v = 0; v < V; v++) {
+            const gl = dLogits[lo + v];
+            if (gl === 0)
+                continue;
+            g.lmHeadBias[v] = g.lmHeadBias[v] + gl;
+            const eo = v * dModel;
+            for (let i = 0; i < dModel; i++) {
+                dNormOut[no + i] = dNormOut[no + i] + gl * w.embedding[eo + i];
+                g.embedding[eo + i] = g.embedding[eo + i] + gl * cache.normOut[no + i];
+            }
+        }
+    }
+    // Final RMSNorm.
+    const dLastHidden = new Float32Array(M * dModel);
+    rmsNormBackward(dNormOut, cache.lastHidden, w.finalNorm, cache.normInv, M, dModel, dLastHidden, g.finalNorm, RMSNORM_EPS);
+    // Layer stack, in reverse.
+    const blockDims = {
+        dModel, dState: dims.dState, dConv: dims.dConv,
+        dInner: dims.dInner, dtRank: dims.dtRank, batch, seqLen,
+    };
+    let d = dLastHidden;
+    for (let i = dims.numLayers - 1; i >= 0; i--) {
+        const layerCache = checkpointed
+            ? mamba1CpuForward(cache.hiddenIn[i], w.layers[i], blockDims).cache
+            : cache.layer[i];
+        d = mamba1CpuBackward(d, w.layers[i], layerCache, blockDims, g.layers[i]);
+    }
+    // Token embedding — the SAME table the LM head reads, so this accumulates
+    // on top of the head's contribution (the tie is why both terms are needed).
+    for (let r = 0; r < M; r++) {
+        const id = tokenIds[r];
+        const eo = id * dModel;
+        const ho = r * dModel;
+        for (let i = 0; i < dModel; i++)
+            g.embedding[eo + i] = g.embedding[eo + i] + d[ho + i];
+    }
+    return { loss, grads: toNamedGrads(g) };
+}
+// ── Parameter shapes (for PEFT) ──────────────────────────────────────────────
+/**
+ * The 2-D shape of a parameter, or `null` when it is a vector (bias / norm gain
+ * / `D_vec`).
+ *
+ * Low-rank adaptation only means anything over a matrix, so this is the single
+ * place that says which parameters an adapter can cover — and it derives the
+ * shapes from `dims` rather than restating the block's layout, so it cannot
+ * drift from {@link cpuWeightsFromNamed}.
+ */
+export function paramMatrixShape(name, dims) {
+    const { dModel, dState: N, dConv: K, dInner: D, dtRank: R } = dims;
+    if (name === 'embedding')
+        return [dims.vocabSize, dModel];
+    const local = name.startsWith('layer') ? name.slice(name.indexOf('.') + 1) : name;
+    switch (local) {
+        case 'wInProj': return [2 * D, dModel];
+        case 'wConv': return [D, K];
+        case 'wXProj': return [R + 2 * N, D];
+        case 'wDtProj': return [D, R];
+        case 'A_log': return [D, N];
+        case 'wOutProj': return [dModel, D];
+        default: return null; // bInProj, bConv, bXProj, bDtProj, D_vec, bOutProj, normWeight, final_norm, lm_head_bias
+    }
+}
+/**
+ * Parameters a low-rank adapter covers by default: the four projections inside
+ * each block.
+ *
+ * Deliberately NOT the token embedding (by far the largest tensor, and tied to
+ * the output head — adapting it low-rank changes the head too) and NOT `A_log`
+ * (the state-transition decay, which WSLA freezes precisely because drifting it
+ * destabilises the SSM).
+ */
+export const DEFAULT_LORA_TARGETS = ['wInProj', 'wXProj', 'wDtProj', 'wOutProj'];
+//# sourceMappingURL=model_cpu.js.map

package/dist/training/model_cpu.js.map

@@ -0,0 +1 @@
+{"version":3,"file":"model_cpu.js","sourceRoot":"","sources":["../../src/training/model_cpu.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;GAcG;AAEH,OAAO,EACH,aAAa,EACb,cAAc,EACd,eAAe,EACf,WAAW,GACd,MAAM,cAAc,CAAC;AACtB,OAAO,EACH,gBAAgB,EAChB,iBAAiB,EACjB,eAAe,EACf,kBAAkB,GAKrB,MAAM,iBAAiB,CAAC;AACzB,OAAO,EAAE,gBAAgB,EAAE,gBAAgB,EAAE,MAAM,eAAe,CAAC;AAEnE,mEAAmE;AACnE,MAAM,CAAC,MAAM,wBAAwB,GAAG,CAAC,QAAQ,CAAU,CAAC;AA6C5D,MAAM,UAAU,iBAAiB,CAAC,IAAkB;IAChD,MAAM,SAAS,GAAkB;QAC7B,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,KAAK,EAAE,IAAI,CAAC,KAAK;QAC3D,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,KAAK,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC;KAChE,CAAC;IACF,OAAO;QACH,SAAS,EAAG,IAAI,YAAY,CAAC,IAAI,CAAC,SAAS,GAAG,IAAI,CAAC,MAAM,CAAC;QAC1D,SAAS,EAAG,IAAI,YAAY,CAAC,IAAI,CAAC,MAAM,CAAC;QACzC,UAAU,EAAE,IAAI,YAAY,CAAC,IAAI,CAAC,SAAS,CAAC;QAC5C,MAAM,EAAM,KAAK,CAAC,IAAI,CAAC,EAAE,MAAM,EAAE,IAAI,CAAC,SAAS,EAAE,EAAE,GAAG,EAAE,CAAC,eAAe,CAAC,SAAS,CAAC,CAAC;KACvF,CAAC;AACN,CAAC;AAED,yFAAyF;AACzF,MAAM,UAAU,YAAY,CAAC,CAAgB;IACzC,MAAM,GAAG,GAAG,IAAI,GAAG,EAAwB,CAAC;IAC5C,GAAG,CAAC,GAAG,CAAC,WAAW,EAAE,CAAC,CAAC,SAAS,CAAC,CAAC;IAClC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QACvC,MAAM,EAAE,GAAG,CAAC,CAAC,MAAM,CAAC,CAAC,CAAE,CAAC;QACxB,KAAK,MAAM,IAAI,IAAI,kBAAkB;YAAE,GAAG,CAAC,GAAG,CAAC,QAAQ,CAAC,IAAI,IAAI,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC;IAClF,CAAC;IACD,GAAG,CAAC,GAAG,CAAC,YAAY,EAAE,CAAC,CAAC,SAAS,CAAC,CAAC;IACnC,GAAG,CAAC,GAAG,CAAC,cAAc,EAAE,CAAC,CAAC,UAAU,CAAC,CAAC;IACtC,OAAO,GAAG,CAAC;AACf,CAAC;AAED,2FAA2F;AAC3F,MAAM,UAAU,YAAY,CAAC,CAAkB;IAC3C,MAAM,GAAG,GAAG,IAAI,GAAG,EAAwB,CAAC;IAC5C,GAAG,CAAC,GAAG,CAAC,WAAW,EAAE,CAAC,CAAC,SAAS,CAAC,CAAC;IAClC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,CAAC,MAAM,CAAC,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;QACvC,MAAM,EAAE,GAAG,CAAC,CAAC,MAAM,CAAC,CAAC,CAAE,CAAC;QACxB,KAAK,MAAM,IAAI,IAAI,kBAAkB;YAAE,GAAG,CAAC,GAAG,CAAC,QAAQ,CAAC,IAAI,IAAI,EAAE,EAAE,EAAE,CAAC,IAAI,CAAC,CAAC,CAAC;IAClF,CAAC;IACD,GAAG,CAAC,GAAG,CAAC,YAAY,EAAE,CAAC,CAAC,SAAS,CAAC,CAAC;IACnC,GAAG,CAAC,GAAG,CAAC,cAAc,EAAE,CAAC,CAAC,UAAU,CAAC,CAAC;IACtC,OAAO,GAAG,CAAC;AACf,CAAC;AAED;;;;GAIG;AACH,MAAM,UAAU,mBAAmB,CAAC,KAAgC,EAAE,IAAkB;IACpF,MAAM,IAAI,GAAG,CAAC,GAAW,EAAE,QAAgB,EAAgB,EAAE;QACzD,MAAM,CAAC,GAAG,KAAK,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC;QACzB,IAAI,CAAC,CAAC;YAAE,MAAM,IAAI,KAAK,CAAC,wCAAwC,GAAG,GAAG,CAAC,CAAC;QACxE,IAAI,CAAC,CAAC,MAAM,KAAK,QAAQ,EAAE,CAAC;YACxB,MAAM,IAAI,KAAK,CAAC,yBAAyB,GAAG,SAAS,CAAC,CAAC,MAAM,uBAAuB,QAAQ,EAAE,CAAC,CAAC;QACpG,CAAC;QACD,OAAO,CAAC,CAAC;IACb,CAAC,CAAC;IACF,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,GAAG,IAAI,CAAC;IACnE,MAAM,KAAK,GAA2C;QAClD,OAAO,EAAE,CAAC,GAAG,CAAC,GAAG,MAAM,EAAE,OAAO,EAAE,CAAC,GAAG,CAAC,EAAE,KAAK,EAAE,CAAC,GAAG,CAAC,EAAE,KAAK,EAAE,CAAC;QAC/D,MAAM,EAAE,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,EAAE,MAAM,EAAE,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,OAAO,EAAE,CAAC,GAAG,CAAC,EAAE,OAAO,EAAE,CAAC;QACtE,KAAK,EAAE,CAAC,GAAG,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,QAAQ,EAAE,MAAM,GAAG,CAAC,EAAE,QAAQ,EAAE,MAAM,EAAE,UAAU,EAAE,MAAM;KACrF,CAAC;IACF,MAAM,MAAM,GAAuB,EAAE,CAAC;IACtC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,SAAS,EAAE,CAAC,EAAE,EAAE,CAAC;QACtC,MAAM,EAAE,GAAG,EAAsB,CAAC;QAClC,KAAK,MAAM,IAAI,IAAI,kBAAkB;YAAE,EAAE,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC,QAAQ,CAAC,IAAI,IAAI,EAAE,EAAE,KAAK,CAAC,IAAI,CAAC,CAAC,CAAC;QACzF,MAAM,CAAC,IAAI,CAAC,EAAE,CAAC,CAAC;IACpB,CAAC;IACD,OAAO;QACH,SAAS,EAAG,IAAI,CAAC,WAAW,EAAE,IAAI,CAAC,SAAS,GAAG,MAAM,CAAC;QACtD,SAAS,EAAG,IAAI,CAAC,YAAY,EAAE,MAAM,CAAC;QACtC,UAAU,EAAE,IAAI,CAAC,cAAc,EAAE,IAAI,CAAC,SAAS,CAAC;QAChD,MAAM;KACT,CAAC;AACN,CAAC;AAWD,SAAS,eAAe,CACpB,QAA2B,EAC3B,CAAkB,EAClB,IAAkB,EAClB,KAAa,EACb,MAAc;AACd,qFAAqF;AACrF,YAAY,GAAG,KAAK;IAEpB,MAAM,EAAE,SAAS,EAAE,MAAM,EAAE,GAAG,IAAI,CAAC;IACnC,MAAM,CAAC,GAAG,KAAK,GAAG,MAAM,CAAC;IAEzB,0BAA0B;IAC1B,IAAI,MAAM,GAAiB,IAAI,YAAY,CAAC,CAAC,GAAG,MAAM,CAAC,CAAC;IACxD,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QACzB,MAAM,EAAE,GAAG,QAAQ,CAAC,CAAC,CAAE,CAAC;QACxB,IAAI,EAAE,GAAG,CAAC,IAAI,EAAE,IAAI,SAAS;YAAE,MAAM,IAAI,KAAK,CAAC,YAAY,EAAE,qBAAqB,SAAS,GAAG,CAAC,CAAC;QAChG,MAAM,CAAC,GAAG,CAAC,CAAC,CAAC,SAAS,CAAC,QAAQ,CAAC,EAAE,GAAG,MAAM,EAAE,CAAC,EAAE,GAAG,CAAC,CAAC,GAAG,MAAM,CAAC,EAAE,CAAC,GAAG,MAAM,CAAC,CAAC;IACjF,CAAC;IAED,MAAM,SAAS,GAAkB;QAC7B,MAAM,EAAE,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,KAAK,EAAE,IAAI,CAAC,KAAK;QAC9C,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,KAAK,EAAE,MAAM;KAC1D,CAAC;IAEF,MAAM,QAAQ,GAAmB,EAAE,CAAC;IACpC,MAAM,WAAW,GAAqB,EAAE,CAAC;IACzC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,IAAI,CAAC,SAAS,EAAE,CAAC,EAAE,EAAE,CAAC;QACtC,QAAQ,CAAC,IAAI,CAAC,MAAM,CAAC,CAAC;QACtB,MAAM,EAAE,MAAM,EAAE,KAAK,EAAE,GAAG,gBAAgB,CAAC,MAAM,EAAE,CAAC,CAAC,MAAM,CAAC,CAAC,CAAE,EAAE,SAAS,CAAC,CAAC;QAC5E,0EAA0E;QAC1E,4EAA4E;QAC5E,0EAA0E;QAC1E,0DAA0D;QAC1D,IAAI,CAAC,YAAY;YAAE,WAAW,CAAC,IAAI,CAAC,KAAK,CAAC,CAAC;QAC3C,MAAM,GAAG,MAAM,CAAC;IACpB,CAAC;IAED,iBAAiB;IACjB,MAAM,OAAO,GAAG,IAAI,YAAY,CAAC,CAAC,GAAG,MAAM,CAAC,CAAC;IAC7C,MAAM,OAAO,GAAG,IAAI,YAAY,CAAC,CAAC,CAAC,CAAC;IACpC,cAAc,CAAC,MAAM,EAAE,CAAC,CAAC,SAAS,EAAE,CAAC,EAAE,MAAM,EAAE,OAAO,EAAE,OAAO,EAAE,WAAW,CAAC,CAAC;IAE9E,6DAA6D;IAC7D,MAAM,MAAM,GAAG,IAAI,YAAY,CAAC,CAAC,GAAG,SAAS,CAAC,CAAC;IAC/C,aAAa,CAAC,OAAO,EAAE,CAAC,CAAC,SAAS,EAAE,CAAC,CAAC,UAAU,EAAE,CAAC,EAAE,MAAM,EAAE,SAAS,EAAE,MAAM,CAAC,CAAC;IAEhF,OAAO,EAAE,QAAQ,EAAE,KAAK,EAAE,WAAW,EAAE,UAAU,EAAE,MAAM,EAAE,OAAO,EAAE,OAAO,EAAE,MAAM,EAAE,CAAC;AAC1F,CAAC;AAED,kEAAkE;AAClE,MAAM,UAAU,eAAe,CAC3B,QAA2B,EAC3B,CAAkB,EAClB,IAAkB,EAClB,KAAa,EACb,MAAc;IAEd,OAAO,eAAe,CAAC,QAAQ,EAAE,CAAC,EAAE,IAAI,EAAE,KAAK,EAAE,MAAM,CAAC,CAAC,MAAM,CAAC;AACpE,CAAC;AAED,4EAA4E;AAC5E,MAAM,UAAU,YAAY,CACxB,QAA2B,EAC3B,OAA0B,EAC1B,CAAkB,EAClB,IAAkB,EAClB,KAAa,EACb,MAAc;IAEd,MAAM,MAAM,GAAG,eAAe,CAAC,QAAQ,EAAE,CAAC,EAAE,IAAI,EAAE,KAAK,EAAE,MAAM,CAAC,CAAC;IACjE,MAAM,CAAC,GAAG,KAAK,GAAG,MAAM,CAAC;IACzB,MAAM,CAAC,GAAG,IAAI,CAAC,SAAS,CAAC;IACzB,IAAI,KAAK,GAAG,CAAC,CAAC;IACd,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QACzB,KAAK,IAAI,gBAAgB,CAAC,MAAM,CAAC,QAAQ,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,EAAE,OAAO,CAAC,CAAC,CAAE,CAAC,CAAC;IAChF,CAAC;IACD,OAAO,KAAK,GAAG,CAAC,CAAC;AACrB,CAAC;AAED;;;;;;GAMG;AACH,MAAM,UAAU,gBAAgB,CAC5B,QAA2B,EAC3B,OAA0B,EAC1B,CAAkB,EAClB,IAAkB,EAClB,KAAa,EACb,MAAc,EACd,OAA2B,EAAE;IAE7B,MAAM,EAAE,SAAS,EAAE,CAAC,EAAE,MAAM,EAAE,GAAG,IAAI,CAAC;IACtC,MAAM,CAAC,GAAG,KAAK,GAAG,MAAM,CAAC;IACzB,MAAM,YAAY,GAAG,IAAI,CAAC,uBAAuB,IAAI,KAAK,CAAC;IAC3D,MAAM,KAAK,GAAG,eAAe,CAAC,QAAQ,EAAE,CAAC,EAAE,IAAI,EAAE,KAAK,EAAE,MAAM,EAAE,YAAY,CAAC,CAAC;IAC9E,MAAM,CAAC,GAAG,iBAAiB,CAAC,IAAI,CAAC,CAAC;IAElC,+CAA+C;IAC/C,IAAI,KAAK,GAAG,CAAC,CAAC;IACd,MAAM,OAAO,GAAG,IAAI,YAAY,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC;IACxC,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QACzB,MAAM,KAAK,GAAG,KAAK,CAAC,MAAM,CAAC,QAAQ,CAAC,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,CAAC;QACxD,MAAM,MAAM,GAAG,OAAO,CAAC,CAAC,CAAE,CAAC;QAC3B,KAAK,IAAI,gBAAgB,CAAC,KAAK,EAAE,MAAM,CAAC,CAAC;QACzC,MAAM,EAAE,GAAG,gBAAgB,CAAC,KAAK,EAAE,MAAM,CAAC,CAAC;QAC3C,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE;YAAE,OAAO,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,GAAG,EAAE,CAAC,CAAC,CAAE,GAAG,CAAC,CAAC;IAChE,CAAC;IACD,MAAM,IAAI,GAAG,KAAK,GAAG,CAAC,CAAC;IAEvB,yDAAyD;IACzD,sEAAsE;IACtE,MAAM,QAAQ,GAAG,IAAI,YAAY,CAAC,CAAC,GAAG,MAAM,CAAC,CAAC;IAC9C,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QACzB,MAAM,EAAE,GAAG,CAAC,GAAG,CAAC,CAAC;QACjB,MAAM,EAAE,GAAG,CAAC,GAAG,MAAM,CAAC;QACtB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;YACzB,MAAM,EAAE,GAAG,OAAO,CAAC,EAAE,GAAG,CAAC,CAAE,CAAC;YAC5B,IAAI,EAAE,KAAK,CAAC;gBAAE,SAAS;YACvB,CAAC,CAAC,UAAU,CAAC,CAAC,CAAC,GAAG,CAAC,CAAC,UAAU,CAAC,CAAC,CAAE,GAAG,EAAE,CAAC;YACxC,MAAM,EAAE,GAAG,CAAC,GAAG,MAAM,CAAC;YACtB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,MAAM,EAAE,CAAC,EAAE,EAAE,CAAC;gBAC9B,QAAQ,CAAC,EAAE,GAAG,CAAC,CAAC,GAAG,QAAQ,CAAC,EAAE,GAAG,CAAC,CAAE,GAAG,EAAE,GAAG,CAAC,CAAC,SAAS,CAAC,EAAE,GAAG,CAAC,CAAE,CAAC;gBACjE,CAAC,CAAC,SAAS,CAAC,EAAE,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,SAAS,CAAC,EAAE,GAAG,CAAC,CAAE,GAAG,EAAE,GAAG,KAAK,CAAC,OAAO,CAAC,EAAE,GAAG,CAAC,CAAE,CAAC;YAC7E,CAAC;QACL,CAAC;IACL,CAAC;IAED,iBAAiB;IACjB,MAAM,WAAW,GAAG,IAAI,YAAY,CAAC,CAAC,GAAG,MAAM,CAAC,CAAC;IACjD,eAAe,CAAC,QAAQ,EAAE,KAAK,CAAC,UAAU,EAAE,CAAC,CAAC,SAAS,EAAE,KAAK,CAAC,OAAO,EAAE,CAAC,EAAE,MAAM,EAC7E,WAAW,EAAE,CAAC,CAAC,SAAS,EAAE,WAAW,CAAC,CAAC;IAE3C,2BAA2B;IAC3B,MAAM,SAAS,GAAkB;QAC7B,MAAM,EAAE,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,KAAK,EAAE,IAAI,CAAC,KAAK;QAC9C,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,MAAM,EAAE,IAAI,CAAC,MAAM,EAAE,KAAK,EAAE,MAAM;KAC1D,CAAC;IACF,IAAI,CAAC,GAAiB,WAAW,CAAC;IAClC,KAAK,IAAI,CAAC,GAAG,IAAI,CAAC,SAAS,GAAG,CAAC,EAAE,CAAC,IAAI,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QAC3C,MAAM,UAAU,GAAG,YAAY;YAC3B,CAAC,CAAC,gBAAgB,CAAC,KAAK,CAAC,QAAQ,CAAC,CAAC,CAAE,EAAE,CAAC,CAAC,MAAM,CAAC,CAAC,CAAE,EAAE,SAAS,CAAC,CAAC,KAAK;YACrE,CAAC,CAAC,KAAK,CAAC,KAAK,CAAC,CAAC,CAAE,CAAC;QACtB,CAAC,GAAG,iBAAiB,CAAC,CAAC,EAAE,CAAC,CAAC,MAAM,CAAC,CAAC,CAAE,EAAE,UAAU,EAAE,SAAS,EAAE,CAAC,CAAC,MAAM,CAAC,CAAC,CAAE,CAAC,CAAC;IAChF,CAAC;IAED,0EAA0E;IAC1E,4EAA4E;IAC5E,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,CAAC,EAAE,CAAC,EAAE,EAAE,CAAC;QACzB,MAAM,EAAE,GAAG,QAAQ,CAAC,CAAC,CAAE,CAAC;QACxB,MAAM,EAAE,GAAG,EAAE,GAAG,MAAM,CAAC;QACvB,MAAM,EAAE,GAAG,CAAC,GAAG,MAAM,CAAC;QACtB,KAAK,IAAI,CAAC,GAAG,CAAC,EAAE,CAAC,GAAG,MAAM,EAAE,CAAC,EAAE;YAAE,CAAC,CAAC,SAAS,CAAC,EAAE,GAAG,CAAC,CAAC,GAAG,CAAC,CAAC,SAAS,CAAC,EAAE,GAAG,CAAC,CAAE,GAAG,CAAC,CAAC,EAAE,GAAG,CAAC,CAAE,CAAC;IAC7F,CAAC;IAED,OAAO,EAAE,IAAI,EAAE,KAAK,EAAE,YAAY,CAAC,CAAC,CAAC,EAAE,CAAC;AAC5C,CAAC;AAED,gFAAgF;AAEhF;;;;;;;;GAQG;AACH,MAAM,UAAU,gBAAgB,CAAC,IAAY,EAAE,IAAkB;IAC7D,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,CAAC,EAAE,KAAK,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,MAAM,EAAE,CAAC,EAAE,GAAG,IAAI,CAAC;IACnE,IAAI,IAAI,KAAK,WAAW;QAAE,OAAO,CAAC,IAAI,CAAC,SAAS,EAAE,MAAM,CAAC,CAAC;IAC1D,MAAM,KAAK,GAAG,IAAI,CAAC,UAAU,CAAC,OAAO,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC,KAAK,CAAC,IAAI,CAAC,OAAO,CAAC,GAAG,CAAC,GAAG,CAAC,CAAC,CAAC,CAAC,CAAC,IAAI,CAAC;IAClF,QAAQ,KAAK,EAAE,CAAC;QACZ,KAAK,SAAS,CAAC,CAAE,OAAO,CAAC,CAAC,GAAG,CAAC,EAAE,MAAM,CAAC,CAAC;QACxC,KAAK,OAAO,CAAC,CAAI,OAAO,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;QAC/B,KAAK,QAAQ,CAAC,CAAG,OAAO,CAAC,CAAC,GAAG,CAAC,GAAG,CAAC,EAAE,CAAC,CAAC,CAAC;QACvC,KAAK,SAAS,CAAC,CAAE,OAAO,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;QAC/B,KAAK,OAAO,CAAC,CAAI,OAAO,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;QAC/B,KAAK,UAAU,CAAC,CAAC,OAAO,CAAC,MAAM,EAAE,CAAC,CAAC,CAAC;QACpC,OAAO,CAAC,CAAS,OAAO,IAAI,CAAC,CAAG,yFAAyF;IAC7H,CAAC;AACL,CAAC;AAED;;;;;;;;GAQG;AACH,MAAM,CAAC,MAAM,oBAAoB,GAAG,CAAC,SAAS,EAAE,QAAQ,EAAE,SAAS,EAAE,UAAU,CAAU,CAAC"}

package/dist/training/trainer.d.ts

@@ -1,11 +1,41 @@
 /**
  * trainer.ts – MambaTrainer class
+ *
+ * REAL backpropagation. Until 2026-08 `_trainStep` handed the SAME `dLogits`
+ * buffer to AdamW for every parameter, so no parameter ever received its own
+ * gradient: `learn()` was SAFE (WSLA freezes the backbone, the trust region
+ * bounds every step, the NaN guard blocks a poisoned weight, and the session
+ * rolls back a regression) but INERT. It now runs a genuine backward pass —
+ * `cpuModelBackward` in `model_cpu.ts` — and each parameter is stepped with its
+ * OWN gradient.
+ *
+ * Where the maths runs, and why:
+ *   • the backward pass is the CPU reference (`cpu_ops.ts` + `mamba1_cpu.ts`),
+ *     a line-for-line mirror of the WGSL forward kernels. It is exact and it is
+ *     PROVEN — `tests/backprop.test.ts` checks every parameter gradient against
+ *     finite differences of the real loss, with no GPU involved;
+ *   • the optimiser is {@link adamwUpdateInPlace}, the CPU mirror of
+ *     `WEIGHT_UPDATE_WGSL`'s `adamw_update`, keeping the trust region and the
+ *     NaN guard the safety story depends on;
+ *   • the model's GPU buffers are written back after every optimiser step, so
+ *     inference (which IS GPU) always serves the trained weights.
+ *
+ * The residual is a WGSL backward pass so training itself runs on the GPU. The
+ * backward kernels exist (`SELECTIVE_SCAN_BACKWARD_WGSL`, `SSD_BACKWARD_WGSL`,
+ * `CONV1D_BACKWARD_WGSL`, `LINEAR_BACKWARD_WGSL`, `ACTIVATIONS_BACKWARD_WGSL`)
+ * but wiring them cannot be validated numerically without a WebGPU device, and
+ * an unvalidated GPU backward is worse than a slower exact one.
  */
+import { type ShardSpec } from '../optim/adamw.js';
+import { LoRAAdapter, type LoRAConfig } from './lora.js';
+import { type CpuModelDims } from './model_cpu.js';
 import { HybridMambaModel, MambaModel } from '../model/mamba_model.js';
 import { BPETokenizer } from '../tokenizer/bpe.js';
 export interface TrainOptions {
     learningRate?: number;
     epochs?: number;
+    /** Sequences trained together per optimiser step. Chunks of equal length are
+     *  grouped into batches of this size; a ragged tail trains on its own. */
     batchSize?: number;
     seqLen?: number;
     maxGradNorm?: number;
@@ -17,11 +47,48 @@ export interface TrainOptions {
     /** 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. */
+    /**
+     * @deprecated The pre-clip gradient L2 norm is now computed exactly on the way
+     * through the optimiser and is ALWAYS reported — it no longer costs a GPU
+     * readback, so there is nothing to opt into. Accepted for API compatibility.
+     */
     trackGradNorm?: boolean;
+    /**
+     * Train LOW-RANK ADAPTERS over the block projections instead of the weights
+     * themselves (see {@link LoRAAdapter}). The base weights are frozen; only A and
+     * B train, so an adapt costs a fraction of the optimiser state and the result is
+     * a small, mergeable delta rather than a whole checkpoint.
+     *
+     * `targets` selects which parameters get an adapter, by their unqualified name
+     * (`wInProj`, `wXProj`, `wDtProj`, `wOutProj`, `wConv`, `A_log`, `embedding`);
+     * omit for {@link DEFAULT_LORA_TARGETS}. A rank larger than a target's smaller
+     * dimension is clamped to it rather than throwing — the small projections
+     * (`wDtProj` is `dInner x dtRank`) would otherwise reject a sensible global rank.
+     */
+    lora?: LoRAConfig & {
+        targets?: readonly string[];
+    };
+    /**
+     * Batches to accumulate gradients over before one optimiser step. Default 1.
+     * The effective batch size is `batchSize * gradientAccumulation` at the memory
+     * cost of `batchSize` — the standard way to train at a batch size the host
+     * cannot hold at once.
+     */
+    gradientAccumulation?: number;
+    /**
+     * Shard the optimiser state (the Adam moments — 2x the parameter bytes, the
+     * dominant cost of a full fine-tune). This instance then OWNS only parameters
+     * where `index % count === index`, and allocates moments for those alone. Run
+     * `count` trainers over the same model and the union of their steps equals one
+     * unsharded step. Same contract as {@link ShardSpec} on {@link AdamW}.
+     */
+    shard?: ShardSpec;
+    /**
+     * Recompute each layer's activations in the backward pass instead of retaining
+     * all of them. Gradients are identical; peak memory drops from O(layers)
+     * activation caches to one, at the cost of a second forward per layer.
+     */
+    activationCheckpointing?: boolean;
     onEpochEnd?: ((epoch: number, loss: number, gradNorm?: number) => void) | null;
 }
 /**
@@ -32,6 +99,18 @@ export interface TrainOptions {
  * non-WSLA mode) for unbounded steps.
  */
 export declare const WSLA_MAX_DELTA = 0.05;
+/** One training example: a token window and its next-token targets. */
+interface TrainChunk {
+    inputs: number[];
+    targets: number[];
+}
+/** A batch of equal-length chunks trained in one optimiser step. */
+interface TrainBatch {
+    inputs: number[];
+    targets: number[];
+    batch: number;
+    seqLen: number;
+}
 export declare class MambaTrainer {
     model: HybridMambaModel;
     tokenizer: BPETokenizer | null;
@@ -40,29 +119,90 @@ export declare class MambaTrainer {
      * 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.
+     * both reuse the SAME `m`/`v` 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;
     private _step;
-    private _adamwPipeline;
-    private _clipReducePipeline;
-    private _clipScalePipeline;
+    private readonly _dims;
+    /** CPU mirror of every model tensor — the gradient engine's working set. */
+    private _mirror;
+    private _mirrorByName;
+    /** Low-rank adapters by parameter name, when training with `lora`. */
+    private _adapters;
+    /** Frozen base weights the adapters sit on top of (LoRA only). */
+    private _loraBase;
+    private _shard;
     constructor(model: HybridMambaModel | MambaModel, tokenizer?: BPETokenizer | null);
+    /**
+     * The low-rank adapters this trainer is training, by parameter name — empty
+     * unless `train({ lora })` was used. Exposed so a caller can serialise just the
+     * adapter (a few hundred KB) instead of the whole checkpoint.
+     */
+    get adapters(): ReadonlyMap<string, LoRAAdapter>;
+    /** Trainable scalar count for the current mode — the LoRA saving, measured. */
+    trainableParamCount(): number;
+    /** Whether this trainer owns parameter `index` under the active shard. */
+    private _owns;
+    /**
+     * Build (once) an adapter for every LoRA-target parameter, snapshot the frozen
+     * base, and return the set. Rank is clamped per parameter to its smaller
+     * dimension so one global rank works across projections of very different shape.
+     */
+    private _installAdapters;
+    /** Write `base + adapter.delta()` into the mirror, so the forward sees the adapted weights. */
+    private _materialiseAdapters;
     /** Get-or-create the Adam first/second moments for a parameter, by name. */
     private _momentFor;
+    /**
+     * Pull every model tensor off the GPU into the CPU mirror. Done once per
+     * `train()` call (not per step) so an externally loaded checkpoint or a
+     * previous session's weights are picked up, after which the mirror is
+     * authoritative and each step writes the updated tensors back.
+     */
+    private _syncMirrorFromGpu;
     train(input: string | number[], opts?: TrainOptions): Promise<number[]>;
+    /**
+     * One optimiser step over one batch: real forward + backward, a global
+     * gradient-norm clip over the trainable set, then AdamW per parameter with
+     * its OWN gradient, written straight back to the GPU buffer.
+     */
     private _trainStep;
+    /**
+     * AdamW over the trainable set. Each parameter is paired with its OWN
+     * gradient by NAME — the bug this replaces passed one shared `dLogits`
+     * buffer to every parameter — and the updated tensor is uploaded to the GPU
+     * so inference serves it.
+     */
     private _adamwStep;
+    /** Get-or-create moments for a synthetic parameter (a LoRA factor), by name. */
+    private _momentNamed;
     /**
-     * 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.
+     * Bytes of optimiser state (Adam m + v) this trainer holds. Under a shard that
+     * is 1/count of the full state; under LoRA it is the adapters' state, not the
+     * weights'. The measurable point of both features.
      */
-    private _clipGradients;
+    optimizerStateBytes(): number;
     evaluate(input: string | number[]): Promise<number>;
 }
+/**
+ * Derive the gradient engine's dimensions from a model, rejecting layer types
+ * the CPU reference cannot differentiate.
+ *
+ * A hybrid schedule (`jamba`/`zamba`, or any explicit schedule containing
+ * `mamba2`/`mamba3`/`attention`) fails LOUDLY here rather than silently training
+ * on gradients that do not exist — which is exactly what the old trainer did.
+ */
+export declare function cpuDimsFor(model: HybridMambaModel): CpuModelDims;
+export declare function buildChunks(ids: number[], seqLen: number): TrainChunk[];
+/**
+ * Group equal-length chunks into batches of at most `batchSize`.
+ *
+ * `batchSize` used to be accepted and ignored — every chunk trained alone. The
+ * batch dimension the model and the gradient engine both take is real, so it is
+ * honoured here; a ragged tail chunk simply trains as its own batch of one.
+ */
+export declare function buildBatches(chunks: TrainChunk[], batchSize: number): TrainBatch[];
+export {};
 //# sourceMappingURL=trainer.d.ts.map

package/dist/training/trainer.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"trainer.d.ts","sourceRoot":"","sources":["../../src/training/trainer.ts"],"names":[],"mappings":"AAAA;;GAEG;AAeH,OAAO,EAAE,gBAAgB,EAAE,UAAU,EAAE,MAAM,yBAAyB,CAAC;AACvE,OAAO,EAAE,YAAY,EAAE,MAAM,qBAAqB,CAAC;AAGnD,MAAM,WAAW,YAAY;IAC3B,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,GAAG,CAAC,EAAE,MAAM,CAAC;IACb,IAAI,CAAC,EAAE,OAAO,CAAC;IACf;uFACmF;IACnF,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB;;;+EAG2E;IAC3E,aAAa,CAAC,EAAE,OAAO,CAAC;IACxB,UAAU,CAAC,EAAE,CAAC,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,QAAQ,CAAC,EAAE,MAAM,KAAK,IAAI,CAAC,GAAG,IAAI,CAAC;CAChF;AAED;;;;;;GAMG;AACH,eAAO,MAAM,cAAc,OAAO,CAAC;AAkBnC,qBAAa,YAAY;IACrB,KAAK,EAAE,gBAAgB,CAAC;IACxB,SAAS,EAAE,YAAY,GAAG,IAAI,CAAC;IAC/B,MAAM,EAAE,SAAS,CAAC;IAClB;;;;;;;OAOG;IACH,OAAO,CAAC,QAAQ,CAA2B;IAC3C,OAAO,CAAC,KAAK,CAAS;IACtB,OAAO,CAAC,cAAc,CAAqB;IAC3C,OAAO,CAAC,mBAAmB,CAAqB;IAChD,OAAO,CAAC,kBAAkB,CAAqB;gBAEnC,KAAK,EAAE,gBAAgB,GAAG,UAAU,EAAE,SAAS,GAAE,YAAY,GAAG,IAAW;IAavF,4EAA4E;IAC5E,OAAO,CAAC,UAAU;IAYZ,KAAK,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM,EAAE,EAAE,IAAI,GAAE,YAAiB,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC;YA2EnE,UAAU;YAqDV,UAAU;IAiCxB;;;;;;OAMG;YACW,cAAc;IAgCtB,QAAQ,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM,EAAE,GAAG,OAAO,CAAC,MAAM,CAAC;CA4B5D"}
+{"version":3,"file":"trainer.d.ts","sourceRoot":"","sources":["../../src/training/trainer.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;GA2BG;AAQH,OAAO,EAAsB,KAAK,SAAS,EAAE,MAAM,mBAAmB,CAAC;AACvE,OAAO,EAAE,WAAW,EAAE,KAAK,UAAU,EAAE,MAAM,WAAW,CAAC;AACzD,OAAO,EAOH,KAAK,YAAY,EAEpB,MAAM,gBAAgB,CAAC;AACxB,OAAO,EAAE,gBAAgB,EAAE,UAAU,EAAE,MAAM,yBAAyB,CAAC;AACvE,OAAO,EAAE,YAAY,EAAE,MAAM,qBAAqB,CAAC;AAGnD,MAAM,WAAW,YAAY;IAC3B,YAAY,CAAC,EAAE,MAAM,CAAC;IACtB,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB;8EAC0E;IAC1E,SAAS,CAAC,EAAE,MAAM,CAAC;IACnB,MAAM,CAAC,EAAE,MAAM,CAAC;IAChB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,WAAW,CAAC,EAAE,MAAM,CAAC;IACrB,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,KAAK,CAAC,EAAE,MAAM,CAAC;IACf,GAAG,CAAC,EAAE,MAAM,CAAC;IACb,IAAI,CAAC,EAAE,OAAO,CAAC;IACf;uFACmF;IACnF,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB;;;;OAIG;IACH,aAAa,CAAC,EAAE,OAAO,CAAC;IACxB;;;;;;;;;;;OAWG;IACH,IAAI,CAAC,EAAE,UAAU,GAAG;QAAE,OAAO,CAAC,EAAE,SAAS,MAAM,EAAE,CAAA;KAAE,CAAC;IACpD;;;;;OAKG;IACH,oBAAoB,CAAC,EAAE,MAAM,CAAC;IAC9B;;;;;;OAMG;IACH,KAAK,CAAC,EAAE,SAAS,CAAC;IAClB;;;;OAIG;IACH,uBAAuB,CAAC,EAAE,OAAO,CAAC;IAClC,UAAU,CAAC,EAAE,CAAC,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,MAAM,EAAE,QAAQ,CAAC,EAAE,MAAM,KAAK,IAAI,CAAC,GAAG,IAAI,CAAC;CAChF;AAED;;;;;;GAMG;AACH,eAAO,MAAM,cAAc,OAAO,CAAC;AAmBnC,uEAAuE;AACvE,UAAU,UAAU;IAAG,MAAM,EAAE,MAAM,EAAE,CAAC;IAAC,OAAO,EAAE,MAAM,EAAE,CAAA;CAAE;AAE5D,oEAAoE;AACpE,UAAU,UAAU;IAAG,MAAM,EAAE,MAAM,EAAE,CAAC;IAAC,OAAO,EAAE,MAAM,EAAE,CAAC;IAAC,KAAK,EAAE,MAAM,CAAC;IAAC,MAAM,EAAE,MAAM,CAAA;CAAE;AAE3F,qBAAa,YAAY;IACrB,KAAK,EAAE,gBAAgB,CAAC;IACxB,SAAS,EAAE,YAAY,GAAG,IAAI,CAAC;IAC/B,MAAM,EAAE,SAAS,CAAC;IAClB;;;;;;;OAOG;IACH,OAAO,CAAC,QAAQ,CAA2B;IAC3C,OAAO,CAAC,KAAK,CAAS;IACtB,OAAO,CAAC,QAAQ,CAAC,KAAK,CAAe;IACrC,4EAA4E;IAC5E,OAAO,CAAC,OAAO,CAAyB;IACxC,OAAO,CAAC,aAAa,CAAmC;IACxD,sEAAsE;IACtE,OAAO,CAAC,SAAS,CAA2B;IAC5C,kEAAkE;IAClE,OAAO,CAAC,SAAS,CAA4B;IAC7C,OAAO,CAAC,MAAM,CAAmB;gBAErB,KAAK,EAAE,gBAAgB,GAAG,UAAU,EAAE,SAAS,GAAE,YAAY,GAAG,IAAW;IAevF;;;;OAIG;IACH,IAAI,QAAQ,IAAI,WAAW,CAAC,MAAM,EAAE,WAAW,CAAC,CAA2B;IAE3E,+EAA+E;IAC/E,mBAAmB,IAAI,MAAM;IAW7B,0EAA0E;IAC1E,OAAO,CAAC,KAAK;IAIb;;;;OAIG;IACH,OAAO,CAAC,gBAAgB;IAsBxB,+FAA+F;IAC/F,OAAO,CAAC,oBAAoB;IAU5B,4EAA4E;IAC5E,OAAO,CAAC,UAAU;IASlB;;;;;OAKG;YACW,kBAAkB;IAW1B,KAAK,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM,EAAE,EAAE,IAAI,GAAE,YAAiB,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC;IAmFjF;;;;OAIG;YACW,UAAU;IAyExB;;;;;OAKG;IACH,OAAO,CAAC,UAAU;IA+ClB,gFAAgF;IAChF,OAAO,CAAC,YAAY;IASpB;;;;OAIG;IACH,mBAAmB,IAAI,MAAM;IAMvB,QAAQ,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM,EAAE,GAAG,OAAO,CAAC,MAAM,CAAC;CA4B5D;AAED;;;;;;;GAOG;AACH,wBAAgB,UAAU,CAAC,KAAK,EAAE,gBAAgB,GAAG,YAAY,CA0BhE;AAED,wBAAgB,WAAW,CAAC,GAAG,EAAE,MAAM,EAAE,EAAE,MAAM,EAAE,MAAM,GAAG,UAAU,EAAE,CAiBvE;AAED;;;;;;GAMG;AACH,wBAAgB,YAAY,CAAC,MAAM,EAAE,UAAU,EAAE,EAAE,SAAS,EAAE,MAAM,GAAG,UAAU,EAAE,CAkBlF"}