@seanhogg/builderforce-memory-engine 2026.7.12 → 2026.9.27

package/dist/lm/evermind_lm.js

@@ -23,6 +23,7 @@
 import { SharedExpertMoE } from "../moe/moe_model.js";
 import { crossEntropyLoss, crossEntropyGrad } from "../training/autograd.js";
 import { AdamW } from "../optim/adamw.js";
+import { DynamicLossScaler, roundFp16 } from "../training/mixed_precision.js";
 import { SeededRng } from "../utils/rng.js";
 import { quantizeFp16, dequantizeFp16 } from "../utils/quantization.js";
 import { appendCrcTrailer, verifyCrcTrailer } from "../utils/crc32.js";
@@ -39,6 +40,15 @@ export const DEFAULT_LM_SEED = 0x45564c4d; // "EVLM"
 const MAGIC = 0x45564c30; // "EVL0"
 export class EvermindLM {
     config;
+    /**
+     * Token-positions pushed through a block since the last {@link resetStats}.
+     *
+     * The cost metric that matters for the tool-calling path: scoring N candidate
+     * tool names used to replay the WHOLE prompt N times, so this counted
+     * N x (prompt + candidate) x layers. With {@link stepDecode} the prompt is paid
+     * for once. Exposed so a benchmark can assert the saving rather than assume it.
+     */
+    _positionsEvaluated = 0;
     /** Tied token embedding / output head: vocabSize × dModel (row-major). */
     emb;
     gEmb;
@@ -112,65 +122,96 @@ export class EvermindLM {
         }
     }
     // ── Forward ────────────────────────────────────────────────────────────────
-    /** Run the model over a token sequence; returns per-position logits + a cache. */
-    forward(tokens) {
-        const { dModel, convKernel, numLayers, vocabSize } = this.config;
-        const T = tokens.length;
-        // Embed.
-        let x = tokens.map((tok) => {
+    /** Embed a token sequence into per-position channel vectors. */
+    _embed(tokens) {
+        const { dModel } = this.config;
+        return tokens.map((tok) => {
             const row = new Float32Array(dModel);
             const off = tok * dModel;
             for (let c = 0; c < dModel; c++)
                 row[c] = this.emb[off + c];
             return row;
         });
-        const layers = [];
-        for (let l = 0; l < numLayers; l++) {
-            const layerIn = x;
-            const ker = this.conv[l];
-            const nConv = this.nConv[l];
-            const nMoe = this.nMoe[l];
-            // Pre-norm → depthwise causal conv → residual.
-            const normedConv = [];
-            const rmsConv = [];
-            for (let t = 0; t < T; t++) {
-                const { y, r } = rmsNorm(layerIn[t], nConv);
-                normedConv.push(y);
-                rmsConv.push(r);
-            }
-            const afterConv = [];
-            for (let t = 0; t < T; t++) {
-                const out = Float32Array.from(layerIn[t]); // residual base
-                for (let c = 0; c < dModel; c++) {
-                    let acc = 0;
-                    for (let j = 0; j < convKernel; j++) {
-                        const ti = t - j;
-                        if (ti >= 0)
-                            acc += ker[c * convKernel + j] * normedConv[ti][c];
-                    }
-                    out[c] = out[c] + acc;
-                }
-                afterConv.push(out);
-            }
-            // Pre-norm → MoE channel mixer → residual.
-            const rmsMoe = [];
-            const moeCache = [];
-            const afterMoe = [];
-            for (let t = 0; t < T; t++) {
-                const { y, r } = rmsNorm(afterConv[t], nMoe);
-                rmsMoe.push(r);
-                const out = Float32Array.from(afterConv[t]); // residual base
-                const mr = this.moe[l].forward(y);
-                for (let c = 0; c < dModel; c++)
-                    out[c] = out[c] + mr.output[c];
-                afterMoe.push(out);
-                moeCache.push(mr.cache);
+    }
+    /**
+     * One (conv + MoE) block: pre-norm → depthwise causal conv → residual, then
+     * pre-norm → MoE channel mixer → residual. Returns the block output and the
+     * activation cache its backward needs. Isolating this is what lets
+     * {@link lossAndBackwardCheckpointed} recompute a layer's activations on demand
+     * instead of retaining every layer's cache at once.
+     */
+    _forwardLayer(l, layerIn) {
+        const T = layerIn.length;
+        const window = [];
+        const normedConv = [];
+        const rmsConv = [];
+        const afterConv = [];
+        const rmsMoe = [];
+        const moeCache = [];
+        const afterMoe = [];
+        for (let t = 0; t < T; t++) {
+            const st = this._stepLayerPosition(l, layerIn[t], window);
+            normedConv.push(st.normed);
+            rmsConv.push(st.rmsConv);
+            afterConv.push(st.afterConv);
+            rmsMoe.push(st.rmsMoe);
+            moeCache.push(st.moeCache);
+            afterMoe.push(st.out);
+        }
+        return { afterMoe, cache: { layerIn, normedConv, rmsConv, afterConv, rmsMoe, moeCache } };
+    }
+    /**
+     * ONE position through one (conv + MoE) block, given the trailing conv window.
+     *
+     * The single place the block's maths lives: the batch path
+     * ({@link _forwardLayer}, which needs the full activation cache for backward)
+     * and the incremental path ({@link stepDecode}, which needs none of it) both
+     * run this, so they cannot drift apart. `window` holds the previous
+     * `convKernel - 1` RMSNormed conv inputs, OLDEST FIRST, and is advanced in
+     * place — a shorter window is exactly the causal zero-padding at the start of
+     * a sequence.
+     */
+    _stepLayerPosition(l, xt, window) {
+        const { dModel, convKernel } = this.config;
+        const ker = this.conv[l];
+        const { y: normed, r: rConv } = rmsNorm(xt, this.nConv[l]);
+        const afterConv = Float32Array.from(xt); // residual base
+        for (let c = 0; c < dModel; c++) {
+            let acc = ker[c * convKernel] * normed[c];
+            for (let j = 1; j < convKernel; j++) {
+                const prev = window[window.length - j]; // position t - j, or undefined before the start
+                if (prev)
+                    acc += ker[c * convKernel + j] * prev[c];
             }
-            layers.push({ layerIn, normedConv, rmsConv, afterConv, rmsMoe, moeCache });
-            x = afterMoe;
+            afterConv[c] = afterConv[c] + acc;
         }
-        // Tied head: logits_t[v] = x_t · emb[v].
-        const logits = x.map((xt) => {
+        window.push(normed);
+        while (window.length > convKernel - 1)
+            window.shift();
+        const { y: normedMoe, r: rMoe } = rmsNorm(afterConv, this.nMoe[l]);
+        const out = Float32Array.from(afterConv); // residual base
+        const mr = this.moe[l].forward(normedMoe);
+        for (let c = 0; c < dModel; c++)
+            out[c] = out[c] + mr.output[c];
+        this._positionsEvaluated++;
+        return {
+            normed, rmsConv: rConv, afterConv, rmsMoe: rMoe,
+            moeCache: mr.cache, out,
+        };
+    }
+    /** One token's embedding row (a copy — callers write into it). */
+    _embedOne(tok) {
+        const { dModel } = this.config;
+        const row = new Float32Array(dModel);
+        const off = tok * dModel;
+        for (let c = 0; c < dModel; c++)
+            row[c] = this.emb[off + c];
+        return row;
+    }
+    /** Tied output head: logits_t[v] = x_t · emb[v]. */
+    _head(x) {
+        const { dModel, vocabSize } = this.config;
+        return x.map((xt) => {
             const lg = new Float32Array(vocabSize);
             for (let v = 0; v < vocabSize; v++) {
                 let acc = 0;
@@ -181,30 +222,92 @@ export class EvermindLM {
             }
             return lg;
         });
-        return { logits, cache: { tokens, layers, finalX: x } };
+    }
+    /** Run the model over a token sequence; returns per-position logits + a cache. */
+    forward(tokens) {
+        let x = this._embed(tokens);
+        const layers = [];
+        for (let l = 0; l < this.config.numLayers; l++) {
+            const { afterMoe, cache } = this._forwardLayer(l, x);
+            layers.push(cache);
+            x = afterMoe;
+        }
+        return { logits: this._head(x), cache: { tokens, layers, finalX: x } };
+    }
+    /**
+     * Produces a single fixed-length embedding vector for a token sequence — the
+     * CPU counterpart of `HybridMambaModel.embed()`, to the same contract.
+     *
+     * Runs the full block stack and takes the final hidden state — i.e. exactly the
+     * state the tied LM head consumes in {@link forward} — then mean-pools across
+     * sequence positions and L2-normalises. The result has length `dModel`, so
+     * cosine similarity between two embeddings reduces to a dot product.
+     *
+     * Like the GPU path it skips the (expensive) head projection: it only needs the
+     * `dModel`-wide hidden state, not `vocabSize` logits. Unlike the GPU path it
+     * needs no WebGPU device at all, which is what lets a headless host (the stdio
+     * MCP server) rank recall by SSM embedding instead of word overlap.
+     *
+     * Streamed one position at a time through the SHARED block implementation
+     * ({@link _stepLayerPosition}, the same routine {@link forward} and
+     * {@link stepDecode} run), so it can never drift from them and it retains no
+     * per-position activation cache.
+     *
+     * The embedding reflects whatever the model currently knows — an untrained
+     * model behaves like a random projection of the token embeddings (still
+     * lexically discriminative), and the representation sharpens as the model is
+     * adapted/distilled.
+     */
+    embed(tokenIds) {
+        const { dModel, numLayers } = this.config;
+        const out = new Float32Array(dModel);
+        const seqLen = tokenIds.length;
+        if (seqLen === 0)
+            return out;
+        // One rolling conv window per layer — a shorter window at the start of the
+        // sequence IS the causal zero-padding, exactly as in stepDecode.
+        const windows = Array.from({ length: numLayers }, () => []);
+        for (let t = 0; t < seqLen; t++) {
+            let x = this._embedOne(tokenIds[t]);
+            for (let l = 0; l < numLayers; l++)
+                x = this._stepLayerPosition(l, x, windows[l]).out;
+            for (let c = 0; c < dModel; c++)
+                out[c] = out[c] + x[c];
+        }
+        // Mean-pool across positions.
+        for (let c = 0; c < dModel; c++)
+            out[c] = out[c] / seqLen;
+        // L2-normalise so cosine similarity reduces to a dot product.
+        let norm = 0;
+        for (let c = 0; c < dModel; c++)
+            norm += out[c] * out[c];
+        norm = Math.sqrt(norm) || 1;
+        for (let c = 0; c < dModel; c++)
+            out[c] = out[c] / norm;
+        return out;
+    }
+    /** Text-level embedding: encode with `codec`, then {@link embed}. */
+    embedText(text, codec) {
+        return this.embed(codec.encode(text));
     }
     // ── Loss + backward ──────────────────────────────────────────────────────────
     /**
-     * Next-token cross-entropy over the sequence (predict tokens[t+1] from
-     * position t), accumulating exact gradients. Returns the mean loss. Call
-     * {@link zeroGrad} before and an optimiser step after.
+     * Head + tied-embedding gradient. Accumulates dL/d(head→emb) into gEmb and
+     * returns the mean next-token loss plus dL/d(finalX). Shared by the full and
+     * checkpointed backward paths so the head maths lives in one place.
      */
-    lossAndBackward(tokens) {
-        const { dModel, convKernel, numLayers, vocabSize } = this.config;
+    _headBackward(tokens, logits, finalX) {
+        const { dModel, vocabSize } = this.config;
         const T = tokens.length;
-        if (T < 2)
-            return 0;
-        const { logits, cache } = this.forward(tokens);
         const predPositions = T - 1; // positions 0..T-2 predict the next token
         const inv = 1 / predPositions;
-        // dL/d(finalX_t) and head gradient into the tied embedding.
         const dX = Array.from({ length: T }, () => new Float32Array(dModel));
         let loss = 0;
         for (let t = 0; t < predPositions; t++) {
             const target = tokens[t + 1];
             loss += crossEntropyLoss(logits[t], target) * inv;
             const dLogit = crossEntropyGrad(logits[t], target); // probs - onehot
-            const xt = cache.finalX[t];
+            const xt = finalX[t];
             for (let v = 0; v < vocabSize; v++) {
                 const g = dLogit[v] * inv;
                 if (g === 0)
@@ -216,62 +319,182 @@ export class EvermindLM {
                 }
             }
         }
-        // Backprop through layers in reverse.
-        for (let l = numLayers - 1; l >= 0; l--) {
-            const lc = cache.layers[l];
-            const ker = this.conv[l];
-            const gker = this.gConv[l];
-            const nConv = this.nConv[l];
-            const gNConv = this.gNConv[l];
-            const nMoe = this.nMoe[l];
-            const gNMoe = this.gNMoe[l];
-            // MoE sub-block: afterMoe = afterConv + MoE(RMSNorm(afterConv, nMoe)).
-            const dAfterConv = [];
-            for (let t = 0; t < T; t++) {
-                const dMoeNormed = this.moe[l].backward(dX[t], lc.moeCache[t]);
-                const { dx, dgain } = rmsNormBackward(dMoeNormed, lc.afterConv[t], lc.rmsMoe[t], nMoe);
-                for (let c = 0; c < dModel; c++)
-                    gNMoe[c] = gNMoe[c] + dgain[c];
-                const d = Float32Array.from(dX[t]); // residual passthrough
-                for (let c = 0; c < dModel; c++)
-                    d[c] = d[c] + dx[c];
-                dAfterConv.push(d);
-            }
-            // Conv sub-block: afterConv = layerIn + conv(RMSNorm(layerIn, nConv)).
-            const dNormedConv = Array.from({ length: T }, () => new Float32Array(dModel));
-            const dLayerIn = dAfterConv.map((v) => Float32Array.from(v)); // residual passthrough
-            for (let t = 0; t < T; t++) {
-                for (let c = 0; c < dModel; c++) {
-                    const dmix = dAfterConv[t][c];
-                    if (dmix === 0)
+        return { loss, dX };
+    }
+    /**
+     * Backward through one (conv + MoE) block given dL/d(block output) and the
+     * block's activation cache. Accumulates conv/norm/MoE gradients and returns
+     * dL/d(block input). The inverse of {@link _forwardLayer}.
+     */
+    _backwardLayer(l, dOut, lc) {
+        const { dModel, convKernel } = this.config;
+        const T = dOut.length;
+        const ker = this.conv[l];
+        const gker = this.gConv[l];
+        const nConv = this.nConv[l];
+        const gNConv = this.gNConv[l];
+        const nMoe = this.nMoe[l];
+        const gNMoe = this.gNMoe[l];
+        // MoE sub-block: afterMoe = afterConv + MoE(RMSNorm(afterConv, nMoe)).
+        const dAfterConv = [];
+        for (let t = 0; t < T; t++) {
+            const dMoeNormed = this.moe[l].backward(dOut[t], lc.moeCache[t]);
+            const { dx, dgain } = rmsNormBackward(dMoeNormed, lc.afterConv[t], lc.rmsMoe[t], nMoe);
+            for (let c = 0; c < dModel; c++)
+                gNMoe[c] = gNMoe[c] + dgain[c];
+            const d = Float32Array.from(dOut[t]); // residual passthrough
+            for (let c = 0; c < dModel; c++)
+                d[c] = d[c] + dx[c];
+            dAfterConv.push(d);
+        }
+        // Conv sub-block: afterConv = layerIn + conv(RMSNorm(layerIn, nConv)).
+        const dNormedConv = Array.from({ length: T }, () => new Float32Array(dModel));
+        const dLayerIn = dAfterConv.map((v) => Float32Array.from(v)); // residual passthrough
+        for (let t = 0; t < T; t++) {
+            for (let c = 0; c < dModel; c++) {
+                const dmix = dAfterConv[t][c];
+                if (dmix === 0)
+                    continue;
+                for (let j = 0; j < convKernel; j++) {
+                    const ti = t - j;
+                    if (ti < 0)
                         continue;
-                    for (let j = 0; j < convKernel; j++) {
-                        const ti = t - j;
-                        if (ti < 0)
-                            continue;
-                        gker[c * convKernel + j] = gker[c * convKernel + j] + dmix * lc.normedConv[ti][c];
-                        dNormedConv[ti][c] = dNormedConv[ti][c] + dmix * ker[c * convKernel + j];
-                    }
-                }
-            }
-            for (let t = 0; t < T; t++) {
-                const { dx, dgain } = rmsNormBackward(dNormedConv[t], lc.layerIn[t], lc.rmsConv[t], nConv);
-                for (let c = 0; c < dModel; c++) {
-                    gNConv[c] = gNConv[c] + dgain[c];
-                    dLayerIn[t][c] = dLayerIn[t][c] + dx[c];
+                    gker[c * convKernel + j] = gker[c * convKernel + j] + dmix * lc.normedConv[ti][c];
+                    dNormedConv[ti][c] = dNormedConv[ti][c] + dmix * ker[c * convKernel + j];
                 }
             }
-            for (let t = 0; t < T; t++)
-                dX[t] = dLayerIn[t];
         }
-        // Embedding lookup: dX at layer-0 input flows into the row for token_t.
         for (let t = 0; t < T; t++) {
+            const { dx, dgain } = rmsNormBackward(dNormedConv[t], lc.layerIn[t], lc.rmsConv[t], nConv);
+            for (let c = 0; c < dModel; c++) {
+                gNConv[c] = gNConv[c] + dgain[c];
+                dLayerIn[t][c] = dLayerIn[t][c] + dx[c];
+            }
+        }
+        return dLayerIn;
+    }
+    /** Embedding lookup: dL/d(layer-0 input) flows into the row for each token. */
+    _embedBackward(tokens, dIn) {
+        const { dModel } = this.config;
+        for (let t = 0; t < tokens.length; t++) {
             const off = tokens[t] * dModel;
             for (let c = 0; c < dModel; c++)
-                this.gEmb[off + c] = this.gEmb[off + c] + dX[t][c];
+                this.gEmb[off + c] = this.gEmb[off + c] + dIn[t][c];
+        }
+    }
+    /**
+     * Next-token cross-entropy over the sequence (predict tokens[t+1] from
+     * position t), accumulating exact gradients. Returns the mean loss. Call
+     * {@link zeroGrad} before and an optimiser step after.
+     */
+    lossAndBackward(tokens) {
+        if (tokens.length < 2)
+            return 0;
+        const { logits, cache } = this.forward(tokens);
+        const { loss, dX } = this._headBackward(tokens, logits, cache.finalX);
+        let d = dX;
+        for (let l = this.config.numLayers - 1; l >= 0; l--)
+            d = this._backwardLayer(l, d, cache.layers[l]);
+        this._embedBackward(tokens, d);
+        return loss;
+    }
+    /**
+     * Activation-checkpointed backward — numerically identical gradients to
+     * {@link lossAndBackward}, but retains only the per-LAYER inputs during the
+     * forward instead of every layer's full activation cache. Each layer's
+     * activations are RECOMPUTED (a cheap extra forward) when its backward runs,
+     * so peak activation memory is one layer's cache, not all of them. This is the
+     * memory-for-compute trade the cookbook pairs with FSDP to fit longer
+     * sequences / bigger models on a constrained device.
+     */
+    lossAndBackwardCheckpointed(tokens) {
+        if (tokens.length < 2)
+            return 0;
+        let x = this._embed(tokens);
+        const layerInputs = [];
+        for (let l = 0; l < this.config.numLayers; l++) {
+            layerInputs.push(x); // keep only the input; drop the cache
+            x = this._forwardLayer(l, x).afterMoe;
         }
+        const logits = this._head(x);
+        const { loss, dX } = this._headBackward(tokens, logits, x);
+        let d = dX;
+        for (let l = this.config.numLayers - 1; l >= 0; l--) {
+            const { cache } = this._forwardLayer(l, layerInputs[l]); // recompute this layer's activations
+            d = this._backwardLayer(l, d, cache);
+        }
+        this._embedBackward(tokens, d);
         return loss;
     }
+    // ── Incremental decoding (prefix cache) ──────────────────────────────────────
+    /** Positions pushed through a block since the last {@link resetStats}. */
+    get positionsEvaluated() { return this._positionsEvaluated; }
+    resetStats() { this._positionsEvaluated = 0; }
+    /** An empty decode state — the start of a sequence. */
+    newDecodeState() {
+        return {
+            convWindow: Array.from({ length: this.config.numLayers }, () => []),
+            length: 0,
+            lastLogits: null,
+        };
+    }
+    /** Deep copy, so one prefix state can seed many independent continuations. */
+    cloneDecodeState(state) {
+        return {
+            convWindow: state.convWindow.map((w) => w.map((v) => Float32Array.from(v))),
+            length: state.length,
+            lastLogits: state.lastLogits ? Float32Array.from(state.lastLogits) : null,
+        };
+    }
+    /**
+     * Consume `tokens` from `state`, returning the per-position logits and the
+     * ADVANCED state. `state` is not mutated — pass {@link cloneDecodeState} output
+     * or a fresh state; this clones internally so the caller's prefix stays reusable.
+     *
+     * Numerically identical to {@link forward} over the concatenated sequence: the
+     * shared {@link _stepLayerPosition} is the only implementation of the block.
+     */
+    stepDecode(state, tokens) {
+        const next = this.cloneDecodeState(state);
+        const logits = [];
+        for (const tok of tokens) {
+            let x = this._embedOne(tok);
+            for (let l = 0; l < this.config.numLayers; l++) {
+                x = this._stepLayerPosition(l, x, next.convWindow[l]).out;
+            }
+            const lg = this._head([x])[0];
+            logits.push(lg);
+            next.lastLogits = lg;
+            next.length++;
+        }
+        return { logits, state: next };
+    }
+    /** Run a prompt once and keep the state, so continuations are cheap. */
+    prefixState(tokens) {
+        return this.stepDecode(this.newDecodeState(), tokens).state;
+    }
+    /**
+     * Mean log-probability the model assigns to `continuation` following a cached
+     * prefix — the tool-decision vote, without replaying the prompt.
+     *
+     * Returns `-Infinity` for an empty continuation, and requires a prefix with at
+     * least one position (there must be a position to predict the first
+     * continuation token FROM).
+     */
+    scoreContinuation(prefix, continuation) {
+        if (continuation.length === 0)
+            return -Infinity;
+        if (!prefix.lastLogits)
+            return -Infinity;
+        // The last prefix position predicts continuation[0]; each continuation
+        // position then predicts the next one, so the final one is never scored.
+        const { logits } = this.stepDecode(prefix, continuation.slice(0, -1));
+        let total = logProbOfToken(prefix.lastLogits, continuation[0]);
+        for (let i = 1; i < continuation.length; i++) {
+            total += logProbOfToken(logits[i - 1], continuation[i]);
+        }
+        return total / continuation.length;
+    }
     // ── Generation ───────────────────────────────────────────────────────────────
     /**
      * Text-level generation: encode the prompt, generate, decode. `codec` is any
@@ -282,20 +505,25 @@ export class EvermindLM {
     generateText(prompt, codec, opts) {
         return codec.decode(this.generate(codec.encode(prompt), opts));
     }
-    /** Greedy / temperature-sampled autoregressive generation. Returns NEW token ids. */
+    /**
+     * Greedy / temperature-sampled autoregressive generation. Returns NEW token ids.
+     *
+     * Carries an incremental decode state, so the prompt is processed ONCE and each
+     * new token costs one position. The previous implementation re-ran the whole
+     * sequence per token, making generation quadratic in the prompt length.
+     */
     generate(prompt, opts) {
         const temperature = opts.temperature ?? 0;
         const rng = temperature > 0 ? new SeededRng((opts.seed ?? 1) >>> 0 || 1) : null;
-        const tokens = [...prompt];
+        let state = this.stepDecode(this.newDecodeState(), prompt.length > 0 ? prompt : [0]).state;
         const produced = [];
         for (let n = 0; n < opts.maxNewTokens; n++) {
-            const { logits } = this.forward(tokens.length > 0 ? tokens : [0]);
-            const last = logits[logits.length - 1];
+            const last = state.lastLogits;
             const next = rng ? sampleTemperature(last, temperature, rng) : argmax(last);
             produced.push(next);
-            tokens.push(next);
             if (opts.stopToken !== undefined && next === opts.stopToken)
                 break;
+            state = this.stepDecode(state, [next]).state;
         }
         return produced;
     }
@@ -437,26 +665,95 @@ export class EvermindLMTrainer {
     model;
     opts;
     adam;
+    scaler;
+    grads;
     constructor(model, opts = {}) {
         this.model = model;
         this.opts = opts;
         this.adam = new AdamW(model, opts);
+        this.grads = model.gradients();
+        this.scaler = opts.mixedPrecision
+            ? new DynamicLossScaler(typeof opts.mixedPrecision === "object" ? opts.mixedPrecision : {})
+            : null;
+    }
+    /** The dynamic loss scaler (mixed-precision mode only) — exposes scale/overflow stats. */
+    get lossScaler() {
+        return this.scaler;
+    }
+    _backward(seq) {
+        return this.opts.checkpoint ? this.model.lossAndBackwardCheckpointed(seq) : this.model.lossAndBackward(seq);
+    }
+    /** Divide accumulated gradients by `n` (accumulation averaging), in place. */
+    _scaleGrads(n) {
+        if (n === 1)
+            return;
+        const inv = 1 / n;
+        for (const g of this.grads)
+            for (let i = 0; i < g.data.length; i++)
+                g.data[i] = g.data[i] * inv;
+    }
+    /** Round accumulated gradients to fp16 precision (mixed-precision simulation), in place. */
+    _fp16Grads() {
+        for (const g of this.grads)
+            for (let i = 0; i < g.data.length; i++)
+                g.data[i] = roundFp16(g.data[i]);
+    }
+    /** Multiply accumulated gradients by `s`, in place. */
+    _mulGrads(s) {
+        if (s === 1)
+            return;
+        for (const g of this.grads)
+            for (let i = 0; i < g.data.length; i++)
+                g.data[i] = g.data[i] * s;
     }
     /** Train on a set of token sequences; returns per-epoch mean loss. */
     fit(sequences) {
         const epochs = this.opts.epochs ?? 1;
+        const accum = Math.max(1, this.opts.accumSteps ?? 1);
         const history = [];
         for (let e = 0; e < epochs; e++) {
             let total = 0;
             let n = 0;
+            let pending = 0;
+            // The scale is fixed for the duration of one accumulation window, so scale-up
+            // (per sequence) and unscale (at flush) always use the SAME factor even as the
+            // controller grows/backs it off between windows.
+            let windowScale = this.scaler ? this.scaler.scale : 1;
+            this.model.zeroGrad();
+            const flush = () => {
+                if (pending === 0)
+                    return;
+                this._scaleGrads(pending); // accumulation averaging
+                if (this.scaler) {
+                    this._fp16Grads(); // half-precision gradients
+                    const overflow = this.scaler.check(this.grads);
+                    if (this.scaler.update(overflow)) {
+                        this._mulGrads(1 / windowScale); // unscale by the SAME factor we scaled by
+                        this.adam.step();
+                    } // else: skip the step, scale has backed off
+                    windowScale = this.scaler.scale; // next window uses the updated scale
+                }
+                else {
+                    this.adam.step();
+                }
+                this.model.zeroGrad();
+                pending = 0;
+            };
             for (const seq of sequences) {
                 if (seq.length < 2)
                     continue;
-                this.model.zeroGrad();
-                total += this.model.lossAndBackward(seq);
-                this.adam.step();
+                if (pending === 0 && this.scaler)
+                    windowScale = this.scaler.scale;
+                const loss = this._backward(seq);
+                // Emulate scaling the loss before backward: the gradient is linear in the
+                // loss, so scaling the just-produced gradients by the loss scale is equivalent.
+                this._mulGrads(windowScale);
+                total += loss;
                 n++;
+                if (++pending >= accum)
+                    flush();
             }
+            flush();
             history.push(n > 0 ? total / n : 0);
         }
         return history;
@@ -520,4 +817,24 @@ function sampleTemperature(logits, temperature, rng) {
     }
     return probs.length - 1;
 }
+/**
+ * Log-probability the model assigned to `id` at a position, from that position's
+ * raw logits. Log-softmax computed max-shifted so a long-tail logit cannot
+ * overflow `exp`.
+ */
+export function logProbOfToken(row, id) {
+    let max = -Infinity;
+    for (let i = 0; i < row.length; i++) {
+        const v = row[i];
+        if (v > max)
+            max = v;
+    }
+    if (!Number.isFinite(max))
+        return -Infinity;
+    let sum = 0;
+    for (let i = 0; i < row.length; i++)
+        sum += Math.exp(row[i] - max);
+    const logit = id >= 0 && id < row.length ? row[id] : -Infinity;
+    return logit - (max + Math.log(sum));
+}
 //# sourceMappingURL=evermind_lm.js.map