@seanhogg/builderforce-memory-engine 2026.7.12 → 2026.9.27

package/src/lm/index.ts

@@ -2,5 +2,15 @@
  * EvermindLM — the generative language model (the runnable "AI").
  */
 
-export { EvermindLM, EvermindLMTrainer, DEFAULT_LM_CONFIG, DEFAULT_LM_SEED } from "./evermind_lm.js";
-export type { EvermindLMConfig, LMGenerateOptions, TextCodec } from "./evermind_lm.js";
+export { EvermindLM, EvermindLMTrainer, DEFAULT_LM_CONFIG, DEFAULT_LM_SEED, logProbOfToken } from "./evermind_lm.js";
+export type { EvermindLMConfig, LMGenerateOptions, TextCodec, EvermindLMTrainOptions, EvermindLMDecodeState } from "./evermind_lm.js";
+
+// CPU text→vector embedder (semantic recall with no GPU).
+export { EvermindTextEmbedder } from "./text_embedder.js";
+export type { TextEmbedder, EmbedderCodec } from "./text_embedder.js";
+
+// PEFT / efficient-training toolkit (LoRA, QLoRA, mixed precision).
+export { LoRAAdapter, EvermindLMLoRA, quantizeBase } from "../training/lora.js";
+export type { LoRAConfig, LoRAFitOptions, BaseQuant } from "../training/lora.js";
+export { DynamicLossScaler, roundFp16, fp16View } from "../training/mixed_precision.js";
+export type { LossScalerOptions } from "../training/mixed_precision.js";

package/src/lm/text_embedder.ts

@@ -0,0 +1,84 @@
+/**
+ * lm/text_embedder.ts — text → vector, on the CPU, from an `EvermindLM`.
+ *
+ * `EvermindLM.embed()` turns TOKENS into a `dModel` vector. Every consumer that
+ * wants SEMANTIC RECALL wants it over TEXT, and needs the vector to be stable
+ * across processes, so it also needs (a) the tokenizer that produced those ids
+ * and (b) a way to tell whether a cached vector came from this exact model. This
+ * class is that pairing, in one place, so no consumer re-derives it.
+ *
+ * It deliberately satisfies the SAME structural shape the memory runtime's
+ * `embed(text)` seam expects (`MemoryStore.recallRanked(query, k, runtime)`
+ * accepts anything with an async `embed(text): Float32Array`), so a headless host
+ * can hand one of these straight to the memory layer and get embedding-ranked
+ * recall with NO GPU — the CPU LM is the whole dependency.
+ */
+
+import type { EvermindLM, TextCodec } from "./evermind_lm.js";
+import { crc32 } from "../utils/crc32.js";
+
+/** The minimal text-embedder seam the memory/recall layers consume. */
+export interface TextEmbedder {
+  /** Embed `text` as an L2-normalised vector of length {@link dimensions}. */
+  embed(text: string): Promise<Float32Array>;
+  /** Vector width (the model's `dModel`). */
+  readonly dimensions: number;
+  /**
+   * Identity of the model+tokenizer that produces these vectors. Vectors are only
+   * comparable within one fingerprint, so a persistent cache MUST discard entries
+   * stamped with a different one.
+   */
+  readonly fingerprint: string;
+}
+
+/** A tokenizer this embedder can encode through (the engine's `BPETokenizer` fits). */
+export type EmbedderCodec = Pick<TextCodec, "encode">;
+
+/**
+ * Pairs an `EvermindLM` with the tokenizer whose ids it was trained on.
+ *
+ * Pure compute and stateless: caching belongs to whoever owns a lifetime (the
+ * memory store's in-process LRU, or a host's on-disk vector cache), not here.
+ */
+export class EvermindTextEmbedder implements TextEmbedder {
+  readonly dimensions: number;
+  readonly fingerprint: string;
+
+  constructor(
+    private readonly model: EvermindLM,
+    private readonly codec: EmbedderCodec,
+    /** Override the derived fingerprint (e.g. with a published model version). */
+    fingerprint?: string,
+  ) {
+    this.dimensions = model.config.dModel;
+    this.fingerprint = fingerprint ?? derivedFingerprint(model);
+  }
+
+  /**
+   * Embed `text`. Async only to match the seam — the work is synchronous CPU, so
+   * a caller that wants the vector without a microtask can use
+   * {@link embedSync}.
+   */
+  async embed(text: string): Promise<Float32Array> {
+    return this.embedSync(text);
+  }
+
+  /** The synchronous form — encode with the codec, then run the CPU LM. */
+  embedSync(text: string): Float32Array {
+    return this.model.embedText(text, this.codec as TextCodec);
+  }
+}
+
+/**
+ * Model identity: architecture plus a CRC-32 over the tied embedding table.
+ *
+ * The embedding table is the one tensor every position's vector flows through, so
+ * an adapted checkpoint (which is what a delta merge produces) always changes it —
+ * making this a cheap but honest "are these vectors still comparable?" stamp.
+ */
+function derivedFingerprint(model: EvermindLM): string {
+  const c = model.config;
+  const bytes = new Uint8Array(model.emb.buffer, model.emb.byteOffset, model.emb.byteLength);
+  const arch = [c.vocabSize, c.dModel, c.numLayers, c.convKernel, c.hiddenDim, c.numExperts, c.topK].join("-");
+  return `evermind-lm/${arch}/${crc32(bytes).toString(16)}`;
+}

package/src/model/attention_block.ts

@@ -36,6 +36,7 @@ import {
 import { LINEAR_FORWARD_WGSL } from '../kernels/linear_projection.js';
 import { gaussianArray } from '../utils/rng.js';
 import { ACTIVATIONS_WGSL }    from '../kernels/activations.js';
+import { COL_SLICE_WGSL, COL_SLICE_ENTRY, dispatchColumnSlice } from '../kernels/slice.js';
 
 import type { SequenceLayer, LayerForwardResult, LayerParam } from './sequence_layer.js';
 
@@ -147,6 +148,7 @@ export class AttentionBlock implements SequenceLayer {
             attn_fwd: createComputePipeline(d, ATTENTION_FORWARD_WGSL,  'attention_forward'),
             attn_val: createComputePipeline(d, ATTENTION_FORWARD_WGSL,  'attention_value'),
             softmax : createComputePipeline(d, SOFTMAX_WGSL,            'softmax_forward'),
+            colSlice: createComputePipeline(d, COL_SLICE_WGSL,          COL_SLICE_ENTRY),
             elAdd   : createComputePipeline(d, ADD_SHADER,              'main'),
         };
 
@@ -194,11 +196,13 @@ export class AttentionBlock implements SequenceLayer {
         const KBuf = createEmptyStorageBuffer(d, M * dModel * 4, true);
         const VBuf = createEmptyStorageBuffer(d, M * dModel * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(qkvOut, 0,               QBuf, 0, M * dModel * 4);
-            enc.copyBufferToBuffer(qkvOut, M * dModel * 4,   KBuf, 0, M * dModel * 4);
-            enc.copyBufferToBuffer(qkvOut, 2 * M * dModel * 4, VBuf, 0, M * dModel * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split. `linear_forward` writes (M, 3*dModel) ROW-major, so Q, K
+            // and V are ranges of columns WITHIN every row — a strided gather, not a
+            // contiguous copy (which only coincides for M == 1).
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, qkvOut, QBuf, M, 3 * dModel, 0, dModel);
+            dispatchColumnSlice(d, slice, qkvOut, KBuf, M, 3 * dModel, dModel, dModel);
+            dispatchColumnSlice(d, slice, qkvOut, VBuf, M, 3 * dModel, 2 * dModel, dModel);
         }
         qkvOut.destroy();
 

package/src/model/mamba1_block.ts

@@ -20,6 +20,7 @@ import { gaussianArray } from '../utils/rng.js';
 import { CONV1D_FORWARD_WGSL }          from '../kernels/conv1d.js';
 import { LINEAR_FORWARD_WGSL }          from '../kernels/linear_projection.js';
 import { ACTIVATIONS_WGSL }             from '../kernels/activations.js';
+import { COL_SLICE_WGSL, COL_SLICE_ENTRY, dispatchColumnSlice } from '../kernels/slice.js';
 
 import type { SequenceLayer, LayerForwardResult, LayerParam } from './sequence_layer.js';
 
@@ -210,6 +211,7 @@ export class Mamba1Block implements SequenceLayer {
             rmsnorm     : createComputePipeline(d, ACTIVATIONS_WGSL,             'rmsnorm_forward'),
             scan_fwd    : createComputePipeline(d, SELECTIVE_SCAN_FORWARD_WGSL,  'forward_scan'),
             scan_reduce : createComputePipeline(d, SELECTIVE_SCAN_FORWARD_WGSL,  'forward_reduce'),
+            colSlice    : createComputePipeline(d, COL_SLICE_WGSL, COL_SLICE_ENTRY),
             elMul       : createComputePipeline(d, MUL_SHADER, 'main'),
             elAdd       : createComputePipeline(d, ADD_SHADER, 'main'),
         };
@@ -253,14 +255,18 @@ export class Mamba1Block implements SequenceLayer {
             dispatchKernel(d, this.pipelines['linear']!, bg, [cdiv(M, 16), cdiv(2 * D, 16), 1]);
         }
 
-        // 3. Split into x and z
+        // 3. Split into x and z.
+        //    in_proj writes (M, 2D) ROW-major, so x and z are interleaved COLUMN
+        //    slices — one per row — not two contiguous halves of the buffer. A
+        //    `copyBufferToBuffer` of the first M*D floats is only the right answer
+        //    when M == 1; for any real sequence it hands the block the first half of
+        //    the ROWS twice. Split with a strided gather instead.
         const xConvIn = createEmptyStorageBuffer(d, M * D * 4, true);
         const zBuf    = createEmptyStorageBuffer(d, M * D * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(inProjOut, 0,         xConvIn, 0, M * D * 4);
-            enc.copyBufferToBuffer(inProjOut, M * D * 4, zBuf,    0, M * D * 4);
-            d.queue.submit([enc.finish()]);
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, inProjOut, xConvIn, M, 2 * D, 0, D);
+            dispatchColumnSlice(d, slice, inProjOut, zBuf, M, 2 * D, D, D);
         }
         inProjOut.destroy();
         cache.zBuf    = zBuf;
@@ -270,7 +276,11 @@ export class Mamba1Block implements SequenceLayer {
         const convOut = createEmptyStorageBuffer(d, M * D * 4, true);
         cache.convOut = convOut;
         {
-            const params = new Uint32Array([L, D, dConv, B]).buffer;
+            // ConvParams is FIVE u32s (seq_len, d_channels, kernel_size, batch,
+            // groups). A four-element uniform is smaller than the struct's minimum
+            // binding size and WebGPU rejects the bind group outright. groups = 1
+            // is standard depthwise, matching Mamba-2/3's call.
+            const params = new Uint32Array([L, D, dConv, B, 1]).buffer;
             const pBuf   = createUniformBuffer(d, params);
             const bg = createBindGroup(d, this.pipelines['conv1d']!,
                 [pBuf, xConvIn, this.gpuWeights['wConv']!, this.gpuWeights['bConv']!, convOut]);
@@ -298,15 +308,17 @@ export class Mamba1Block implements SequenceLayer {
             dispatchKernel(d, this.pipelines['linear']!, bg, [cdiv(M, 16), cdiv(R + 2 * N, 16), 1]);
         }
 
+        //    Same story as the (x, z) split: Δ, B and C are COLUMN slices of the
+        //    (M, R+2N) projection, so they need a strided gather, not a prefix copy.
         const dtRaw = createEmptyStorageBuffer(d, M * R * 4, true);
         const B_raw = createEmptyStorageBuffer(d, B * L * N * 4, true);
         const C_raw = createEmptyStorageBuffer(d, B * L * N * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(xProjOut, 0,               dtRaw, 0, M * R * 4);
-            enc.copyBufferToBuffer(xProjOut, M * R * 4,       B_raw, 0, B * L * N * 4);
-            enc.copyBufferToBuffer(xProjOut, M * (R + N) * 4, C_raw, 0, B * L * N * 4);
-            d.queue.submit([enc.finish()]);
+            const P = R + 2 * N;
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, xProjOut, dtRaw, M, P, 0, R);
+            dispatchColumnSlice(d, slice, xProjOut, B_raw, M, P, R, N);
+            dispatchColumnSlice(d, slice, xProjOut, C_raw, M, P, R + N, N);
         }
         xProjOut.destroy();
         cache.B_raw = B_raw;
@@ -335,7 +347,11 @@ export class Mamba1Block implements SequenceLayer {
             const bg1 = createBindGroup(d, this.pipelines['scan_fwd']!,
                 [pBuf, siluOut, deltaFull, this.gpuWeights['A_log']!, B_raw, C_raw,
                  this.gpuWeights['D_vec']!, scanY, hCache]);
-            dispatchKernel(d, this.pipelines['scan_fwd']!, bg1, [cdiv(D, 8), cdiv(N, 8), B]);
+            // One workgroup per (d, n, batch): the kernel reads `wgid.x` as d and
+            // `wgid.y` as n, so the grid must be D x N x B. Dispatching
+            // ceil(D/8) x ceil(N/8) left all but a 1/64th corner of the state
+            // matrix unscanned (and unwritten in h_cache).
+            dispatchKernel(d, this.pipelines['scan_fwd']!, bg1, [D, N, B]);
 
             const bg2 = createBindGroup(d, this.pipelines['scan_reduce']!,
                 [pBuf, siluOut, deltaFull, this.gpuWeights['A_log']!, B_raw, C_raw,

package/src/model/mamba2_block.ts

@@ -24,6 +24,7 @@ import {
 import { SSD_FORWARD_WGSL }    from '../kernels/ssd.js';
 import { gaussianArray } from '../utils/rng.js';
 import { CONV1D_FORWARD_WGSL } from '../kernels/conv1d.js';
+import { COL_SLICE_WGSL, COL_SLICE_ENTRY, dispatchColumnSlice } from '../kernels/slice.js';
 import { LINEAR_FORWARD_WGSL } from '../kernels/linear_projection.js';
 import { ACTIVATIONS_WGSL }    from '../kernels/activations.js';
 
@@ -127,6 +128,7 @@ export class Mamba2Block implements SequenceLayer {
         this.pipelines = {
             linear   : createComputePipeline(d, LINEAR_FORWARD_WGSL,  'linear_forward'),
             conv1d   : createComputePipeline(d, CONV1D_FORWARD_WGSL,  'conv1d_forward'),
+            colSlice   : createComputePipeline(d, COL_SLICE_WGSL, COL_SLICE_ENTRY),
             rmsnorm  : createComputePipeline(d, ACTIVATIONS_WGSL,     'rmsnorm_forward'),
             ssd_fwd  : createComputePipeline(d, SSD_FORWARD_WGSL,     'ssd_chunk_forward'),
             elAdd    : createComputePipeline(d, ADD_SHADER,           'main'),
@@ -181,10 +183,13 @@ export class Mamba2Block implements SequenceLayer {
         const xConvBuf = createEmptyStorageBuffer(d, M * convD * 4, true);
         const dtBuf    = createEmptyStorageBuffer(d, M * H * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(inProjOut, 0,              xConvBuf, 0, M * convD * 4);
-            enc.copyBufferToBuffer(inProjOut, M * convD * 4,  dtBuf,    0, M * H * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split. `linear_forward` writes (M, convD + H) ROW-major, so each
+            // part is a range of columns WITHIN every row — a strided gather, not a
+            // contiguous copy (which only coincides for M == 1).
+            const slice = this.pipelines['colSlice']!;
+            const stride = convD + H;
+            dispatchColumnSlice(d, slice, inProjOut, xConvBuf, M, stride, 0, convD);
+            dispatchColumnSlice(d, slice, inProjOut, dtBuf, M, stride, convD, H);
         }
         inProjOut.destroy();
 
@@ -204,11 +209,11 @@ export class Mamba2Block implements SequenceLayer {
         const bProjBuf = createEmptyStorageBuffer(d, M * G * N * 4, true);
         const cProjBuf = createEmptyStorageBuffer(d, M * G * N * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(convOut, 0,            xSsdBuf,  0, M * D * 4);
-            enc.copyBufferToBuffer(convOut, M * D * 4,    bProjBuf, 0, M * G * N * 4);
-            enc.copyBufferToBuffer(convOut, M * (D + G * N) * 4, cProjBuf, 0, M * G * N * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split of the (M, convD) conv output — same reasoning as above.
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, convOut, xSsdBuf, M, convD, 0, D);
+            dispatchColumnSlice(d, slice, convOut, bProjBuf, M, convD, D, G * N);
+            dispatchColumnSlice(d, slice, convOut, cProjBuf, M, convD, D + G * N, G * N);
         }
         convOut.destroy();
 

package/src/model/mamba3_block.ts

@@ -33,6 +33,7 @@ import {
 import { COMPLEX_SSD_FORWARD_WGSL } from '../kernels/complex_ssd.js';
 import { gaussianArray } from '../utils/rng.js';
 import { CONV1D_FORWARD_WGSL }      from '../kernels/conv1d.js';
+import { COL_SLICE_WGSL, COL_SLICE_ENTRY, dispatchColumnSlice } from '../kernels/slice.js';
 import { LINEAR_FORWARD_WGSL }      from '../kernels/linear_projection.js';
 import { ACTIVATIONS_WGSL }         from '../kernels/activations.js';
 
@@ -145,6 +146,7 @@ export class Mamba3Block implements SequenceLayer {
         this.pipelines = {
             linear     : createComputePipeline(d, LINEAR_FORWARD_WGSL,       'linear_forward'),
             conv1d     : createComputePipeline(d, CONV1D_FORWARD_WGSL,       'conv1d_forward'),
+            colSlice   : createComputePipeline(d, COL_SLICE_WGSL, COL_SLICE_ENTRY),
             rmsnorm    : createComputePipeline(d, ACTIVATIONS_WGSL,          'rmsnorm_forward'),
             cssd_fwd   : createComputePipeline(d, COMPLEX_SSD_FORWARD_WGSL,  'complex_ssd_forward'),
             elAdd      : createComputePipeline(d, ADD_SHADER,                'main'),
@@ -198,10 +200,13 @@ export class Mamba3Block implements SequenceLayer {
         const xConvBuf = createEmptyStorageBuffer(d, M * convD * 4, true);
         const dtBuf    = createEmptyStorageBuffer(d, M * H * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(inProjOut, 0,             xConvBuf, 0, M * convD * 4);
-            enc.copyBufferToBuffer(inProjOut, M * convD * 4, dtBuf,    0, M * H * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split. `linear_forward` writes (M, convD + H) ROW-major, so each
+            // part is a range of columns WITHIN every row — a strided gather, not a
+            // contiguous copy (which only coincides for M == 1).
+            const slice = this.pipelines['colSlice']!;
+            const stride = convD + H;
+            dispatchColumnSlice(d, slice, inProjOut, xConvBuf, M, stride, 0, convD);
+            dispatchColumnSlice(d, slice, inProjOut, dtBuf, M, stride, convD, H);
         }
         inProjOut.destroy();
 
@@ -221,11 +226,11 @@ export class Mamba3Block implements SequenceLayer {
         const bProjBuf = createEmptyStorageBuffer(d, M * G * Nc * 2 * 4, true);
         const cProjBuf = createEmptyStorageBuffer(d, M * G * Nc * 2 * 4, true);
         {
-            const enc = d.createCommandEncoder();
-            enc.copyBufferToBuffer(convOut, 0,                       xSsdBuf,  0, M * D * 4);
-            enc.copyBufferToBuffer(convOut, M * D * 4,               bProjBuf, 0, M * G * Nc * 2 * 4);
-            enc.copyBufferToBuffer(convOut, M * (D + G * Nc * 2) * 4, cProjBuf, 0, M * G * Nc * 2 * 4);
-            d.queue.submit([enc.finish()]);
+            // COLUMN split of the (M, convD) conv output — same reasoning as above.
+            const slice = this.pipelines['colSlice']!;
+            dispatchColumnSlice(d, slice, convOut, xSsdBuf, M, convD, 0, D);
+            dispatchColumnSlice(d, slice, convOut, bProjBuf, M, convD, D, G * Nc * 2);
+            dispatchColumnSlice(d, slice, convOut, cProjBuf, M, convD, D + G * Nc * 2, G * Nc * 2);
         }
         convOut.destroy();
 

package/src/moe/index.ts

@@ -13,12 +13,3 @@ export type { MoEConfig, MoEParam, RouteResult } from "./moe_model.js";
 export { MoETrainer } from "./moe_trainer.js";
 export type { MoESample, MoETrainOptions, MoEEpochResult } from "./moe_trainer.js";
 
-export { EvermindModelPackage } from "./moe_package.js";
-export type {
-  EvermindModelManifest,
-  EvermindModelCard,
-  EvermindModelType,
-  EvermindModality,
-  PackageMeta,
-  ValidationResult,
-} from "./moe_package.js";

package/src/optim/adamw.ts

@@ -15,19 +15,42 @@ export interface OptimTarget {
   gradients(): OptimParam[];
 }
 
+/**
+ * Optimizer-state sharding (ZeRO-1 analog). When set, this optimiser instance
+ * OWNS only parameter tensors where `paramIndex % count === index`, and only
+ * allocates the AdamW moments (m, v — 2× the parameter bytes, the dominant cost
+ * of full fine-tuning per the cookbook's memory accounting) for its shard.
+ *
+ * Run `count` instances over the SAME model, one per `index`, and the union of
+ * their steps equals one unsharded step — but each holds only 1/count of the
+ * optimizer state. That is exactly how FSDP/ZeRO fits a larger model on the same
+ * hardware; here it is the seam a multi-device trainer partitions on, and even
+ * single-process it caps resident moment memory.
+ */
+export interface ShardSpec {
+  /** This shard's rank, 0-based. */
+  index: number;
+  /** Total number of shards. */
+  count: number;
+}
+
 export interface AdamWOptions {
   lr?: number;
   beta1?: number;
   beta2?: number;
   eps?: number;
   weightDecay?: number;
+  /** Optimizer-state sharding. Omit for the full (single-owner) optimiser. */
+  shard?: ShardSpec;
 }
 
 export class AdamW {
-  private readonly m: Float32Array[] = [];
-  private readonly v: Float32Array[] = [];
+  // Per-parameter moments, or null for tensors this shard does not own (no memory).
+  private readonly m: (Float32Array | null)[] = [];
+  private readonly v: (Float32Array | null)[] = [];
   private t = 0;
-  private readonly opt: Required<AdamWOptions>;
+  private readonly opt: Required<Omit<AdamWOptions, "shard">>;
+  private readonly shard: ShardSpec | null;
 
   constructor(
     private readonly target: OptimTarget,
@@ -40,13 +63,29 @@ export class AdamW {
       eps: options.eps ?? 1e-8,
       weightDecay: options.weightDecay ?? 0,
     };
-    for (const p of target.parameters()) {
-      this.m.push(new Float32Array(p.data.length));
-      this.v.push(new Float32Array(p.data.length));
+    const shard = options.shard ?? null;
+    if (shard && (shard.count <= 0 || shard.index < 0 || shard.index >= shard.count)) {
+      throw new Error(`AdamW: invalid shard ${shard.index}/${shard.count}`);
     }
+    this.shard = shard;
+    const params = target.parameters();
+    for (let p = 0; p < params.length; p++) {
+      if (this._owns(p)) {
+        this.m.push(new Float32Array(params[p]!.data.length));
+        this.v.push(new Float32Array(params[p]!.data.length));
+      } else {
+        this.m.push(null);
+        this.v.push(null);
+      }
+    }
+  }
+
+  /** Whether this shard owns (and updates) parameter tensor `p`. */
+  private _owns(p: number): boolean {
+    return this.shard === null || p % this.shard.count === this.shard.index;
   }
 
-  /** One optimiser step from the currently-accumulated gradients. */
+  /** One optimiser step from the currently-accumulated gradients (owned tensors only). */
   step(): void {
     this.t++;
     const { lr, beta1, beta2, eps, weightDecay } = this.opt;
@@ -55,6 +94,7 @@ export class AdamW {
     const bc1 = 1 - Math.pow(beta1, this.t);
     const bc2 = 1 - Math.pow(beta2, this.t);
     for (let p = 0; p < params.length; p++) {
+      if (!this._owns(p)) continue;
       const w = params[p]!.data;
       const g = grads[p]!.data;
       const m = this.m[p]!;
@@ -69,4 +109,72 @@ export class AdamW {
       }
     }
   }
+
+  /** Bytes of optimizer state this shard holds (m + v) — 1/count of the full state. */
+  stateBytes(): number {
+    let n = 0;
+    for (const arr of this.m) if (arr) n += arr.byteLength;
+    for (const arr of this.v) if (arr) n += arr.byteLength;
+    return n;
+  }
+}
+
+// ── The CPU mirror of WEIGHT_UPDATE_WGSL ─────────────────────────────────────
+
+/** Hyperparameters of one {@link adamwUpdateInPlace} step. */
+export interface AdamWKernelStep {
+  lr: number;
+  beta1: number;
+  beta2: number;
+  eps: number;
+  weightDecay: number;
+  /** beta1^t — the precomputed bias-correction term (matches the kernel uniform). */
+  beta1_t: number;
+  /** beta2^t. */
+  beta2_t: number;
+  /** Trust region: max |Δθ| per step. 0 disables the bound. */
+  maxDelta: number;
+  /** Global gradient-clip factor, applied to `grad` on the fly (1 = no clipping). */
+  gradScale?: number;
+}
+
+/**
+ * One AdamW step over a single parameter tensor — the exact CPU mirror of the
+ * `adamw_update` entry point in `WEIGHT_UPDATE_WGSL`, including its two safety
+ * properties:
+ *
+ *   • the **NaN/Inf guard** — a non-finite step is dropped rather than written
+ *     into a weight, so one bad gradient cannot permanently poison the model;
+ *   • the **trust region** — `maxDelta` bounds how far a single step can move a
+ *     weight, which is what keeps repeated write-through adapts reversible.
+ *
+ * Decoupled weight decay (`θ * (1 - lr·wd)`), matching the kernel — deliberately
+ * NOT the coupled form {@link AdamW} uses for the MoE/EvermindLM trainers.
+ *
+ * `param`, `m` and `v` are updated in place.
+ */
+export function adamwUpdateInPlace(
+  param: Float32Array,
+  grad: Float32Array,
+  m: Float32Array,
+  v: Float32Array,
+  hp: AdamWKernelStep,
+): void {
+  const { lr, beta1, beta2, eps, weightDecay, beta1_t, beta2_t, maxDelta } = hp;
+  const gradScale = hp.gradScale ?? 1;
+  const bc1 = 1 - beta1_t;
+  const bc2 = 1 - beta2_t;
+  for (let i = 0; i < param.length; i++) {
+    const g = grad[i]! * gradScale;
+    const mNew = beta1 * m[i]! + (1 - beta1) * g;
+    const vNew = beta2 * v[i]! + (1 - beta2) * g * g;
+    m[i] = mNew;
+    v[i] = vNew;
+
+    let step = (lr * (mNew / bc1)) / (Math.sqrt(vNew / bc2) + eps);
+    if (!Number.isFinite(step)) step = 0;
+    if (maxDelta > 0) step = Math.min(maxDelta, Math.max(-maxDelta, step));
+
+    param[i] = param[i]! * (1 - lr * weightDecay) - step;
+  }
 }

package/src/tokenizer/bpe.ts

@@ -52,6 +52,14 @@ 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;
 
+/** A round-trippable snapshot of a {@link BPETokenizer} — see `toObject`/`deserialize`. */
+export interface BPETokenizerSpec {
+    vocab: Record<string, number>;
+    /** Merge rules in rank order ("a b" strings). */
+    merges: string[];
+    specials?: { bos?: string; eos?: string; pad?: string; unk?: string };
+}
+
 export class BPETokenizer {
     vocab: Map<string, number>;
     idToToken: Map<number, string>;
@@ -225,6 +233,53 @@ export class BPETokenizer {
         this.loadFromObjects(vocabObj, mergeArr);
     }
 
+    /**
+     * The tokenizer as a plain, JSON-serialisable object: vocab, merges and the
+     * four special-token strings.
+     *
+     * This is what lets an `.evermind` package EMBED its tokenizer instead of
+     * expecting a consumer to source the matching vocab separately — token ids are
+     * meaningless without the exact vocabulary that produced them, so shipping a
+     * checkpoint without one is shipping half an artifact.
+     */
+    toObject(): BPETokenizerSpec {
+        return {
+            vocab: Object.fromEntries(this.vocab),
+            merges: [...this.merges.entries()].sort((a, b) => a[1] - b[1]).map(([m]) => m),
+            specials: { bos: this.bosToken, eos: this.eosToken, pad: this.padToken, unk: this.unkToken },
+        };
+    }
+
+    /** Restore a tokenizer from {@link toObject}. */
+    loadFromSpec(spec: BPETokenizerSpec): void {
+        const sp = spec.specials;
+        if (sp?.bos) this.bosToken = sp.bos;
+        if (sp?.eos) this.eosToken = sp.eos;
+        if (sp?.pad) this.padToken = sp.pad;
+        if (sp?.unk) this.unkToken = sp.unk;
+        this.loadFromObjects(spec.vocab, spec.merges);
+    }
+
+    /** UTF-8 JSON bytes of {@link toObject} — the `.evermind` tokenizer section. */
+    serialize(): ArrayBuffer {
+        const bytes = new TextEncoder().encode(JSON.stringify(this.toObject()));
+        // Copy into a standalone ArrayBuffer (the view may be a slice of a pool).
+        const out = new ArrayBuffer(bytes.byteLength);
+        new Uint8Array(out).set(bytes);
+        return out;
+    }
+
+    /** Inverse of {@link serialize}. */
+    static deserialize(buffer: ArrayBuffer): BPETokenizer {
+        const spec = JSON.parse(new TextDecoder().decode(new Uint8Array(buffer))) as BPETokenizerSpec;
+        if (!spec || typeof spec.vocab !== 'object' || !Array.isArray(spec.merges)) {
+            throw new Error('BPETokenizer.deserialize: not a serialised BPE tokenizer (expected { vocab, merges })');
+        }
+        const tok = new BPETokenizer();
+        tok.loadFromSpec(spec);
+        return tok;
+    }
+
     /**
      * Fetch and import a Hugging Face `tokenizer.json` by URL — the IDE's
      * "import merges from a repo" path. Resolves once the tokenizer is loaded.