@seanhogg/builderforce-memory-engine 2026.9.29 → 2026.9.31

package/src/import/foreign/infer.ts

@@ -0,0 +1,62 @@
+/**
+ * import/foreign/infer.ts — a Mamba-1 checkpoint's config, read off its tensors.
+ *
+ * `state-spaces/mamba-*` repositories publish `model.safetensors` with a
+ * `config.json` in the original Mamba-SSM vocabulary (`d_model`, `n_layer`, …)
+ * that no adapter reads, and a dropped safetensors file often arrives with no
+ * config at all. Every geometry fact the port needs is in the tensor shapes, so
+ * this derives the Hugging Face `MambaConfig` the {@link mambaAdapter} expects
+ * instead of anyone restating the derivation.
+ */
+
+import type { NamedTensor } from "../../export/tensors.js";
+import type { ForeignConfig } from "./adapter.js";
+import { normaliseSourceName } from "./plan.js";
+
+function fail(detail: string): never {
+  throw new Error(`import/foreign: cannot read a Mamba-1 config from these tensors — ${detail}`);
+}
+
+/**
+ * The `MambaForCausalLM` config implied by a Mamba-1 checkpoint's tensors:
+ * embedding `[vocab, d_model]`, per-layer `in_proj [2·d_inner, d_model]`,
+ * `conv1d [d_inner, (1,) d_conv]`, `A_log [d_inner, d_state]` and
+ * `x_proj [dt_rank + 2·d_state, d_inner]`.
+ */
+export function mambaConfigFromTensors(tensors: readonly NamedTensor[]): ForeignConfig {
+  const byName = new Map(tensors.map((t) => [normaliseSourceName(t.name), t]));
+  const shape = (name: string) => byName.get(name)?.shape;
+
+  const emb = shape("backbone.embeddings.weight") ?? shape("backbone.embedding.weight");
+  if (!emb || emb.length !== 2) fail("no backbone.embedding(s).weight [vocab, d_model]");
+  const [vocabSize, dModel] = emb as [number, number];
+
+  let numLayers = 0;
+  while (byName.has(`backbone.layers.${numLayers}.mixer.in_proj.weight`)) numLayers++;
+  if (numLayers === 0) fail("no backbone.layers.N.mixer.in_proj.weight");
+
+  const m = "backbone.layers.0.mixer";
+  const inProj = shape(`${m}.in_proj.weight`)!;
+  const conv = shape(`${m}.conv1d.weight`);
+  const aLog = shape(`${m}.A_log`);
+  const xProj = shape(`${m}.x_proj.weight`);
+  if (!conv || !aLog || !xProj) fail(`layer 0 is missing conv1d.weight, A_log or x_proj.weight`);
+
+  const dInner = inProj[0]! / 2;
+  const dState = aLog[1]!;
+  const dConv = conv[conv.length - 1]!;
+  const dtRank = xProj[0]! - 2 * dState;
+  if (!Number.isInteger(dInner) || dtRank <= 0) fail(`inconsistent shapes (in_proj ${inProj.join("×")}, x_proj ${xProj.join("×")})`);
+
+  return {
+    model_type: "mamba",
+    architectures: ["MambaForCausalLM"],
+    vocab_size: vocabSize,
+    hidden_size: dModel,
+    num_hidden_layers: numLayers,
+    intermediate_size: dInner,
+    state_size: dState,
+    conv_kernel: dConv,
+    time_step_rank: dtRank,
+  };
+}

package/src/import/foreign/registry.ts

@@ -12,12 +12,13 @@ import type { NamedTensor } from "../../export/tensors.js";
 import type { HybridMambaModelConfig } from "../../model/mamba_model.js";
 import { safetensorsToTensors } from "../safetensors.js";
 import { executePortPlan, type PortDiscard, type PortedTensors, type SynthesisedTarget } from "./plan.js";
-import { falconMambaAdapter } from "./falcon_mamba.js";
+import { falconMambaAdapter, mambaAdapter } from "./falcon_mamba.js";
 import { codestralMambaAdapter } from "./codestral_mamba.js";
 import type { ForeignConfig, ForeignMambaAdapter } from "./adapter.js";
 
 /** Every foreign SSM architecture this engine can warm-start from. */
 export const FOREIGN_MAMBA_ADAPTERS: readonly ForeignMambaAdapter[] = [
+  mambaAdapter,
   falconMambaAdapter,
   codestralMambaAdapter,
 ];

package/src/import/foreign/vocab.ts

@@ -0,0 +1,50 @@
+/**
+ * import/foreign/vocab.ts — fit a ported checkpoint to a different vocabulary size.
+ *
+ * A ported model keeps its source tokenizer's vocabulary. To pair it with another
+ * tokenizer the embedding table (tied to the LM head) must have that tokenizer's row
+ * count: surplus rows are dropped, missing rows are drawn small and random so the
+ * new tokens start near the embedding cloud's centre. The new rows are recorded on
+ * the port as synthesised, never passed off as ported.
+ */
+
+import { SeededRng, gaussianArray } from "../../utils/rng.js";
+import type { PortedCheckpoint } from "./registry.js";
+
+export interface ResizeVocabOptions {
+  /** Standard deviation of new rows. Default 0.02, the usual embedding init. */
+  std?: number;
+  /** Seed for the new rows, so the same resize produces the same checkpoint. Default 1. */
+  seed?: number;
+}
+
+/** `ported` with an embedding of `vocabSize` rows and a model config to match. */
+export function resizePortedVocab(ported: PortedCheckpoint, vocabSize: number, opts: ResizeVocabOptions = {}): PortedCheckpoint {
+  if (!Number.isInteger(vocabSize) || vocabSize <= 0) {
+    throw new Error(`import/foreign: vocabulary size must be a positive integer, got ${vocabSize}`);
+  }
+  const from = ported.modelConfig.vocabSize;
+  if (vocabSize === from) return ported;
+
+  const emb = ported.weights.get("embedding");
+  if (!emb) throw new Error('import/foreign: the port has no "embedding" to resize');
+  const dModel = ported.modelConfig.dModel;
+
+  const resized = new Float32Array(vocabSize * dModel);
+  resized.set(emb.subarray(0, Math.min(from, vocabSize) * dModel));
+  const synthesisedTargets = [...ported.synthesisedTargets];
+  if (vocabSize > from) {
+    const std = opts.std ?? 0.02;
+    const rng = new SeededRng(opts.seed ?? 1);
+    resized.set(gaussianArray((vocabSize - from) * dModel, std, () => rng.next()), from * dModel);
+    synthesisedTargets.push({
+      target: "embedding",
+      value: std,
+      reason: `rows ${from}..${vocabSize - 1} drawn from N(0, ${std}²) to reach a ${vocabSize}-token vocabulary`,
+    });
+  }
+
+  const weights = new Map(ported.weights);
+  weights.set("embedding", resized);
+  return { ...ported, weights, synthesisedTargets, modelConfig: { ...ported.modelConfig, vocabSize } };
+}

package/src/import/index.ts

@@ -25,7 +25,11 @@ export {
   executePortPlan,
   normaliseSourceName,
   falconMambaAdapter,
+  mambaAdapter,
+  mamba1Adapter,
   codestralMambaAdapter,
+  mambaConfigFromTensors,
+  resizePortedVocab,
 } from "./foreign/index.js";
 export type {
   ForeignConfig,
@@ -38,4 +42,5 @@ export type {
   PortedCheckpoint,
   PortedTensors,
   SynthesisedTarget,
+  ResizeVocabOptions,
 } from "./foreign/index.js";

package/src/import/safetensors.ts

@@ -4,7 +4,7 @@
  * The exact inverse of {@link ../export/safetensors.ts}: parse the 8-byte
  * little-endian u64 header length, the JSON header (name → {dtype, shape,
  * data_offsets}), then decode each tensor's raw bytes into a Float32Array.
- * F32 and F16 are supported (F16 is widened to f32 on read). This is the reader
+ * F32, F16 and BF16 are supported (both halves are widened to f32 on read). This is the reader
  * half of the warm-start / weight-port path — export was one-directional before.
  */
 
@@ -63,5 +63,12 @@ function decodeTensor(seg: Uint8Array, dtype: string, numel: number, name: strin
     for (let i = 0; i < numel; i++) data[i] = fp16ToFloat(dv.getUint16(i * 2, true));
     return data;
   }
-  throw new Error(`safetensors: unsupported dtype "${dtype}" for "${name}" (expected F32 or F16)`);
+  if (d === "BF16" || d === "BFLOAT16") {
+    // bfloat16 is the top half of an f32: shift it back into place.
+    if (seg.length < numel * 2) throw new Error(`safetensors: "${name}" BF16 bytes ${seg.length} < ${numel * 2}`);
+    const bits = new Uint32Array(data.buffer);
+    for (let i = 0; i < numel; i++) bits[i] = dv.getUint16(i * 2, true) << 16;
+    return data;
+  }
+  throw new Error(`safetensors: unsupported dtype "${dtype}" for "${name}" (expected F32, F16 or BF16)`);
 }

package/src/index.ts

@@ -181,7 +181,7 @@ export {
 } from './import/index.js';
 export type { ImportOptions } from './import/index.js';
 
-// ── Foreign weight port (Falcon-Mamba / Codestral-Mamba → HybridMambaModel) ────
+// ── Foreign weight port (Mamba / Falcon-Mamba / Codestral-Mamba → HybridMambaModel) ──
 // Transformer checkpoints are NOT portable — distillation is the only route, and
 // `foreignMambaAdapterFor` rejects them saying so.
 
@@ -194,7 +194,11 @@ export {
     executePortPlan,
     normaliseSourceName,
     falconMambaAdapter,
+    mambaAdapter,
+    mamba1Adapter,
     codestralMambaAdapter,
+    mambaConfigFromTensors,
+    resizePortedVocab,
 } from './import/index.js';
 export type {
     ForeignConfig,
@@ -207,6 +211,7 @@ export type {
     PortedCheckpoint,
     PortedTensors,
     SynthesisedTarget,
+    ResizeVocabOptions,
 } from './import/index.js';
 
 // ── Benchmarking (held-out perplexity / accuracy / throughput + A/B) ──────────

package/src/moe/moe_model.ts

@@ -188,18 +188,20 @@ export class LoadBalanceAccumulator {
     for (const e of route.experts) this.counts[e] = this.counts[e]! + 1;
     for (let e = 0; e < this.numExperts; e++) this.probSum[e] = this.probSum[e]! + route.probs[e]!;
   }
+  /** `f_e`: the fraction of all dispatches that went to each expert (sums to 1 once anything is observed). */
+  dispatchFractions(): Float32Array {
+    const dispatched = this.counts.reduce((a, b) => a + b, 0) || 1; // = tokens·topK
+    return Float32Array.from(this.counts, (c) => c / dispatched);
+  }
   /** The load-balance loss over everything observed so far (0 if no tokens). */
   loss(): number {
     if (this.tokens === 0) return 0;
-    const E = this.numExperts;
-    const dispatched = this.counts.reduce((a, b) => a + b, 0) || 1; // = tokens·topK
+    const f = this.dispatchFractions();
     let sum = 0;
-    for (let e = 0; e < E; e++) {
-      const f = this.counts[e]! / dispatched; // fraction of dispatches to e
-      const p = this.probSum[e]! / this.tokens; // mean router prob for e
-      sum += f * p;
+    for (let e = 0; e < this.numExperts; e++) {
+      sum += f[e]! * (this.probSum[e]! / this.tokens); // f_e · mean router prob for e
     }
-    return E * sum;
+    return this.numExperts * sum;
   }
 }
 

package/src/moe/moe_trainer.ts

@@ -8,7 +8,7 @@
  * WebGPU optimiser kernel would accelerate.
  */
 
-import { SharedExpertMoE } from "./moe_model.js";
+import { LoadBalanceAccumulator, SharedExpertMoE } from "./moe_model.js";
 import { AdamW } from "../optim/adamw.js";
 
 export interface MoESample {
@@ -87,7 +87,7 @@ export class MoETrainer {
       // Forward + task backward, retaining (x, probs) for the batch aux gradient.
       const xs: Float32Array[] = [];
       const probsList: Float32Array[] = [];
-      const counts = new Float32Array(numExperts);
+      const balance = new LoadBalanceAccumulator(numExperts);
       let batchLoss = 0;
 
       for (const s of batch) {
@@ -101,12 +101,11 @@ export class MoETrainer {
         this.model.backward(dOut, f.cache);
         xs.push(f.cache.x);
         probsList.push(f.route.probs);
-        for (const ex of f.route.experts) counts[ex] = counts[ex]! + 1;
+        balance.observe(f.route);
       }
 
       // Load-balancing aux gradient (batch-level): f = dispatch fractions.
-      const dispatched = counts.reduce((a, b) => a + b, 0) || 1;
-      const fVec = Float32Array.from(counts, (c) => c / dispatched);
+      const fVec = balance.dispatchFractions();
       const scale = (this.opt.auxWeight * numExperts) / batch.length;
       for (let i = 0; i < xs.length; i++) {
         this.model.auxGradStep(xs[i]!, probsList[i]!, fVec, scale);
@@ -117,9 +116,7 @@ export class MoETrainer {
       this.adam.step();
 
       epochLoss += batchLoss;
-      lastAux = numExperts * fVec.reduce((sum, f, e) => sum + f * (probsList.length
-        ? probsList.reduce((s, p) => s + p[e]!, 0) / probsList.length
-        : 0), 0);
+      lastAux = balance.loss();
     }
 
     return { loss: epochLoss / Math.max(1, samples.length), auxLoss: lastAux };

package/src/utils/rng.ts

@@ -50,16 +50,20 @@ export function setInitSeed(seed: number | undefined): void {
     }
 }
 
-/** Box–Muller Gaussian sample from the active source. */
-export function randn(std = 1): number {
-    const u1 = Math.max(_next(), 1e-12);
-    const u2 = _next();
+/**
+ * Box–Muller Gaussian sample. Draws from the active init source unless `next` names
+ * another uniform [0,1) source — a caller-owned {@link SeededRng}, say, when the draw
+ * must be reproducible without touching the process-wide seed.
+ */
+export function randn(std = 1, next: () => number = _next): number {
+    const u1 = Math.max(next(), 1e-12);
+    const u2 = next();
     return std * Math.sqrt(-2 * Math.log(u1)) * Math.cos(2 * Math.PI * u2);
 }
 
 /** Returns a Float32Array of `n` Gaussian samples with the given standard deviation. */
-export function gaussianArray(n: number, std: number): Float32Array {
+export function gaussianArray(n: number, std: number, next?: () => number): Float32Array {
     const a = new Float32Array(n);
-    for (let i = 0; i < n; i++) a[i] = randn(std);
+    for (let i = 0; i < n; i++) a[i] = randn(std, next);
     return a;
 }