@seanhogg/builderforce-memory-engine 2026.7.12 → 2026.9.27

package/src/import/foreign/codestral_mamba.ts

@@ -0,0 +1,224 @@
+/**
+ * import/foreign/codestral_mamba.ts — weight port for Codestral-Mamba
+ * (`mistralai/Mamba-Codestral-7B-v0.1`, `Mamba2ForCausalLM`,
+ * `model_type: "mamba2"`, Mamba-2 / SSD).
+ *
+ * Source layout confirmed against the repo's published HF `config.json`, its
+ * `model.safetensors.index.json` (579 tensors = 9 per layer × 64 + 3 globals),
+ * the header of `consolidated.safetensors`, and
+ * `transformers/models/mamba2/modeling_mamba2.py`:
+ *
+ *   backbone.embeddings.weight                    [vocab, d_model]
+ *   backbone.layers.{N}.norm.weight               [d_model]
+ *   backbone.layers.{N}.mixer.in_proj.weight      [2*d_inner + 2*n_groups*d_state + n_heads, d_model]
+ *   backbone.layers.{N}.mixer.conv1d.weight       [conv_dim, 1, d_conv]
+ *   backbone.layers.{N}.mixer.conv1d.bias         [conv_dim]
+ *   backbone.layers.{N}.mixer.dt_bias             [n_heads]
+ *   backbone.layers.{N}.mixer.A_log               [n_heads]
+ *   backbone.layers.{N}.mixer.D                   [n_heads]
+ *   backbone.layers.{N}.mixer.norm.weight         [d_inner]
+ *   backbone.layers.{N}.mixer.out_proj.weight     [d_model, d_inner]
+ *   backbone.norm_f.weight                        [d_model]
+ *   lm_head.weight                                [vocab, d_model]
+ *
+ * where `conv_dim = d_inner + 2 * n_groups * d_state`. Mamba-2 has no `x_proj`
+ * and no `dt_proj`: B, C and dt all come out of the fused `in_proj`.
+ *
+ * Three reconciliations against `Mamba2Block`:
+ *
+ *  1. **`in_proj` row order.** Upstream splits
+ *     `[gate (d_inner), x|B|C (conv_dim), dt (n_heads)]`; this engine's fused
+ *     projection is `[x (d_inner), B (n_groups*d_state), C (n_groups*d_state),
+ *     dt (n_heads)]` — the SAME tail, minus the leading gate rows. So the port
+ *     is one contiguous row slice starting at `d_inner`, and the gate rows are
+ *     declared as a discard (see 3).
+ *  2. **`A_log` convention.** Upstream uses `A = -exp(A_log)`; `ssd.ts` uses
+ *     `A = -softplus(A_log)`. The port re-expresses the same `A` as
+ *     `softplus⁻¹(exp(A_log))` so decay rates are preserved exactly, rather than
+ *     copying the raw value and silently changing every head's decay.
+ *  3. **The gate branch.** `Mamba2Block` has no `z` gate: it RMS-norms the scan
+ *     output and projects out, where upstream computes
+ *     `norm(y * silu(z))`. The `d_inner × d_model` gate rows therefore have no
+ *     target. That is a REAL loss of capability, not a formatting detail, and it
+ *     is reported as a `PortDiscard` on every layer so a caller sees it.
+ *
+ * Also mapped without transform: `mixer.norm.weight` → the block's inner
+ * `normWeight` (applied ungated here, per 3), `backbone.layers.N.norm.weight` →
+ * `preNormWeight`, `conv1d.{weight,bias}` (depthwise, taps reversed) → the fused
+ * `[x, B, C]` conv, and `D` / `dt_bias` per head. Linear weights are
+ * `[out, in]` row-major on both sides, so no transposes.
+ *
+ * `tie_word_embeddings: false`, so `lm_head.weight` is a real stored tensor with
+ * no target here (`HybridMambaModel` ties its head to the embedding) — it is
+ * reported as unmapped.
+ */
+
+import type { HybridMambaModelConfig, LayerSpec } from "../../model/mamba_model.js";
+import type { PortDiscard, PortPlan, PortRule } from "./plan.js";
+import {
+  expandFactor,
+  firstInt,
+  readInnerSize,
+  requireInt,
+  type ForeignConfig,
+  type ForeignMambaAdapter,
+  type PortTarget,
+} from "./adapter.js";
+
+const ID = "codestral_mamba";
+
+/** `backbone.embeddings` (HF) vs `backbone.embedding` (Mistral-native). */
+const EMBEDDING_NAMES = ["backbone.embeddings.weight", "backbone.embedding.weight"];
+
+interface Mamba2Geometry {
+  vocabSize: number;
+  dModel: number;
+  numLayers: number;
+  dInner: number;
+  dState: number;
+  dConv: number;
+  nHeads: number;
+  nGroups: number;
+  chunkLen: number;
+}
+
+function rules(a: Mamba2Geometry): PortRule[] {
+  const { vocabSize: V, dModel: E, numLayers: L, dInner: D, dState: N, dConv: K, nHeads: H, nGroups: G } = a;
+  const convDim = D + 2 * G * N;
+  /** Rows this engine keeps out of the fused projection: [x, B, C, dt]. */
+  const keptRows = convDim + H;
+
+  const out: PortRule[] = [
+    { target: "embedding", shape: [V, E], from: EMBEDDING_NAMES, op: { kind: "copy" } },
+  ];
+
+  for (let i = 0; i < L; i++) {
+    const m = `backbone.layers.${i}.mixer`;
+    const t = `layer${i}`;
+    out.push(
+      {
+        target: `${t}.wInProj`,
+        shape: [keptRows, E],
+        from: [`${m}.in_proj.weight`],
+        // Drop the leading gate rows; keep [x, B, C, dt] verbatim.
+        op: { kind: "rowSlice", start: D, count: keptRows },
+        note: "upstream rows are [gate, x, B, C, dt]; this block has no gate branch",
+      },
+      {
+        target: `${t}.wConv`,
+        shape: [convDim, K],
+        from: [`${m}.conv1d.weight`],
+        op: { kind: "convTaps" },
+        note: "depthwise [conv_dim, 1, d_conv] → [conv_dim, d_conv] with the tap axis reversed",
+      },
+      { target: `${t}.bConv`, shape: [convDim], from: [`${m}.conv1d.bias`], op: { kind: "copy" } },
+      {
+        target: `${t}.A_log`,
+        shape: [H],
+        from: [`${m}.A_log`],
+        op: { kind: "negExpToSoftplusInv" },
+        note: "upstream A = -exp(A_log); this engine's SSD kernel uses A = -softplus(A_log)",
+      },
+      { target: `${t}.dt_bias`, shape: [H], from: [`${m}.dt_bias`], op: { kind: "copy" } },
+      { target: `${t}.D_vec`, shape: [H], from: [`${m}.D`], op: { kind: "copy" } },
+      { target: `${t}.wOutProj`, shape: [E, D], from: [`${m}.out_proj.weight`], op: { kind: "copy" } },
+      {
+        target: `${t}.normWeight`,
+        shape: [D],
+        from: [`${m}.norm.weight`],
+        op: { kind: "copy" },
+        note: "upstream this gain belongs to a GATED RMSNorm; applied ungated here",
+      },
+      {
+        target: `${t}.preNormWeight`,
+        shape: [E],
+        from: [`backbone.layers.${i}.norm.weight`],
+        op: { kind: "copy" },
+      },
+    );
+  }
+
+  out.push({ target: "final_norm", shape: [E], from: ["backbone.norm_f.weight"], op: { kind: "copy" } });
+  return out;
+}
+
+function discards(a: Mamba2Geometry): PortDiscard[] {
+  const out: PortDiscard[] = [];
+  for (let i = 0; i < a.numLayers; i++) {
+    out.push({
+      source: `backbone.layers.${i}.mixer.in_proj.weight`,
+      elements: a.dInner * a.dModel,
+      reason:
+        "the leading d_inner gate (z) rows — Mamba2Block computes norm(y) rather than " +
+        "norm(y * silu(z)), so it has no gate branch to receive them",
+    });
+  }
+  return out;
+}
+
+export const codestralMambaAdapter: ForeignMambaAdapter = {
+  id: ID,
+  label: "Codestral-Mamba (Mamba-2 / SSD)",
+  modelTypes: ["mamba2"],
+  // `MambaCodestralForCausalLM` is not a transformers class — the published
+  // config says `Mamba2ForCausalLM` — but it is a common mis-citation, so it is
+  // accepted here rather than sending a correct checkpoint down the error path.
+  architectures: ["Mamba2ForCausalLM", "MambaCodestralForCausalLM"],
+
+  describe(config: ForeignConfig): PortTarget {
+    const dModel = requireInt(config, "hidden_size");
+    const numLayers = requireInt(config, "num_hidden_layers");
+    const vocabSize = requireInt(config, "vocab_size");
+    const dState = requireInt(config, "state_size");
+    const dConv = requireInt(config, "conv_kernel");
+    const nHeads = requireInt(config, "num_heads");
+    const nGroups = firstInt(config, ["n_groups"]) ?? 1;
+    const chunkLen = firstInt(config, ["chunk_size"]) ?? 256;
+    const dInner = readInnerSize(config, dModel);
+    const expand = expandFactor(dInner, dModel, "Codestral-Mamba");
+
+    // Upstream's own invariant (`hidden_size * expand == num_heads * head_dim`);
+    // Mamba2Block enforces the same divisibility, so check it here where the
+    // message can name the config fields that disagree.
+    const headDim = firstInt(config, ["head_dim"]);
+    if (dInner % nHeads !== 0) {
+      throw new Error(
+        `import/foreign: Codestral-Mamba d_inner ${dInner} is not divisible by num_heads ${nHeads}`,
+      );
+    }
+    if (headDim !== undefined && headDim * nHeads !== dInner) {
+      throw new Error(
+        `import/foreign: Codestral-Mamba config is inconsistent — head_dim ${headDim} × ` +
+          `num_heads ${nHeads} = ${headDim * nHeads}, but d_inner is ${dInner}`,
+      );
+    }
+
+    const layers: LayerSpec[] = Array.from({ length: numLayers }, () => ({ type: "mamba2" as const }));
+    const modelConfig: HybridMambaModelConfig = {
+      vocabSize,
+      dModel,
+      numLayers,
+      layers,
+      dState,
+      dConv,
+      expand,
+      nHeads,
+      nGroups,
+      chunkLen,
+    };
+
+    const geometry: Mamba2Geometry = {
+      vocabSize,
+      dModel,
+      numLayers,
+      dInner,
+      dState,
+      dConv,
+      nHeads,
+      nGroups,
+      chunkLen,
+    };
+
+    return { modelConfig, plan: { adapter: ID, rules: rules(geometry), discards: discards(geometry) } };
+  },
+};

package/src/import/foreign/falcon_mamba.ts

@@ -0,0 +1,173 @@
+/**
+ * import/foreign/falcon_mamba.ts — weight port for Falcon-Mamba
+ * (`FalconMambaForCausalLM`, `model_type: "falcon_mamba"`, Mamba-1 / S6).
+ *
+ * Source layout confirmed against `tiiuae/falcon-mamba-7b`'s published
+ * `config.json`, `model.safetensors.index.json` (643 tensors = 10 per layer × 64
+ * + 3 globals) and the safetensors header of shard 1, plus
+ * `transformers/models/falcon_mamba/modeling_falcon_mamba.py`:
+ *
+ *   backbone.embeddings.weight                       [vocab, d_model]
+ *   backbone.layers.{N}.norm.weight                  [d_model]
+ *   backbone.layers.{N}.mixer.in_proj.weight         [2*d_inner, d_model]
+ *   backbone.layers.{N}.mixer.conv1d.weight          [d_inner, 1, d_conv]
+ *   backbone.layers.{N}.mixer.conv1d.bias            [d_inner]
+ *   backbone.layers.{N}.mixer.x_proj.weight          [dt_rank + 2*d_state, d_inner]
+ *   backbone.layers.{N}.mixer.dt_proj.weight         [d_inner, dt_rank]
+ *   backbone.layers.{N}.mixer.dt_proj.bias           [d_inner]
+ *   backbone.layers.{N}.mixer.A_log                  [d_inner, d_state]
+ *   backbone.layers.{N}.mixer.D                      [d_inner]
+ *   backbone.layers.{N}.mixer.out_proj.weight        [d_model, d_inner]
+ *   backbone.norm_f.weight                           [d_model]
+ *   lm_head.weight                                   [vocab, d_model]
+ *
+ * Conventions that line up with `Mamba1Block` as-is, so most rules are copies:
+ *   • `in_proj` rows are `[x, gate]` (`hidden_states, gate = …chunk(2, dim=1)`),
+ *     and `Mamba1Block.forward` also takes the first half as `x` and the second
+ *     as the gate `z`.
+ *   • `x_proj` rows are `[dt, B, C]`, matching the block's split order.
+ *   • `A = -exp(A_log)` in both — `selective_scan.ts` computes
+ *     `a_cont = -exp(clamp(a_log, -10, 5))`.
+ *   • `dt_proj.bias` is added inside `softplus(dt + bias)` upstream; here it is
+ *     the `dt_proj` linear's bias and the scan applies `softplus` afterwards —
+ *     algebraically the same value.
+ *   • Linear weights are `[out, in]` row-major on both sides (no transpose).
+ *
+ * Fidelity caveats recorded on the plan rather than hidden:
+ *   • `use_bias: false` upstream, so `in_proj` / `x_proj` / `out_proj` ship no
+ *     bias; this engine's blocks always own one, and they are filled with zeros
+ *     (an exact no-op, not an approximation).
+ *   • Falcon-Mamba RMS-normalises `dt`, `B` and `C` inside the mixer. Those
+ *     norms are `FalconMambaWeightlessRMSNorm` with `persistent=False` buffers,
+ *     so they contribute NO tensors to the checkpoint — nothing is dropped from
+ *     the file — but this engine's S6 scan has no such normalisation step, so a
+ *     ported model is not numerically identical to the original.
+ *   • `tie_word_embeddings: false`, so `lm_head.weight` is a real stored tensor.
+ *     `HybridMambaModel` always ties its LM head to the embedding table, so that
+ *     tensor has no target and is reported as unmapped.
+ */
+
+import type { HybridMambaModelConfig } from "../../model/mamba_model.js";
+import type { PortPlan, PortRule } from "./plan.js";
+import {
+  expandFactor,
+  readInnerSize,
+  readTimeStepRank,
+  requireInt,
+  type ForeignConfig,
+  type ForeignMambaAdapter,
+  type PortTarget,
+} from "./adapter.js";
+
+const ID = "falcon_mamba";
+
+/** `backbone.embeddings` (HF) vs `backbone.embedding` (Mamba-SSM native). */
+const EMBEDDING_NAMES = ["backbone.embeddings.weight", "backbone.embedding.weight"];
+
+function rules(a: {
+  vocabSize: number;
+  dModel: number;
+  numLayers: number;
+  dInner: number;
+  dState: number;
+  dConv: number;
+  dtRank: number;
+}): PortRule[] {
+  const { vocabSize: V, dModel: E, numLayers: L, dInner: D, dState: N, dConv: K, dtRank: R } = a;
+  const out: PortRule[] = [
+    { target: "embedding", shape: [V, E], from: EMBEDDING_NAMES, op: { kind: "copy" } },
+  ];
+
+  for (let i = 0; i < L; i++) {
+    const m = `backbone.layers.${i}.mixer`;
+    const t = `layer${i}`;
+    out.push(
+      { target: `${t}.wInProj`, shape: [2 * D, E], from: [`${m}.in_proj.weight`], op: { kind: "copy" } },
+      {
+        target: `${t}.bInProj`,
+        shape: [2 * D],
+        from: [],
+        op: { kind: "fill", value: 0 },
+        note: "Falcon-Mamba sets use_bias=false — in_proj ships no bias; zeros are exact",
+      },
+      {
+        target: `${t}.wConv`,
+        shape: [D, K],
+        from: [`${m}.conv1d.weight`],
+        op: { kind: "convTaps" },
+        note: "depthwise [d_inner, 1, d_conv] → [d_inner, d_conv] with the tap axis reversed",
+      },
+      { target: `${t}.bConv`, shape: [D], from: [`${m}.conv1d.bias`], op: { kind: "copy" } },
+      { target: `${t}.wXProj`, shape: [R + 2 * N, D], from: [`${m}.x_proj.weight`], op: { kind: "copy" } },
+      {
+        target: `${t}.bXProj`,
+        shape: [R + 2 * N],
+        from: [],
+        op: { kind: "fill", value: 0 },
+        note: "x_proj is bias-free upstream; zeros are exact",
+      },
+      { target: `${t}.wDtProj`, shape: [D, R], from: [`${m}.dt_proj.weight`], op: { kind: "copy" } },
+      { target: `${t}.bDtProj`, shape: [D], from: [`${m}.dt_proj.bias`], op: { kind: "copy" } },
+      { target: `${t}.A_log`, shape: [D, N], from: [`${m}.A_log`], op: { kind: "copy" } },
+      { target: `${t}.D_vec`, shape: [D], from: [`${m}.D`], op: { kind: "copy" } },
+      { target: `${t}.wOutProj`, shape: [E, D], from: [`${m}.out_proj.weight`], op: { kind: "copy" } },
+      {
+        target: `${t}.bOutProj`,
+        shape: [E],
+        from: [],
+        op: { kind: "fill", value: 0 },
+        note: "out_proj is bias-free upstream; zeros are exact",
+      },
+      {
+        target: `${t}.normWeight`,
+        shape: [E],
+        from: [`backbone.layers.${i}.norm.weight`],
+        op: { kind: "copy" },
+      },
+    );
+  }
+
+  out.push({ target: "final_norm", shape: [E], from: ["backbone.norm_f.weight"], op: { kind: "copy" } });
+  return out;
+}
+
+export const falconMambaAdapter: ForeignMambaAdapter = {
+  id: ID,
+  label: "Falcon-Mamba (Mamba-1 / S6)",
+  modelTypes: ["falcon_mamba"],
+  architectures: ["FalconMambaForCausalLM"],
+
+  describe(config: ForeignConfig): PortTarget {
+    const dModel = requireInt(config, "hidden_size");
+    const numLayers = requireInt(config, "num_hidden_layers");
+    const vocabSize = requireInt(config, "vocab_size");
+    const dState = requireInt(config, "state_size");
+    const dConv = requireInt(config, "conv_kernel");
+    const dInner = readInnerSize(config, dModel);
+    const expand = expandFactor(dInner, dModel, "Falcon-Mamba");
+    const dtRank = readTimeStepRank(config, dModel);
+
+    const modelConfig: HybridMambaModelConfig = {
+      vocabSize,
+      dModel,
+      numLayers,
+      dState,
+      dConv,
+      expand,
+      // Every layer is a Mamba-1 mixer (the model's default schedule), but the
+      // dt rank and conv bias are per-block settings the shorthand can't carry.
+      defaultMamba1: { dtRank, biasConv: true },
+    };
+
+    const plan: PortPlan = {
+      adapter: ID,
+      rules: rules({ vocabSize, dModel, numLayers, dInner, dState, dConv, dtRank }),
+      // Nothing is partially consumed: every Falcon-Mamba tensor this plan
+      // touches is taken whole. `lm_head.weight` is untouched entirely and
+      // surfaces through `unmappedSources`.
+      discards: [],
+    };
+
+    return { modelConfig, plan };
+  },
+};

package/src/import/foreign/index.ts

@@ -0,0 +1,37 @@
+/**
+ * import/foreign/index.ts — the foreign-checkpoint weight-port surface.
+ *
+ * Warm-start this engine's `HybridMambaModel` from a published Mamba checkpoint:
+ * pass the checkpoint's `config.json` and its tensors, get back the model config
+ * to build plus a name-keyed set of weights, together with an explicit account of
+ * anything that did not map cleanly.
+ *
+ * Supported: Falcon-Mamba (Mamba-1) and Codestral-Mamba (Mamba-2). Transformer
+ * checkpoints are NOT portable at any tensor naming — distillation is the only
+ * route, and {@link foreignMambaAdapterFor} says so when it rejects one.
+ */
+
+export {
+  FOREIGN_MAMBA_ADAPTERS,
+  foreignMambaAdapterFor,
+  portForeignMamba,
+  portForeignMambaSafetensors,
+} from "./registry.js";
+export type { PortedCheckpoint } from "./registry.js";
+
+export { applyPortedWeights } from "./apply.js";
+
+export { executePortPlan, normaliseSourceName } from "./plan.js";
+export type {
+  PortDiscard,
+  PortOp,
+  PortPlan,
+  PortRule,
+  PortedTensors,
+  SynthesisedTarget,
+} from "./plan.js";
+
+export type { ForeignConfig, ForeignMambaAdapter, PortTarget } from "./adapter.js";
+
+export { falconMambaAdapter } from "./falcon_mamba.js";
+export { codestralMambaAdapter } from "./codestral_mamba.js";