qwen35-loader.test.ts

/**
 * Validates the Qwen3.5 real-weights mapping WITHOUT the 2.34 GB download: a
 * mock safetensors reader returns the real tensor names/shapes with known
 * data, and we assert the transforms (conv weight reshape, Linear transpose,
 * config normalization) assemble into a runnable model. Numeric parity vs
 * `transformers` needs the real checkpoint — see verify-parity.ts.
 */
import { beforeAll, describe, expect, it } from "vitest";
import * as tf from "@tensorflow/tfjs";
import { loadQwen35, normalizeQwen35Config } from "./qwen35-loader.ts";
import type { TensorReader } from "../lfm2-hybrid/lfm2-loader.ts";

// real Qwen/Qwen3.5-0.8B text backbone config fields (subset), as fetched from HF.
const realCfg = {
  model_type: "qwen3_5_text",
  hidden_size: 1024,
  num_hidden_layers: 24,
  num_attention_heads: 8,
  num_key_value_heads: 2,
  head_dim: 256,
  linear_num_key_heads: 16,
  linear_num_value_heads: 16,
  linear_key_head_dim: 128,
  linear_value_head_dim: 128,
  linear_conv_kernel_dim: 4,
  intermediate_size: 3584,
  vocab_size: 248320,
  rms_norm_eps: 1e-6,
  rope_parameters: { rope_theta: 1e7, partial_rotary_factor: 0.25 },
  eos_token_id: 248044,
  layer_types: [
    "linear_attention", "linear_attention", "linear_attention", "full_attention",
    "linear_attention", "linear_attention", "linear_attention", "full_attention",
    "linear_attention", "linear_attention", "linear_attention", "full_attention",
    "linear_attention", "linear_attention", "linear_attention", "full_attention",
    "linear_attention", "linear_attention", "linear_attention", "full_attention",
    "linear_attention", "linear_attention", "linear_attention", "full_attention",
  ],
};

describe("Qwen3.5 real-weights loader", () => {
  beforeAll(async () => {
    await tf.setBackend("cpu");
    await tf.ready();
  });

  it("normalizes the real config: 18 linear + 6 attention, rotaryDim 64", () => {
    const cfg = normalizeQwen35Config(realCfg as any);
    expect(cfg.layerTypes.filter((t) => t === "attention").length).toBe(6);
    expect(cfg.layerTypes.filter((t) => t === "linear").length).toBe(18);
    expect([3, 7, 11, 15, 19, 23].every((i) => cfg.layerTypes[i] === "attention")).toBe(true);
    expect(cfg.headDim).toBe(256);
    expect(cfg.rotaryDim).toBe(64); // 256 * 0.25
    expect(cfg.linearHeads).toBe(16);
    expect(cfg.convKernel).toBe(4);
    expect(cfg.eosTokenIds).toEqual([248044]);
  });

  it("assembles a runnable model from a mock reader (all names/shapes line up)", async () => {
    const cfg = normalizeQwen35Config({
      ...realCfg,
      hidden_size: 32,
      num_hidden_layers: 2,
      layer_types: ["linear_attention", "full_attention"],
      num_attention_heads: 4,
      num_key_value_heads: 2,
      head_dim: 16,
      linear_num_key_heads: 2,
      linear_num_value_heads: 2,
      linear_key_head_dim: 8,
      linear_value_head_dim: 8,
      linear_conv_kernel_dim: 4,
      intermediate_size: 64,
      vocab_size: 40,
    } as any);
    const H = 32, I = 64, V = 40;
    const keyDim = 2 * 8, valueDim = 2 * 8, convDim = keyDim * 2 + valueDim; // 16,16,48
    const qOut = 4 * 16 * 2, kvOut = 2 * 16, oIn = 4 * 16; // attention dims (headDim=16)
    const shapes: Record<string, number[]> = {
      "model.language_model.embed_tokens.weight": [V, H],
      "model.language_model.norm.weight": [H],
      "model.language_model.layers.0.input_layernorm.weight": [H],
      "model.language_model.layers.0.post_attention_layernorm.weight": [H],
      "model.language_model.layers.0.mlp.gate_proj.weight": [I, H],
      "model.language_model.layers.0.mlp.up_proj.weight": [I, H],
      "model.language_model.layers.0.mlp.down_proj.weight": [H, I],
      "model.language_model.layers.0.linear_attn.in_proj_qkv.weight": [convDim, H],
      "model.language_model.layers.0.linear_attn.in_proj_z.weight": [valueDim, H],
      "model.language_model.layers.0.linear_attn.in_proj_a.weight": [2, H],
      "model.language_model.layers.0.linear_attn.in_proj_b.weight": [2, H],
      "model.language_model.layers.0.linear_attn.conv1d.weight": [convDim, 1, 4],
      "model.language_model.layers.0.linear_attn.dt_bias": [2],
      "model.language_model.layers.0.linear_attn.A_log": [2],
      "model.language_model.layers.0.linear_attn.norm.weight": [8],
      "model.language_model.layers.0.linear_attn.out_proj.weight": [H, valueDim],
      "model.language_model.layers.1.input_layernorm.weight": [H],
      "model.language_model.layers.1.post_attention_layernorm.weight": [H],
      "model.language_model.layers.1.mlp.gate_proj.weight": [I, H],
      "model.language_model.layers.1.mlp.up_proj.weight": [I, H],
      "model.language_model.layers.1.mlp.down_proj.weight": [H, I],
      "model.language_model.layers.1.self_attn.q_proj.weight": [qOut, H],
      "model.language_model.layers.1.self_attn.k_proj.weight": [kvOut, H],
      "model.language_model.layers.1.self_attn.v_proj.weight": [kvOut, H],
      "model.language_model.layers.1.self_attn.o_proj.weight": [H, oIn],
      "model.language_model.layers.1.self_attn.q_norm.weight": [16],
      "model.language_model.layers.1.self_attn.k_norm.weight": [16],
    };
    let seed = 1;
    const rnd = () => ((seed = (seed * 1103515245 + 12345) % 2147483648) / 2147483648 - 0.5);
    const reader: TensorReader = {
      index: { tensors: new Map(Object.keys(shapes).map((n) => [n, true])) },
      read: async (name: string) => {
        const shape = shapes[name];
        if (!shape) throw new Error(`mock reader: unexpected tensor ${name}`);
        const n = shape.reduce((a, b) => a * b, 1);
        return { data: Float32Array.from({ length: n }, () => rnd()), shape };
      },
    } as TensorReader;
    const model = await loadQwen35(cfg, reader);
    expect(model.layers[0]!.type).toBe("linear");
    expect(model.layers[1]!.type).toBe("attention");
    const ids = tf.tensor2d([[1, 2, 3]], [1, 3], "int32");
    const hidden = model.forwardHidden(ids);
    expect(hidden.shape).toEqual([1, 3, H]);
    expect(Number.isFinite(hidden.square().sum().dataSync()[0]!)).toBe(true);
    ids.dispose();
    hidden.dispose();
    model.dispose();
  });
});