5e03008268
The stored logit hashes depend on the platform's math library as well
as on the code, but a mismatch blamed only the code. Each platform's
lines now carry a `math-library` line naming the library they were
made with (macOS 27.2's libsystem_m, Ubuntu GLIBC 2.39 on Linux arm64,
2.43 on Linux x86_64). A mismatch names it, says whether the ways of
running the model agree with each other (if they do, this
environment's library may round differently), and says how to tell.
examples/cohere2_reference.rs is the generator: the test's checkpoints,
scenario and hash, through forward_prompt and forward_token only, so it
runs unchanged at 30c4cf1ee. Run there, with and without rayon, it
reproduces all 78 stored lines on all three platforms, and it prints
the `math-library` line for a new one.
With JOULECLAW_FABRIC=1 in the environment, the main check ran in a
child process and the 1- and 3-thread checks never ran at all. The
test process now checks nothing itself: it starts one child per
configuration, each with the fabric off, and each child must report
its check.
Tests failing first: the mismatch message did not name the library;
a test run with JOULECLAW_FABRIC=1 reported 0 of its 3 checks. Both
pass now, with the switch unset and set.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
280 lines
10 KiB
Rust
280 lines
10 KiB
Rust
//! Generator of `src/cohere2_moe_reference.txt`, the stored reference
|
|
//! the test `cohere2_moe::tests::stored_reference` checks.
|
|
//!
|
|
//! It builds the test's three synthetic Cohere2-MoE checkpoints (the
|
|
//! same weights, byte for byte), runs the test's scenario through
|
|
//! `forward_prompt` and `forward_token`, and prints one line per output
|
|
//! in the file's format, plus a `math-library` line naming this
|
|
//! system's math library, on which the bits depend as well as on the
|
|
//! code.
|
|
//!
|
|
//! The stored lines are this program's output at commit 30c4cf1ee (main
|
|
//! before the CPU passes), whose `forward_prompt` and `forward_token`
|
|
//! are the code the CPU passes must match. It uses only API that commit
|
|
//! has, so to add a platform:
|
|
//!
|
|
//! ```bash
|
|
//! git worktree add --detach /tmp/base 30c4cf1ee
|
|
//! cp examples/cohere2_reference.rs /tmp/base/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/
|
|
//! cd /tmp/base/jouleclaw/jouleclaw-rs
|
|
//! cargo run -q -p jouleclaw-loader-gguf --example cohere2_reference
|
|
//! cargo run -q -p jouleclaw-loader-gguf --features rayon --example cohere2_reference # must print the same
|
|
//! ```
|
|
//!
|
|
//! and append its lines to the file. Run at a later commit, it prints
|
|
//! that commit's lines, which the test requires to be the same.
|
|
|
|
use jouleclaw_loader_gguf::cohere2_moe::{Cohere2MoeConfig, Cohere2MoeModel};
|
|
use jouleclaw_loader_gguf::safetensors_reader::read_safetensors_file;
|
|
|
|
fn tiny_config() -> Cohere2MoeConfig {
|
|
Cohere2MoeConfig {
|
|
hidden_size: 8,
|
|
num_hidden_layers: 4,
|
|
vocab_size: 13,
|
|
layer_types: vec![
|
|
"full_attention".into(),
|
|
"sliding_attention".into(),
|
|
"sliding_attention".into(),
|
|
"full_attention".into(),
|
|
],
|
|
num_attention_heads: 2,
|
|
num_key_value_heads: 1,
|
|
head_dim: 4,
|
|
rope_theta: 10_000.0,
|
|
sliding_window: 3,
|
|
num_experts: 4,
|
|
num_experts_per_tok: 2,
|
|
intermediate_size: 6,
|
|
expert_selection_fn: "sigmoid".into(),
|
|
norm_topk_prob: false,
|
|
first_k_dense_replace: 1,
|
|
prefix_dense_intermediate_size: Some(10),
|
|
prefix_dense_sliding_window_pattern: 1,
|
|
rms_norm_eps: 1e-6,
|
|
logit_scale: 0.5,
|
|
}
|
|
}
|
|
|
|
fn medium_config() -> Cohere2MoeConfig {
|
|
Cohere2MoeConfig {
|
|
hidden_size: 64,
|
|
num_hidden_layers: 4,
|
|
vocab_size: 37,
|
|
layer_types: vec![
|
|
"full_attention".into(),
|
|
"sliding_attention".into(),
|
|
"sliding_attention".into(),
|
|
"full_attention".into(),
|
|
],
|
|
num_attention_heads: 4,
|
|
num_key_value_heads: 2,
|
|
head_dim: 32,
|
|
rope_theta: 10_000.0,
|
|
sliding_window: 20,
|
|
num_experts: 8,
|
|
num_experts_per_tok: 3,
|
|
intermediate_size: 24,
|
|
expert_selection_fn: "sigmoid".into(),
|
|
norm_topk_prob: false,
|
|
first_k_dense_replace: 1,
|
|
prefix_dense_intermediate_size: Some(40),
|
|
prefix_dense_sliding_window_pattern: 1,
|
|
rms_norm_eps: 1e-6,
|
|
logit_scale: 0.5,
|
|
}
|
|
}
|
|
|
|
fn wide_config() -> Cohere2MoeConfig {
|
|
Cohere2MoeConfig {
|
|
hidden_size: 256,
|
|
num_hidden_layers: 3,
|
|
vocab_size: 1031,
|
|
layer_types: vec![
|
|
"full_attention".into(),
|
|
"sliding_attention".into(),
|
|
"full_attention".into(),
|
|
],
|
|
num_attention_heads: 8,
|
|
num_key_value_heads: 2,
|
|
head_dim: 32,
|
|
rope_theta: 10_000.0,
|
|
sliding_window: 16,
|
|
num_experts: 8,
|
|
num_experts_per_tok: 2,
|
|
intermediate_size: 96,
|
|
expert_selection_fn: "sigmoid".into(),
|
|
norm_topk_prob: false,
|
|
first_k_dense_replace: 1,
|
|
prefix_dense_intermediate_size: Some(2048),
|
|
prefix_dense_sliding_window_pattern: 1,
|
|
rms_norm_eps: 1e-6,
|
|
logit_scale: 0.5,
|
|
}
|
|
}
|
|
|
|
/// A scenario: `prompt` tokens prefilled, `decode` steps, `extend`
|
|
/// tokens fed into the live state, two more steps.
|
|
struct Shape {
|
|
prompt: u32,
|
|
decode: u32,
|
|
extend: u32,
|
|
}
|
|
|
|
/// Deterministic weights for a tensor, from its name.
|
|
fn fill(name: &str, n: usize) -> Vec<f32> {
|
|
let mut h = 1469598103934665603u64;
|
|
for b in name.bytes() {
|
|
h = (h ^ b as u64).wrapping_mul(1099511628211);
|
|
}
|
|
(0..n)
|
|
.map(|i| {
|
|
let mut x = h.wrapping_add((i as u64).wrapping_mul(0x9E3779B97F4A7C15));
|
|
x ^= x >> 33;
|
|
x = x.wrapping_mul(0xFF51AFD7ED558CCD);
|
|
x ^= x >> 33;
|
|
((x % 2000) as f32 / 1000.0 - 1.0) * 0.4
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
/// The checkpoint for `cfg` as a BF16 safetensors file at `path`.
|
|
fn build_checkpoint(path: &std::path::Path, cfg: &Cohere2MoeConfig) {
|
|
let h = cfg.hidden_size as u64;
|
|
let mut tensors: Vec<(String, Vec<u64>, Vec<f32>)> = Vec::new();
|
|
let mut push = |name: String, shape: Vec<u64>| {
|
|
let n: u64 = shape.iter().product();
|
|
let vals = fill(&name, n as usize);
|
|
tensors.push((name, shape, vals));
|
|
};
|
|
push("model.embed_tokens.weight".into(), vec![cfg.vocab_size as u64, h]);
|
|
push("model.norm.weight".into(), vec![h]);
|
|
let q_dim = (cfg.num_attention_heads * cfg.head_dim) as u64;
|
|
let kv_dim = (cfg.num_key_value_heads * cfg.head_dim) as u64;
|
|
for i in 0..cfg.num_hidden_layers {
|
|
let p = format!("model.layers.{i}");
|
|
push(format!("{p}.input_layernorm.weight"), vec![h]);
|
|
push(format!("{p}.self_attn.q_proj.weight"), vec![q_dim, h]);
|
|
push(format!("{p}.self_attn.k_proj.weight"), vec![kv_dim, h]);
|
|
push(format!("{p}.self_attn.v_proj.weight"), vec![kv_dim, h]);
|
|
push(format!("{p}.self_attn.o_proj.weight"), vec![h, q_dim]);
|
|
if i < cfg.first_k_dense_replace {
|
|
let pf = cfg.prefix_dense_intermediate_size.unwrap() as u64;
|
|
push(format!("{p}.mlp.gate_proj.weight"), vec![pf, h]);
|
|
push(format!("{p}.mlp.up_proj.weight"), vec![pf, h]);
|
|
push(format!("{p}.mlp.down_proj.weight"), vec![h, pf]);
|
|
} else {
|
|
push(format!("{p}.mlp.gate.weight"), vec![cfg.num_experts as u64, h]);
|
|
for e in 0..cfg.num_experts {
|
|
let f = cfg.intermediate_size as u64;
|
|
push(format!("{p}.mlp.experts.{e}.gate_proj.weight"), vec![f, h]);
|
|
push(format!("{p}.mlp.experts.{e}.up_proj.weight"), vec![f, h]);
|
|
push(format!("{p}.mlp.experts.{e}.down_proj.weight"), vec![h, f]);
|
|
}
|
|
}
|
|
}
|
|
let mut header = String::from("{");
|
|
let mut data: Vec<u8> = Vec::new();
|
|
for (i, (name, shape, vals)) in tensors.iter().enumerate() {
|
|
let start = data.len();
|
|
for &v in vals {
|
|
// BF16 by truncation, as the test's checkpoints are.
|
|
data.extend_from_slice(&((v.to_bits() >> 16) as u16).to_le_bytes());
|
|
}
|
|
let end = data.len();
|
|
if i > 0 {
|
|
header.push(',');
|
|
}
|
|
header.push_str(&format!(
|
|
"\"{name}\":{{\"dtype\":\"BF16\",\"shape\":{shape:?},\"data_offsets\":[{start},{end}]}}"
|
|
));
|
|
}
|
|
header.push('}');
|
|
let mut out = Vec::new();
|
|
out.extend_from_slice(&(header.len() as u64).to_le_bytes());
|
|
out.extend_from_slice(header.as_bytes());
|
|
out.extend_from_slice(&data);
|
|
std::fs::write(path, out).expect("write the checkpoint");
|
|
}
|
|
|
|
/// FNV-1a over every logit's bits, little-endian (north_bench's hash).
|
|
fn logits_hash(logits: &[f32]) -> u64 {
|
|
let mut h = 0xcbf2_9ce4_8422_2325u64;
|
|
for v in logits {
|
|
for b in v.to_bits().to_le_bytes() {
|
|
h = (h ^ b as u64).wrapping_mul(0x0000_0100_0000_01b3);
|
|
}
|
|
}
|
|
h
|
|
}
|
|
|
|
/// The scenario's lines for one model: `<model> <output> <hash>`.
|
|
fn scenario(name: &str, model: &Cohere2MoeModel, shape: &Shape) -> Vec<String> {
|
|
let v = model.config.vocab_size as u32;
|
|
let mut out = Vec::new();
|
|
let mut line = |what: String, l: &[f32]| out.push(format!("{name} {what} {:016x}", logits_hash(l)));
|
|
let mut state = model.new_state();
|
|
let prompt: Vec<u32> = (0..shape.prompt).map(|i| (i * 7 + 3) % v).collect();
|
|
line("prefill".into(), &model.forward_prompt(&prompt, &mut state));
|
|
for i in 0..shape.decode {
|
|
line(format!("decode{i}"), &model.forward_token((i * 5 + 2) % v, &mut state));
|
|
}
|
|
// The chunk is fed one decode step at a time (`forward_extend`,
|
|
// which batches it, gives the same logits and does not exist at
|
|
// the commit the stored lines come from).
|
|
let mut fed = Vec::new();
|
|
for i in 0..shape.extend {
|
|
fed = model.forward_token((i * 11 + 5) % v, &mut state);
|
|
}
|
|
line("extend".into(), &fed);
|
|
for i in 0..2 {
|
|
line(format!("after{i}"), &model.forward_token((i * 3 + 1) % v, &mut state));
|
|
}
|
|
out
|
|
}
|
|
|
|
/// This system's math library, as `math-library` names it.
|
|
fn math_library() -> String {
|
|
let run = |cmd: &str, args: &[&str]| {
|
|
std::process::Command::new(cmd)
|
|
.args(args)
|
|
.output()
|
|
.ok()
|
|
.map(|o| String::from_utf8_lossy(&o.stdout).lines().next().unwrap_or("").trim().to_string())
|
|
.unwrap_or_default()
|
|
};
|
|
match std::env::consts::OS {
|
|
"macos" => format!(
|
|
"the system math library (libsystem_m) of macOS {} ({})",
|
|
run("sw_vers", &["-productVersion"]),
|
|
run("sw_vers", &["-buildVersion"])
|
|
),
|
|
"linux" => {
|
|
// `ldd (Ubuntu GLIBC 2.39-0ubuntu8.9) 2.39`: the C library's name and version.
|
|
let l = run("ldd", &["--version"]);
|
|
let libc = l.split_once('(').and_then(|(_, r)| r.rsplit_once(')')).map_or(l.as_str(), |(n, _)| n);
|
|
format!("the libm of {libc}")
|
|
}
|
|
os => format!("the system math library of {os}"),
|
|
}
|
|
}
|
|
|
|
fn main() {
|
|
let platform = format!("{}-{}", std::env::consts::OS, std::env::consts::ARCH);
|
|
println!("math-library {platform} {}", math_library());
|
|
let models = [
|
|
("tiny", tiny_config(), Shape { prompt: 5, decode: 4, extend: 4 }),
|
|
("medium", medium_config(), Shape { prompt: 45, decode: 6, extend: 25 }),
|
|
("wide", wide_config(), Shape { prompt: 24, decode: 4, extend: 20 }),
|
|
];
|
|
for (name, cfg, shape) in models {
|
|
let path = std::env::temp_dir().join(format!("cohere2-reference-{name}-{}.safetensors", std::process::id()));
|
|
build_checkpoint(&path, &cfg);
|
|
let st = read_safetensors_file(&path).expect("read the checkpoint");
|
|
let model = Cohere2MoeModel::bind(cfg, st).expect("bind the checkpoint");
|
|
let _ = std::fs::remove_file(&path);
|
|
for line in scenario(name, &model, &shape) {
|
|
println!("{platform} {line}");
|
|
}
|
|
}
|
|
}
|