@seanhogg/builderforce-memory-engine 2026.6.18

package/src/model/sequence_layer.ts

@@ -0,0 +1,29 @@
+/**
+ * sequence_layer.ts – Common interface implemented by all block types.
+ *
+ * Mamba1Block, Mamba2Block, Mamba3Block, and AttentionBlock all implement
+ * SequenceLayer so that HybridMambaModel can iterate layers generically.
+ */
+
+export interface LayerForwardResult {
+    output : GPUBuffer;
+    cache  : unknown;  // type-specific per layer variant
+}
+
+export interface LayerParam {
+    buf   : GPUBuffer;
+    numel : number;
+    name  : string;
+}
+
+export type LayerType = 'mamba1' | 'mamba2' | 'mamba3' | 'attention';
+
+export interface SequenceLayer {
+    readonly layerType: LayerType;
+
+    forward(xBuf: GPUBuffer, batch: number, seqLen: number): LayerForwardResult;
+    parameters(): LayerParam[];
+    getTrainableParams(): LayerParam[];
+    setWSLAMode(enabled: boolean): void;
+    destroy(): void;
+}

package/src/tokenizer/bpe.ts

@@ -0,0 +1,186 @@
+/**
+ * bpe.ts – Browser-side Byte Pair Encoding (BPE) tokenizer.
+ */
+
+export interface BPEEncodeOptions {
+  addBos?: boolean;
+  addEos?: boolean;
+}
+
+export type PadSide = 'right' | 'left';
+
+function buildByteEncoder(): Map<number, string> {
+    const enc = new Map<number, string>();
+    const ranges: [number, number][] = [
+        [0x21, 0x7E],
+        [0xA1, 0xAC],
+        [0xAE, 0xFF],
+    ];
+    let n = 0;
+    for (const [lo, hi] of ranges) {
+        for (let b = lo; b <= hi; b++) {
+            enc.set(b, String.fromCodePoint(b));
+        }
+    }
+    for (let b = 0; b < 256; b++) {
+        if (!enc.has(b)) {
+            enc.set(b, String.fromCodePoint(256 + n));
+            n++;
+        }
+    }
+    return enc;
+}
+
+const BYTE_ENCODER = buildByteEncoder();
+const BYTE_DECODER = new Map([...BYTE_ENCODER].map(([k, v]) => [v, k]));
+
+const PRE_TOKENIZE_RE =
+    /(?:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+/gu;
+
+export class BPETokenizer {
+    vocab: Map<string, number>;
+    idToToken: Map<number, string>;
+    merges: Map<string, number>;
+    bosToken: string;
+    eosToken: string;
+    padToken: string;
+    unkToken: string;
+    bosId: number | null;
+    eosId: number | null;
+    padId: number | null;
+
+    constructor() {
+        this.vocab      = new Map();
+        this.idToToken  = new Map();
+        this.merges     = new Map();
+        this.bosToken   = '<|im_start|>';
+        this.eosToken   = '<|im_end|>';
+        this.padToken   = '<|endoftext|>';
+        this.unkToken   = '<unk>';
+        this.bosId      = null;
+        this.eosId      = null;
+        this.padId      = null;
+    }
+
+    async load(vocab: string | Record<string, number>, merges: string | string[]): Promise<void> {
+        let vocabObj: Record<string, number>;
+        if (typeof vocab === 'string') {
+            const res = await fetch(vocab);
+            vocabObj = await res.json() as Record<string, number>;
+        } else {
+            vocabObj = vocab;
+        }
+        this.vocab     = new Map(Object.entries(vocabObj).map(([k, v]) => [k, Number(v)]));
+        this.idToToken = new Map([...this.vocab].map(([k, v]) => [v, k]));
+
+        let mergeLines: string[];
+        if (typeof merges === 'string') {
+            const res = await fetch(merges);
+            const txt = await res.text();
+            mergeLines = txt.split('\n').filter(l => l && !l.startsWith('#'));
+        } else {
+            mergeLines = merges;
+        }
+        this.merges = new Map();
+        mergeLines.forEach((line, rank) => {
+            this.merges.set(line.trim(), rank);
+        });
+
+        this.bosId = this.vocab.get(this.bosToken) ?? null;
+        this.eosId = this.vocab.get(this.eosToken) ?? null;
+        this.padId = this.vocab.get(this.padToken) ?? null;
+    }
+
+    loadFromObjects(vocabObj: Record<string, number>, mergeArr: string[]): void {
+        this.vocab     = new Map(Object.entries(vocabObj).map(([k, v]) => [k, Number(v)]));
+        this.idToToken = new Map([...this.vocab].map(([k, v]) => [v, k]));
+        this.merges    = new Map(mergeArr.map((m, i) => [m, i]));
+        this.bosId = this.vocab.get(this.bosToken) ?? null;
+        this.eosId = this.vocab.get(this.eosToken) ?? null;
+        this.padId = this.vocab.get(this.padToken) ?? null;
+    }
+
+    encode(text: string, opts: BPEEncodeOptions = {}): number[] {
+        const words = text.match(PRE_TOKENIZE_RE) ?? [];
+        const ids: number[]   = [];
+
+        if (opts.addBos && this.bosId !== null) ids.push(this.bosId);
+
+        for (const word of words) {
+            const bytes    = new TextEncoder().encode(word);
+            const byteStr  = Array.from(bytes).map(b => BYTE_ENCODER.get(b) ?? '?').join('');
+            const bpeTokens = this._bpe(byteStr);
+
+            for (const tok of bpeTokens) {
+                const id = this.vocab.get(tok);
+                if (id !== undefined) {
+                    ids.push(id);
+                } else {
+                    for (const ch of tok) {
+                        const cid = this.vocab.get(ch);
+                        if (cid !== undefined) ids.push(cid);
+                    }
+                }
+            }
+        }
+
+        if (opts.addEos && this.eosId !== null) ids.push(this.eosId);
+        return ids;
+    }
+
+    decode(ids: number[]): string {
+        let byteStr = '';
+        for (const id of ids) {
+            const tok = this.idToToken.get(id);
+            if (tok !== undefined) byteStr += tok;
+        }
+        const bytes = new Uint8Array(
+            [...byteStr].map(ch => BYTE_DECODER.get(ch) ?? ch.codePointAt(0) ?? 0)
+        );
+        try {
+            return new TextDecoder('utf-8').decode(bytes);
+        } catch {
+            return byteStr;
+        }
+    }
+
+    _bpe(word: string): string[] {
+        if (this.vocab.has(word)) return [word];
+
+        let symbols = [...word];
+
+        while (symbols.length > 1) {
+            let bestRank = Infinity;
+            let bestIdx  = -1;
+
+            for (let i = 0; i < symbols.length - 1; i++) {
+                const pair = symbols[i] + ' ' + symbols[i + 1];
+                const rank = this.merges.get(pair);
+                if (rank !== undefined && rank < bestRank) {
+                    bestRank = rank;
+                    bestIdx  = i;
+                }
+            }
+
+            if (bestIdx === -1) break;
+
+            const merged = symbols[bestIdx]! + symbols[bestIdx + 1]!;
+            symbols = [
+                ...symbols.slice(0, bestIdx),
+                merged,
+                ...symbols.slice(bestIdx + 2),
+            ];
+        }
+
+        return symbols;
+    }
+
+    padOrTruncate(ids: number[], maxLen: number, side: PadSide = 'right'): number[] {
+        if (ids.length >= maxLen) return ids.slice(0, maxLen);
+        const padId = this.padId ?? 0;
+        const pad   = new Array<number>(maxLen - ids.length).fill(padId);
+        return side === 'right' ? [...ids, ...pad] : [...pad, ...ids];
+    }
+
+    get vocabSize(): number { return this.vocab.size; }
+}

package/src/training/autograd.ts

@@ -0,0 +1,135 @@
+/**
+ * autograd.ts – Lightweight tape-based automatic differentiation engine.
+ */
+
+/* eslint-disable @typescript-eslint/no-explicit-any */
+const _gpu = globalThis as any;
+
+interface TapeEntry {
+  backward: () => void | Promise<void>;
+}
+
+let _tape: TapeEntry[] = [];
+let _gradEnabled = true;
+
+export class Tensor {
+    data: GPUBuffer | null;
+    shape: number[];
+    numel: number;
+    requiresGrad: boolean;
+    grad: GPUBuffer | null;
+    _gradFn: number | null;
+
+    constructor(data: GPUBuffer | null, shape: number[], requiresGrad = false) {
+        this.data         = data;
+        this.shape        = shape;
+        this.numel        = shape.reduce((a, b) => a * b, 1);
+        this.requiresGrad = requiresGrad;
+        this.grad         = null;
+        this._gradFn      = null;
+    }
+
+    get byteSize(): number { return this.numel * 4; }
+
+    zeroGrad(device: GPUDevice): void {
+        if (this.grad) {
+            device.queue.writeBuffer(this.grad, 0, new Float32Array(this.numel));
+        }
+    }
+
+    destroy(): void {
+        this.data?.destroy();
+        this.grad?.destroy();
+        this.data = null;
+        this.grad = null;
+    }
+}
+
+export function enableGrad(): void  { _gradEnabled = true;  }
+export function noGrad(): void      { _gradEnabled = false; }
+export function clearTape(): void   { _tape = []; }
+
+export function recordOperation(backwardFn: () => void | Promise<void>): number {
+    if (!_gradEnabled) return -1;
+    _tape.push({ backward: backwardFn });
+    return _tape.length - 1;
+}
+
+export async function backward(): Promise<void> {
+    for (let i = _tape.length - 1; i >= 0; i--) {
+        await _tape[i]!.backward();
+    }
+    clearTape();
+}
+
+export function ensureGradBuffer(device: GPUDevice, tensor: Tensor): void {
+    if (!tensor.grad) {
+        const STORAGE_USAGE: number = (_gpu.GPUBufferUsage?.STORAGE ?? 0x80) |
+                                      (_gpu.GPUBufferUsage?.COPY_DST ?? 0x08) |
+                                      (_gpu.GPUBufferUsage?.COPY_SRC ?? 0x04);
+        tensor.grad = device.createBuffer({
+            size  : tensor.byteSize,
+            usage : STORAGE_USAGE,
+        });
+        device.queue.writeBuffer(tensor.grad, 0, new Float32Array(tensor.numel));
+    }
+}
+
+export function allocateGradients(device: GPUDevice, tensors: Tensor[]): void {
+    for (const t of tensors) {
+        if (t.requiresGrad) ensureGradBuffer(device, t);
+    }
+}
+
+export function zeroGradients(device: GPUDevice, tensors: Tensor[]): void {
+    for (const t of tensors) {
+        if (t.grad) {
+            device.queue.writeBuffer(t.grad, 0, new Float32Array(t.numel));
+        }
+    }
+}
+
+export function onesLikeScalar(device: GPUDevice): GPUBuffer {
+    const USAGE: number = (_gpu.GPUBufferUsage?.STORAGE ?? 0x80) |
+                          (_gpu.GPUBufferUsage?.COPY_DST ?? 0x08);
+    const buf = device.createBuffer({
+        size  : 4,
+        usage : USAGE,
+        mappedAtCreation: true,
+    });
+    new Float32Array(buf.getMappedRange()).set([1.0]);
+    buf.unmap();
+    return buf;
+}
+
+export function crossEntropyLoss(logits: Float32Array, targetId: number): number {
+    let maxLogit = -Infinity;
+    for (let i = 0; i < logits.length; i++) {
+        if (logits[i]! > maxLogit) maxLogit = logits[i]!;
+    }
+    let sumExp = 0;
+    for (let i = 0; i < logits.length; i++) {
+        sumExp += Math.exp(logits[i]! - maxLogit);
+    }
+    const logSumExp = Math.log(sumExp) + maxLogit;
+    return logSumExp - logits[targetId]!;
+}
+
+export function crossEntropyGrad(logits: Float32Array, targetId: number): Float32Array {
+    let maxLogit = -Infinity;
+    for (let i = 0; i < logits.length; i++) {
+        if (logits[i]! > maxLogit) maxLogit = logits[i]!;
+    }
+    let sumExp = 0;
+    const exp_shifted = new Float32Array(logits.length);
+    for (let i = 0; i < logits.length; i++) {
+        exp_shifted[i] = Math.exp(logits[i]! - maxLogit);
+        sumExp += exp_shifted[i]!;
+    }
+    const probs = new Float32Array(logits.length);
+    for (let i = 0; i < logits.length; i++) {
+        probs[i] = exp_shifted[i]! / sumExp;
+    }
+    probs[targetId] = (probs[targetId] ?? 0) - 1.0;
+    return probs;
+}

package/src/training/trainer.ts

@@ -0,0 +1,309 @@
+/**
+ * trainer.ts – MambaTrainer class
+ */
+
+import {
+    createUniformBuffer,
+    createStorageBuffer,
+    createEmptyStorageBuffer,
+    createComputePipeline,
+    createBindGroup,
+    dispatchKernel,
+    cdiv,
+} from '../utils/gpu_utils.js';
+
+import { crossEntropyLoss, crossEntropyGrad } from './autograd.js';
+import { WEIGHT_UPDATE_WGSL, GRAD_CLIP_WGSL } from '../kernels/weight_update.js';
+import { HybridMambaModel, MambaModel } from '../model/mamba_model.js';
+import { BPETokenizer } from '../tokenizer/bpe.js';
+import type { LayerParam as BlockParam } from '../model/sequence_layer.js';
+
+export interface TrainOptions {
+  learningRate?: number;
+  epochs?: number;
+  batchSize?: number;
+  seqLen?: number;
+  maxGradNorm?: number;
+  weightDecay?: number;
+  beta1?: number;
+  beta2?: number;
+  eps?: number;
+  wsla?: boolean;
+  onEpochEnd?: ((epoch: number, loss: number) => void) | null;
+}
+
+interface AdamMoments {
+  m: GPUBuffer;
+  v: GPUBuffer;
+}
+
+interface AdamHyperparams {
+  learningRate: number;
+  weightDecay: number;
+  beta1: number;
+  beta2: number;
+  eps: number;
+  beta1_t: number;
+  beta2_t: number;
+}
+
+export class MambaTrainer {
+    model: HybridMambaModel;
+    tokenizer: BPETokenizer | null;
+    device: GPUDevice;
+    private _moments: AdamMoments[] | null;
+    private _step: number;
+    private _adamwPipeline: GPUComputePipeline;
+    private _clipReducePipeline: GPUComputePipeline;
+    private _clipScalePipeline: GPUComputePipeline;
+
+    constructor(model: HybridMambaModel | MambaModel, tokenizer: BPETokenizer | null = null) {
+        this.model     = model;
+        this.tokenizer = tokenizer;
+        this.device    = model.device;
+
+        this._moments = null;
+        this._step = 0;
+
+        this._adamwPipeline   = createComputePipeline(this.device, WEIGHT_UPDATE_WGSL, 'adamw_update');
+        this._clipReducePipeline = createComputePipeline(this.device, GRAD_CLIP_WGSL, 'grad_norm_reduce');
+        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),
+        }));
+    }
+
+    async train(input: string | number[], opts: TrainOptions = {}): Promise<number[]> {
+        const {
+            learningRate = 1e-4,
+            epochs       = 5,
+            batchSize    = 1,
+            seqLen       = 512,
+            maxGradNorm  = 1.0,
+            weightDecay  = 0.01,
+            beta1        = 0.9,
+            beta2        = 0.999,
+            eps          = 1e-8,
+            wsla         = false,
+            onEpochEnd   = null,
+        } = opts;
+
+        if (wsla) this.model.setWSLAMode(true);
+
+        let tokenIds: number[];
+        if (typeof input === 'string') {
+            if (!this.tokenizer) {
+                throw new Error(
+                    'MambaTrainer requires a tokenizer when input is a string. ' +
+                    'Pass a BPETokenizer instance as the second constructor argument.'
+                );
+            }
+            tokenIds = this.tokenizer.encode(input);
+        } else {
+            tokenIds = Array.from(input);
+        }
+
+        if (tokenIds.length < 2) {
+            throw new Error('Input must contain at least 2 tokens to form a training pair.');
+        }
+
+        const chunks = buildChunks(tokenIds, seqLen);
+        if (chunks.length === 0) {
+            throw new Error('Input is too short to form any training chunk.');
+        }
+
+        this._initMoments();
+
+        const epochLosses: number[] = [];
+
+        for (let epoch = 0; epoch < epochs; epoch++) {
+            let epochLoss = 0;
+            let numSteps  = 0;
+
+            for (const { inputs, targets } of chunks) {
+                const loss = await this._trainStep(
+                    inputs, targets, batchSize,
+                    { learningRate, maxGradNorm, weightDecay, beta1, beta2, eps, wsla }
+                );
+                epochLoss += loss;
+                numSteps++;
+            }
+
+            const avgLoss = epochLoss / numSteps;
+            epochLosses.push(avgLoss);
+
+            if (onEpochEnd) onEpochEnd(epoch + 1, avgLoss);
+        }
+
+        if (wsla) this.model.setWSLAMode(false);
+        return epochLosses;
+    }
+
+    private async _trainStep(
+        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;
+
+        this._step++;
+        const seqLen    = inputs.length;
+        const vocabSize = this.model.config.vocabSize;
+
+        const { logits, gpuLogits } = await this.model.forward(
+            new Uint32Array(inputs), batch, seqLen
+        );
+
+        let totalLoss = 0;
+        const dLogits = new Float32Array(batch * seqLen * vocabSize);
+
+        for (let i = 0; i < seqLen; i++) {
+            const offset = i * vocabSize;
+            const logitSlice = logits.slice(offset, offset + vocabSize);
+            const target = targets[i]!;
+            totalLoss += crossEntropyLoss(logitSlice, target);
+            const grad  = crossEntropyGrad(logitSlice, target);
+            for (let v = 0; v < vocabSize; v++) {
+                dLogits[offset + v] = grad[v]! / seqLen;
+            }
+        }
+        const loss = totalLoss / seqLen;
+
+        const dLogitsBuf = createStorageBuffer(this.device, dLogits, false);
+
+        await this._clipGradients(dLogitsBuf, dLogits.length, maxGradNorm);
+
+        const params  = this.model.parameters();
+        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 }
+        );
+
+        dLogitsBuf.destroy();
+        gpuLogits.destroy();
+
+        return loss;
+    }
+
+    private async _adamwStep(
+        params: BlockParam[],
+        gradBufs: GPUBuffer[],
+        hp: AdamHyperparams
+    ): Promise<void> {
+        const { learningRate, weightDecay, beta1, beta2, eps, beta1_t, beta2_t } = hp;
+
+        for (let i = 0; i < params.length; i++) {
+            const p       = params[i]!;
+            const gradBuf = gradBufs[Math.min(i, gradBufs.length - 1)]!;
+
+            if (!gradBuf || gradBuf.size < p.numel * 4) continue;
+
+            const paramsBuf = createUniformBuffer(this.device, packAdamParams(
+                p.numel, learningRate, beta1, beta2, eps, weightDecay, beta1_t, beta2_t
+            ));
+
+            const bg = createBindGroup(this.device, this._adamwPipeline, [
+                paramsBuf,
+                p.buf,
+                gradBuf,
+                this._moments![i]!.m,
+                this._moments![i]!.v,
+            ]);
+
+            dispatchKernel(this.device, this._adamwPipeline, bg,
+                [cdiv(p.numel, 256), 1, 1]);
+
+            paramsBuf.destroy();
+        }
+    }
+
+    private async _clipGradients(gradBuf: GPUBuffer, numel: number, maxNorm: number): Promise<void> {
+        const normSqBuf = createEmptyStorageBuffer(this.device, 4, true);
+        this.device.queue.writeBuffer(normSqBuf, 0, new Float32Array([0.0]));
+
+        const clipParams = new ArrayBuffer(8);
+        new Uint32Array(clipParams, 0, 1).set([numel]);
+        new Float32Array(clipParams, 4, 1).set([maxNorm * maxNorm]);
+        const pBuf = createUniformBuffer(this.device, clipParams);
+
+        const bg1 = createBindGroup(this.device, this._clipReducePipeline,
+            [pBuf, gradBuf, normSqBuf]);
+        dispatchKernel(this.device, this._clipReducePipeline, bg1,
+            [cdiv(numel, 256), 1, 1]);
+
+        const bg2 = createBindGroup(this.device, this._clipScalePipeline,
+            [pBuf, gradBuf, normSqBuf]);
+        dispatchKernel(this.device, this._clipScalePipeline, bg2,
+            [cdiv(numel, 256), 1, 1]);
+
+        pBuf.destroy();
+        normSqBuf.destroy();
+    }
+
+    async evaluate(input: string | number[]): Promise<number> {
+        let tokenIds: number[];
+        if (typeof input === 'string') {
+            if (!this.tokenizer) throw new Error('Tokenizer required for string input.');
+            tokenIds = this.tokenizer.encode(input);
+        } else {
+            tokenIds = Array.from(input);
+        }
+
+        const seqLen    = tokenIds.length;
+        const vocabSize = this.model.config.vocabSize;
+
+        const { logits } = await this.model.forward(
+            new Uint32Array(tokenIds.slice(0, -1)), 1, seqLen - 1
+        );
+
+        let totalLoss = 0;
+        for (let i = 0; i < seqLen - 1; i++) {
+            const offset = i * vocabSize;
+            totalLoss += crossEntropyLoss(
+                logits.slice(offset, offset + vocabSize),
+                tokenIds[i + 1]!
+            );
+        }
+
+        const avgLoss = totalLoss / (seqLen - 1);
+        return Math.exp(avgLoss);
+    }
+}
+
+function buildChunks(ids: number[], seqLen: number): Array<{inputs: number[], targets: number[]}> {
+    const chunks: Array<{inputs: number[], targets: number[]}> = [];
+    for (let start = 0; start + seqLen < ids.length; start += seqLen) {
+        chunks.push({
+            inputs : ids.slice(start, start + seqLen),
+            targets: ids.slice(start + 1, start + seqLen + 1),
+        });
+    }
+    const rem = ids.length % seqLen;
+    if (rem > 1) {
+        const start = ids.length - rem;
+        chunks.push({
+            inputs : ids.slice(start, -1),
+            targets: ids.slice(start + 1),
+        });
+    }
+    return chunks;
+}
+
+function packAdamParams(
+    numElements: number, lr: number, beta1: number, beta2: number,
+    eps: number, weightDecay: number, beta1_t: number, beta2_t: number
+): ArrayBuffer {
+    const buf = new ArrayBuffer(32);
+    new Uint32Array(buf, 0, 1).set([numElements]);
+    new Float32Array(buf, 4, 7).set([lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t]);
+    return buf;
+}