lfm2.test.ts

/**
 * End-to-end gate for the LFM2.5 hybrid forward/backward on the CPU backend.
 * Proves the wiring that the main engine lacked: a per-layer conv|attention
 * dispatch that (a) runs, (b) is differentiable through BOTH mixer types,
 * (c) trains a fact to convergence, (d) generates it back faithfully, and
 * (e) has a conv-KV-cache whose incremental decode matches full recompute.
 */
import { beforeAll, describe, expect, it } from "vitest";
import * as tf from "@tensorflow/tfjs";
import { Lfm2Model, type Lfm2Config, type LayerWeights, type LayerCache } from "./lfm2-model";

// Tiny model with the real interleave pattern (both mixer types present).
function tinyConfig(): Lfm2Config {
  return {
    hiddenSize: 32,
    numLayers: 6,
    layerTypes: ["conv", "conv", "attention", "conv", "attention", "conv"],
    numHeads: 4,
    numKvHeads: 2,
    headDim: 8,
    intermediateSize: 64,
    vocabSize: 40,
    rmsEps: 1e-5,
    ropeTheta: 1e6,
    convDim: 32, // conv_dim == hidden in LFM2
    convK: 3,
    eosTokenIds: [39],
  };
}

function lcg(seed: number) {
  let s = seed;
  return () => ((s = (s * 1103515245 + 12345) % 2147483648) / 2147483648 - 0.5);
}

function buildModel(cfg: Lfm2Config, seed: number, trainable = false): Lfm2Model {
  const m = new Lfm2Model(cfg);
  const rnd = lcg(seed);
  const wrap = (t: tf.Tensor) => (trainable ? tf.variable(t) : tf.keep(t)) as tf.Tensor;
  const mat = (r: number, c: number, scale = 0.3) =>
    wrap(tf.tensor2d(Array.from({ length: r * c }, () => rnd() * scale), [r, c]));
  const vec = (n: number) => wrap(tf.tensor1d(Array.from({ length: n }, () => 1 + rnd() * 0.1)));
  const { hiddenSize: H, numHeads, numKvHeads, headDim, intermediateSize: I, convDim: C, convK: K } = cfg;
  m.embed = wrap(tf.tensor2d(Array.from({ length: cfg.vocabSize * H }, () => rnd()), [cfg.vocabSize, H]));
  m.finalNorm = vec(H);
  for (let i = 0; i < cfg.numLayers; i++) {
    const common = { operatorNorm: vec(H), ffnNorm: vec(H), w1: mat(H, I), w2: mat(I, H), w3: mat(H, I) };
    if (cfg.layerTypes[i] === "attention") {
      m.layers.push({
        type: "attention",
        ...common,
        qProj: mat(H, numHeads * headDim),
        kProj: mat(H, numKvHeads * headDim),
        vProj: mat(H, numKvHeads * headDim),
        oProj: mat(numHeads * headDim, H),
        qNorm: vec(headDim),
        kNorm: vec(headDim),
      } as LayerWeights);
    } else {
      m.layers.push({
        type: "conv",
        ...common,
        inProj: mat(H, 3 * C),
        convWeight: mat(K, C),
        outProj: mat(C, H),
      } as LayerWeights);
    }
  }
  return m;
}

const PROMPT = [1, 2, 3, 4, 5];
const COMPLETION = [10, 11, 12, 39];

function buildStep(promptIds: number[], completionIds: number[]) {
  const all = [...promptIds, ...completionIds];
  const inputIds = all.slice(0, -1);
  const targetIds = all.slice(1);
  const mask = new Float32Array(targetIds.length);
  for (let i = promptIds.length - 1; i < targetIds.length; i++) mask[i] = 1;
  return {
    ids: tf.tensor2d([inputIds], [1, inputIds.length], "int32"),
    targets: tf.tensor1d(targetIds, "int32"),
    mask: tf.tensor1d(mask, "float32"),
  };
}

function loss(model: Lfm2Model, step: ReturnType<typeof buildStep>): tf.Scalar {
  const hidden = model.forwardHidden(step.ids);
  const T = hidden.shape[1] as number;
  const flat = hidden.reshape([T, -1]) as tf.Tensor2D;
  const logits = flat.matMul(model.embed as tf.Tensor2D, false, true);
  const nll = tf.losses.softmaxCrossEntropy(
    tf.oneHot(step.targets as tf.Tensor1D, model.cfg.vocabSize),
    logits,
    undefined,
    undefined,
    tf.Reduction.NONE,
  ) as tf.Tensor1D;
  const l = nll.mul(step.mask).sum().div(step.mask.sum().add(1e-8)) as tf.Scalar;
  hidden.dispose();
  return l;
}

// greedy generation via full recompute (no cache) — reference path
async function genFull(model: Lfm2Model, promptIds: number[], maxTokens: number): Promise<number[]> {
  const out: number[] = [];
  let seq = [...promptIds];
  for (let i = 0; i < maxTokens; i++) {
    const next = tf.tidy(() => {
      const ids = tf.tensor2d([seq], [1, seq.length], "int32");
      const hidden = model.forwardHidden(ids);
      return model.lastLogits(hidden).argMax(-1).dataSync()[0]!;
    });
    if (model.cfg.eosTokenIds.includes(next)) break;
    out.push(next);
    seq.push(next);
  }
  return out;
}

// greedy generation via incremental caches (conv-state + KV) — production path
async function genCached(model: Lfm2Model, promptIds: number[], maxTokens: number): Promise<number[]> {
  const out: number[] = [];
  let caches: LayerCache[] = [];
  const write: LayerCache[] = [];
  let hidden = model.forwardHidden(tf.tensor2d([promptIds], [1, promptIds.length], "int32"), null, write);
  caches = write.slice();
  for (let i = 0; i < maxTokens; i++) {
    const next = tf.tidy(() => model.lastLogits(hidden).argMax(-1).dataSync()[0]!);
    hidden.dispose();
    if (model.cfg.eosTokenIds.includes(next)) break;
    out.push(next);
    const w2: LayerCache[] = [];
    hidden = model.forwardHidden(tf.tensor2d([[next]], [1, 1], "int32"), caches, w2);
    // free old caches, adopt new
    for (const c of caches) (c.type === "attn" ? [c.k, c.v] : [c.bx]).forEach((t) => t.dispose());
    caches = w2.slice();
  }
  hidden.dispose();
  for (const c of caches) (c.type === "attn" ? [c.k, c.v] : [c.bx]).forEach((t) => t.dispose());
  return out;
}

describe("LFM2 hybrid forward/backward", () => {
  beforeAll(async () => {
    await tf.setBackend("cpu");
    await tf.ready();
  });

  it("forward produces correct-shape hidden states through both mixer types", () => {
    const cfg = tinyConfig();
    const model = buildModel(cfg, 1);
    const ids = tf.tensor2d([PROMPT], [1, PROMPT.length], "int32");
    const hidden = model.forwardHidden(ids);
    expect(hidden.shape).toEqual([1, PROMPT.length, cfg.hiddenSize]);
    expect(Number.isFinite(hidden.square().sum().dataSync()[0]!)).toBe(true);
    ids.dispose();
    hidden.dispose();
    model.dispose();
  });

  it("has finite, nonzero gradients through conv AND attention weights", () => {
    const cfg = tinyConfig();
    const model = buildModel(cfg, 2, /*trainable*/ true);
    const step = buildStep(PROMPT, COMPLETION);
    const convLayer = model.layers[0]!; // conv
    const attnLayer = model.layers[2]!; // attention
    const probes: Record<string, tf.Variable> = {
      conv_inProj: (convLayer as unknown as { inProj: tf.Variable }).inProj,
      conv_convWeight: (convLayer as unknown as { convWeight: tf.Variable }).convWeight,
      conv_outProj: (convLayer as unknown as { outProj: tf.Variable }).outProj,
      attn_qProj: (attnLayer as unknown as { qProj: tf.Variable }).qProj,
      attn_vProj: (attnLayer as unknown as { vProj: tf.Variable }).vProj,
    };
    const { grads, value } = tf.variableGrads(() => loss(model, step), Object.values(probes));
    expect(Number.isFinite(value.dataSync()[0]!)).toBe(true);
    for (const [name, g] of Object.entries(grads)) {
      const abs = tf.tidy(() => g.abs().sum().dataSync()[0]!);
      expect(Number.isFinite(abs), `${name} grad finite`).toBe(true);
      expect(abs, `${name} grad nonzero`).toBeGreaterThan(0);
    }
    value.dispose();
    Object.values(grads).forEach((g) => g.dispose());
    tf.dispose([step.ids, step.targets, step.mask]);
    model.dispose();
  });

  it("trains a fact to convergence and generates it back", async () => {
    const cfg = tinyConfig();
    const model = buildModel(cfg, 3, /*trainable*/ true);
    const vars = model.layers.flatMap((l) => Object.values(l).filter((t): t is tf.Variable => t instanceof tf.Variable));
    vars.push(model.embed as tf.Variable, model.finalNorm as tf.Variable);
    const opt = tf.train.adam(3e-3);
    const step = buildStep(PROMPT, COMPLETION);
    let last = NaN;
    for (let i = 0; i < 150; i++) {
      const c = opt.minimize(() => loss(model, step), true, vars)!;
      last = c.dataSync()[0]!;
      c.dispose();
    }
    const gen = await genFull(model, PROMPT, 6);
    expect(last).toBeLessThan(0.05);
    expect(gen.slice(0, 3)).toEqual([10, 11, 12]);
    opt.dispose();
    tf.dispose([step.ids, step.targets, step.mask]);
    model.dispose();
  }, 60_000);

  it("generation is faithful: greedy first token == teacher-forced argmax", () => {
    const cfg = tinyConfig();
    const model = buildModel(cfg, 4);
    const step = buildStep(PROMPT, COMPLETION);
    const tfArg = tf.tidy(() => {
      const hidden = model.forwardHidden(step.ids);
      const promptOnly = hidden.slice([0, 0, 0], [1, PROMPT.length, -1]);
      return model.lastLogits(promptOnly).argMax(-1).dataSync()[0]!;
    });
    const genFirst = tf.tidy(() => {
      const hidden = model.forwardHidden(tf.tensor2d([PROMPT], [1, PROMPT.length], "int32"));
      return model.lastLogits(hidden).argMax(-1).dataSync()[0]!;
    });
    expect(genFirst).toBe(tfArg);
    tf.dispose([step.ids, step.targets, step.mask]);
    model.dispose();
  });

  it("conv-KV-cache decode matches full recompute (the cache is correct)", async () => {
    const cfg = tinyConfig();
    const model = buildModel(cfg, 5);
    const full = await genFull(model, PROMPT, 8);
    const cached = await genCached(model, PROMPT, 8);
    expect(cached).toEqual(full);
    expect(full.length).toBeGreaterThan(0); // it actually generated something
    model.dispose();
  }, 30_000);
});