helexa/crates/neuron/tests/qwen3_next_parity.rs

//! Numerical parity for the qwen3_next path (#92) against the HF
//! transformers reference, via the tiny self-contained fixture
//! generated by `script/dump_qwen3_next_tiny.py`.
//!
//! The fixture directory carries the WHOLE checkpoint (tiny
//! random-weight `Qwen3NextForCausalLM`: config.json +
//! model.safetensors, a few hundred KB) plus the reference
//! final-position logits, so this test needs no snapshot, no env var,
//! and runs in CI. It pins: flat-config normalisation, the `model.*`
//! weight prefix, the fused `in_proj_qkvz`/`in_proj_ba` de-interleave,
//! hybrid full/linear layer interleaving, and the MoE block (routing,
//! per-expert SwiGLU, shared expert + sigmoid gate).
//!
//! Self-skips (with a loud eprintln) while the fixture has not been
//! generated yet — regeneration instructions in the script docstring.

use candle_core::{DType, Device, Tensor};
use serde::Deserialize;
use std::path::{Path, PathBuf};

#[derive(Deserialize)]
struct Manifest {
    token_ids: Vec<u32>,
    files: std::collections::HashMap<String, FileEntry>,
}

#[derive(Deserialize)]
struct FileEntry {
    file: String,
    shape: Vec<usize>,
}

fn fixture_dir() -> PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/numerical/qwen3_next-tiny")
}

fn read_f32(path: &Path) -> Vec<f32> {
    let bytes = std::fs::read(path).unwrap_or_else(|e| panic!("read {path:?}: {e}"));
    bytes
        .chunks_exact(4)
        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
        .collect()
}

#[test]
fn tiny_qwen3_next_logits_match_hf_reference() {
    let dir = fixture_dir();
    let manifest_path = dir.join("manifest.json");
    if !manifest_path.exists() {
        eprintln!(
            "SKIP qwen3_next parity: fixture not generated yet — run \
             script/dump_qwen3_next_tiny.py --out {} on a host with \
             torch + transformers>=4.57",
            dir.display()
        );
        return;
    }

    let manifest: Manifest =
        serde_json::from_str(&std::fs::read_to_string(&manifest_path).expect("read manifest"))
            .expect("parse manifest");
    let logits_entry = &manifest.files["logits"];
    let reference = read_f32(&dir.join(&logits_entry.file));
    assert_eq!(
        reference.len(),
        logits_entry.shape.iter().product::<usize>()
    );

    let config_json = std::fs::read_to_string(dir.join("config.json")).expect("read config.json");
    let cfg = neuron::harness::arch::qwen3_5::Config::from_config_json(&config_json)
        .expect("normalise qwen3_next config");

    let dev = Device::Cpu;
    let st_path = dir.join("model.safetensors");
    // SAFETY: mmap of committed fixture files; nothing mutates them.
    let vb = unsafe {
        candle_nn::var_builder::ShardedSafeTensors::var_builder(
            std::slice::from_ref(&st_path),
            DType::F32,
            &dev,
        )
        .expect("build ShardedVarBuilder over fixture checkpoint")
    };
    let mut model = neuron::harness::arch::qwen3_5::Qwen3_5ForCausalLM::new(cfg, vb)
        .expect("load tiny qwen3_next checkpoint through neuron");

    let input = Tensor::new(manifest.token_ids.as_slice(), &dev)
        .unwrap()
        .unsqueeze(0)
        .unwrap();
    let logits = model.forward(&input, 0).expect("forward");
    let got: Vec<f32> = logits.flatten_all().unwrap().to_vec1().unwrap();
    assert_eq!(got.len(), reference.len());

    // f32-vs-f32 through an 8-layer doll-house model: agreement should
    // be tight (the qwen3_5 text fixtures observe max_abs ≈ 0.000,
    // cosine ≈ 1.0). Thresholds sit far above rounding noise and far
    // below any real wiring bug (a swapped de-interleave region, a
    // topk-before-softmax, a missing shared-expert gate all blow past
    // them instantly).
    let max_abs = got
        .iter()
        .zip(&reference)
        .map(|(a, b)| (a - b).abs())
        .fold(0f32, f32::max);
    let dot: f64 = got
        .iter()
        .zip(&reference)
        .map(|(a, b)| (*a as f64) * (*b as f64))
        .sum();
    let na: f64 = got.iter().map(|a| (*a as f64).powi(2)).sum::<f64>().sqrt();
    let nb: f64 = reference
        .iter()
        .map(|b| (*b as f64).powi(2))
        .sum::<f64>()
        .sqrt();
    let cosine = dot / (na * nb);

    eprintln!("qwen3_next parity: max_abs={max_abs:.6} cosine={cosine:.8}");
    assert!(max_abs < 1e-3, "max abs diff {max_abs} exceeds 1e-3");
    assert!(cosine > 0.9999, "cosine {cosine} below 0.9999");
}