qwen35.test.ts

/**
 * End-to-end gate for the Qwen3.5 linear_attn/attention hybrid forward/backward
 * on the CPU backend — closes the `linear-attn-unsupported` gap the toolkit's
 * classifier flags for this model (see ../toolkit/classifier.ts and
 * ../qwen3.5-linear-attn/RESULTS.md). Proves the same five things
 * lfm2-hybrid/lfm2.test.ts proved for LFM2.5's conv/attention hybrid, here for
 * Gated-DeltaNet/attention: (a) forward runs through both mixer types, (b) is
 * differentiable through both, (c) trains a fact to convergence, (d) generates
 * it back faithfully, (e) the recurrent-state + conv-history cache's
 * incremental decode matches full recompute.
 */
import { beforeAll, describe, expect, it } from "vitest";
import * as tf from "@tensorflow/tfjs";
import { Qwen35Model, type Qwen35Config, type LayerWeights, type LayerCache } from "./qwen35-model.ts";

// Tiny model with the real interleave pattern (both mixer types present).
function tinyConfig(): Qwen35Config {
  return {
    hiddenSize: 32,
    numLayers: 6,
    layerTypes: ["linear", "linear", "attention", "linear", "attention", "linear"],
    numHeads: 4,
    numKvHeads: 2,
    headDim: 8,
    linearHeads: 2,
    keyHeadDim: 4,
    valueHeadDim: 4,
    convKernel: 4,
    intermediateSize: 64,
    vocabSize: 40,
    rmsEps: 1e-6,
    ropeTheta: 1e7,
    rotaryDim: 4, // partial rotary, < headDim(8) — exercises the pass-through path too
    eosTokenIds: [39],
  };
}

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

function buildModel(cfg: Qwen35Config, seed: number, trainable = false): Qwen35Model {
  const m = new Qwen35Model(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, base = 0) => wrap(tf.tensor1d(Array.from({ length: n }, () => base + rnd() * 0.1)));
  const {
    hiddenSize: H, numHeads, numKvHeads, headDim, intermediateSize: I,
    linearHeads: V, keyHeadDim: Kd, valueHeadDim: Vd, convKernel: K,
  } = cfg;
  const keyDim = V * Kd, valueDim = V * Vd, convDim = keyDim * 2 + valueDim;
  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 = {
      inputLayernorm: vec(H),
      postAttentionLayernorm: vec(H),
      gateProj: mat(H, I),
      upProj: mat(H, I),
      downProj: mat(I, H),
    };
    if (cfg.layerTypes[i] === "attention") {
      m.layers.push({
        type: "attention",
        ...common,
        qProj: mat(H, numHeads * headDim * 2),
        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: "linear",
        ...common,
        inProjQkv: mat(H, convDim),
        inProjZ: mat(H, valueDim),
        inProjA: mat(H, V),
        inProjB: mat(H, V),
        convWeight: mat(K, convDim, 0.5),
        dtBias: vec(V, 1), // keep decay gate away from a knife-edge at seed 0
        aLog: vec(V, 0.5),
        gatedNorm: vec(Vd),
        outProj: mat(valueDim, 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: Qwen35Model, 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: Qwen35Model, 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;
}

function disposeCache(c: LayerCache) {
  (c.type === "attn" ? [c.k, c.v] : [c.convHist, c.state]).forEach((t) => t.dispose());
}

// greedy generation via incremental caches (recurrent-state + conv-history + KV) — production path
async function genCached(model: Qwen35Model, 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);
    caches.forEach(disposeCache);
    caches = w2.slice();
  }
  hidden.dispose();
  caches.forEach(disposeCache);
  return out;
}

describe("Qwen3.5 linear_attn/attention 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 linear_attn AND attention weights", () => {
    const cfg = tinyConfig();
    const model = buildModel(cfg, 2, /*trainable*/ true);
    const step = buildStep(PROMPT, COMPLETION);
    const linLayer = model.layers[0]!; // linear
    const attnLayer = model.layers[2]!; // attention
    const probes: Record<string, tf.Variable> = {
      lin_inProjQkv: (linLayer as unknown as { inProjQkv: tf.Variable }).inProjQkv,
      lin_convWeight: (linLayer as unknown as { convWeight: tf.Variable }).convWeight,
      lin_aLog: (linLayer as unknown as { aLog: tf.Variable }).aLog,
      lin_outProj: (linLayer 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("recurrent-state + conv-history cache decode matches full recompute", 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);
});