1021 lines
32 KiB
Rust
1021 lines
32 KiB
Rust
use super::*;
|
||
|
||
pub(crate) fn validate_model_artifact(
|
||
path: &Path,
|
||
expected: ModelChoice,
|
||
support: bool,
|
||
) -> Result<(), String> {
|
||
if support {
|
||
let model = Gguf::open(path)?;
|
||
validate_dspark(&model, &FLASH)
|
||
} else {
|
||
let model = Model::open_main(path, expected)?;
|
||
let summary = model.summary();
|
||
if summary.tensor_count == 0
|
||
|| model
|
||
.render_prompt("", "", ReasoningMode::Direct)
|
||
.is_empty()
|
||
{
|
||
return Err("model intake produced an empty tensor directory or prompt".into());
|
||
}
|
||
let _ = model.tokenize("");
|
||
let eos = model.eos_token();
|
||
let _ = model.token_bytes(eos);
|
||
if !model.is_stop_token(eos) {
|
||
return Err("tokenizer EOS marker is not a stop token".into());
|
||
}
|
||
model.tensor_data("token_embd.weight")?;
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
pub(super) fn validate_main(model: &Gguf, expected: ModelChoice) -> Result<Shape, String> {
|
||
let family = if model.bytes("general.architecture").ok() == Some(b"glm-dsa") {
|
||
ModelFamily::Glm
|
||
} else {
|
||
ModelFamily::DeepSeek
|
||
};
|
||
let shape = match family {
|
||
ModelFamily::Glm => GLM,
|
||
ModelFamily::DeepSeek => match model.u32("deepseek4.block_count")? {
|
||
43 => FLASH,
|
||
61 => PRO,
|
||
layers => return Err(format!("unsupported DeepSeek layer count: {layers}")),
|
||
},
|
||
};
|
||
if shape.model != expected {
|
||
return Err(format!(
|
||
"{} contains {}, but the selected model is {expected}",
|
||
model.path().display(),
|
||
shape.model
|
||
));
|
||
}
|
||
validate_metadata(model, &shape)?;
|
||
validate_tensors(model, &shape)?;
|
||
Ok(shape)
|
||
}
|
||
|
||
fn validate_metadata(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||
let prefix = if shape.family == ModelFamily::Glm {
|
||
"glm-dsa"
|
||
} else {
|
||
"deepseek4"
|
||
};
|
||
for (key, expected) in [
|
||
("block_count", u64::from(shape.layers)),
|
||
("embedding_length", shape.embd),
|
||
("vocab_size", shape.vocab),
|
||
("attention.head_count", shape.heads),
|
||
("attention.head_count_kv", shape.head_kv),
|
||
("attention.key_length", shape.head_dim),
|
||
("attention.value_length", shape.value_dim),
|
||
("rope.dimension_count", shape.rot),
|
||
("attention.q_lora_rank", shape.lora_q),
|
||
("expert_count", shape.experts),
|
||
("expert_used_count", shape.experts_used),
|
||
("expert_feed_forward_length", shape.ff_expert),
|
||
("expert_shared_count", shape.expert_shared),
|
||
("attention.indexer.head_count", shape.indexer_heads),
|
||
("attention.indexer.key_length", shape.indexer_head_dim),
|
||
("attention.indexer.top_k", shape.indexer_top_k),
|
||
] {
|
||
expect_u64(model, &format!("{prefix}.{key}"), expected)?;
|
||
}
|
||
expect_float(
|
||
model,
|
||
&format!("{prefix}.attention.layer_norm_rms_epsilon"),
|
||
shape.rms_epsilon,
|
||
)?;
|
||
expect_float(
|
||
model,
|
||
&format!("{prefix}.expert_weights_scale"),
|
||
shape.expert_weight_scale,
|
||
)?;
|
||
if !model.boolean(&format!("{prefix}.expert_weights_norm"))? {
|
||
return Err(format!("{prefix}.expert_weights_norm must be true"));
|
||
}
|
||
expect_float(model, &format!("{prefix}.rope.freq_base"), shape.rope_base)?;
|
||
|
||
if shape.family == ModelFamily::Glm {
|
||
for (key, expected) in [
|
||
("context_length", shape.original_context),
|
||
("feed_forward_length", shape.ff_dense),
|
||
("attention.kv_lora_rank", shape.kv_lora),
|
||
("attention.key_length_mla", shape.key_mla),
|
||
("attention.value_length_mla", shape.value_mla),
|
||
("expert_group_count", 1),
|
||
("expert_group_used_count", 1),
|
||
("expert_gating_func", 2),
|
||
("leading_dense_block_count", u64::from(shape.leading_dense)),
|
||
("nextn_predict_layers", u64::from(shape.nextn)),
|
||
] {
|
||
expect_u64(model, &format!("{prefix}.{key}"), expected)?;
|
||
}
|
||
return Ok(());
|
||
}
|
||
|
||
for (key, expected) in [
|
||
("attention.output_group_count", shape.out_groups),
|
||
("attention.output_lora_rank", shape.lora_o),
|
||
("hash_layer_count", u64::from(shape.hash_layers)),
|
||
("attention.sliding_window", shape.sliding_window),
|
||
("hyper_connection.count", shape.hc),
|
||
("hyper_connection.sinkhorn_iterations", shape.hc_sinkhorn),
|
||
] {
|
||
expect_u64(model, &format!("{prefix}.{key}"), expected)?;
|
||
}
|
||
for key in ["expert_group_count", "expert_group_used_count"] {
|
||
if let Some(Value::U32(value)) = model.metadata.get(&format!("{prefix}.{key}"))
|
||
&& *value != 0
|
||
{
|
||
return Err(format!("{prefix}.{key} must be zero"));
|
||
}
|
||
}
|
||
for (key, expected) in [
|
||
("hyper_connection.epsilon", shape.hc_epsilon),
|
||
(
|
||
"attention.compress_rope_freq_base",
|
||
shape.compress_rope_base,
|
||
),
|
||
] {
|
||
expect_float(model, &format!("{prefix}.{key}"), expected)?;
|
||
}
|
||
for (key, expected) in [
|
||
("rope.scaling.factor", shape.rope_scale),
|
||
("rope.scaling.yarn_beta_fast", shape.rope_beta_fast),
|
||
("rope.scaling.yarn_beta_slow", shape.rope_beta_slow),
|
||
] {
|
||
if model.metadata.contains_key(&format!("{prefix}.{key}")) {
|
||
expect_float(model, &format!("{prefix}.{key}"), expected)?;
|
||
}
|
||
}
|
||
if model
|
||
.metadata
|
||
.contains_key("deepseek4.rope.scaling.original_context_length")
|
||
{
|
||
expect_u64(
|
||
model,
|
||
"deepseek4.rope.scaling.original_context_length",
|
||
shape.original_context,
|
||
)?;
|
||
}
|
||
let ratios = model.u32s("deepseek4.attention.compress_ratios")?;
|
||
let clamps = model.f32s("deepseek4.swiglu_clamp_exp")?;
|
||
if ratios.len() < shape.layers as usize || clamps.len() < shape.layers as usize {
|
||
return Err("DeepSeek per-layer metadata is shorter than the layer count".into());
|
||
}
|
||
for layer in 0..shape.layers as usize {
|
||
let expected = compression_ratio(shape, layer as u32);
|
||
if ratios[layer] != expected {
|
||
return Err(format!(
|
||
"layer {layer} compression ratio is {}, expected {expected}",
|
||
ratios[layer]
|
||
));
|
||
}
|
||
if !float_eq(clamps[layer], shape.swiglu_clamp) {
|
||
return Err(format!("layer {layer} has an invalid SwiGLU clamp"));
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn validate_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||
match shape.family {
|
||
ModelFamily::DeepSeek => validate_deepseek_tensors(model, shape),
|
||
ModelFamily::Glm => validate_glm_tensors(model, shape),
|
||
}
|
||
}
|
||
|
||
fn validate_deepseek_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||
let hc_dim = shape.embd * shape.hc;
|
||
let hc_mix = 2 * shape.hc + shape.hc * shape.hc;
|
||
let q_dim = shape.heads * shape.head_dim;
|
||
let output_low = shape.out_groups * shape.lora_o;
|
||
expect(
|
||
model,
|
||
"token_embd.weight",
|
||
&[F16],
|
||
&[shape.embd, shape.vocab],
|
||
)?;
|
||
expect(model, "output_hc_base.weight", &[F32], &[shape.hc])?;
|
||
expect(model, "output_hc_fn.weight", &[F16], &[hc_dim, shape.hc])?;
|
||
expect(model, "output_hc_scale.weight", &[F32], &[1])?;
|
||
expect(model, "output_norm.weight", &[F32], &[shape.embd])?;
|
||
expect(model, "output.weight", DENSE, &[shape.embd, shape.vocab])?;
|
||
|
||
for layer in 0..shape.layers {
|
||
let name = |suffix: &str| format!("blk.{layer}.{suffix}");
|
||
expect(model, &name("hc_attn_fn.weight"), &[F16], &[hc_dim, hc_mix])?;
|
||
expect(model, &name("hc_attn_scale.weight"), &[F32], &[3])?;
|
||
expect(model, &name("hc_attn_base.weight"), &[F32], &[hc_mix])?;
|
||
expect(model, &name("attn_norm.weight"), &[F32], &[shape.embd])?;
|
||
expect(
|
||
model,
|
||
&name("attn_q_a.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.lora_q],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_q_a_norm.weight"),
|
||
&[F32],
|
||
&[shape.lora_q],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_q_b.weight"),
|
||
DENSE,
|
||
&[shape.lora_q, q_dim],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_kv.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.head_dim],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_kv_a_norm.weight"),
|
||
&[F32],
|
||
&[shape.head_dim],
|
||
)?;
|
||
expect(model, &name("attn_sinks.weight"), &[F32], &[shape.heads])?;
|
||
expect(
|
||
model,
|
||
&name("attn_output_a.weight"),
|
||
DENSE,
|
||
&[
|
||
shape.head_dim * (shape.heads / shape.out_groups),
|
||
output_low,
|
||
],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_output_b.weight"),
|
||
DENSE,
|
||
&[output_low, shape.embd],
|
||
)?;
|
||
|
||
let ratio = compression_ratio(shape, layer);
|
||
if ratio != 0 {
|
||
let compression_width = if ratio == 4 { 2 } else { 1 } * shape.head_dim;
|
||
expect(
|
||
model,
|
||
&name("attn_compressor_ape.weight"),
|
||
&[F16],
|
||
&[compression_width, u64::from(ratio)],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_compressor_kv.weight"),
|
||
&[F16],
|
||
&[shape.embd, compression_width],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_compressor_gate.weight"),
|
||
&[F16],
|
||
&[shape.embd, compression_width],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_compressor_norm.weight"),
|
||
&[F32],
|
||
&[shape.head_dim],
|
||
)?;
|
||
}
|
||
if ratio == 4 {
|
||
let index_q = shape.indexer_heads * shape.indexer_head_dim;
|
||
let index_width = 2 * shape.indexer_head_dim;
|
||
expect(
|
||
model,
|
||
&name("indexer.attn_q_b.weight"),
|
||
&[F16, Q8_0],
|
||
&[shape.lora_q, index_q],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer.proj.weight"),
|
||
&[F16],
|
||
&[shape.embd, shape.indexer_heads],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer_compressor_ape.weight"),
|
||
&[F16],
|
||
&[index_width, 4],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer_compressor_kv.weight"),
|
||
&[F16],
|
||
&[shape.embd, index_width],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer_compressor_gate.weight"),
|
||
&[F16],
|
||
&[shape.embd, index_width],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer_compressor_norm.weight"),
|
||
&[F32],
|
||
&[shape.indexer_head_dim],
|
||
)?;
|
||
}
|
||
expect(model, &name("hc_ffn_fn.weight"), &[F16], &[hc_dim, hc_mix])?;
|
||
expect(model, &name("hc_ffn_scale.weight"), &[F32], &[3])?;
|
||
expect(model, &name("hc_ffn_base.weight"), &[F32], &[hc_mix])?;
|
||
expect(model, &name("ffn_norm.weight"), &[F32], &[shape.embd])?;
|
||
expect(
|
||
model,
|
||
&name("ffn_gate_inp.weight"),
|
||
&[F16],
|
||
&[shape.embd, shape.experts],
|
||
)?;
|
||
expect_optional(model, &name("exp_probs_b.bias"), &[F32], &[shape.experts])?;
|
||
expect(
|
||
model,
|
||
&name("ffn_gate_exps.weight"),
|
||
ROUTED,
|
||
&[shape.embd, shape.ff_expert, shape.experts],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_up_exps.weight"),
|
||
ROUTED,
|
||
&[shape.embd, shape.ff_expert, shape.experts],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_down_exps.weight"),
|
||
ROUTED,
|
||
&[shape.ff_expert, shape.embd, shape.experts],
|
||
)?;
|
||
same_type(
|
||
model,
|
||
&name("ffn_gate_exps.weight"),
|
||
&name("ffn_up_exps.weight"),
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_gate_shexp.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.ff_expert],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_up_shexp.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.ff_expert],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_down_shexp.weight"),
|
||
DENSE,
|
||
&[shape.ff_expert, shape.embd],
|
||
)?;
|
||
if layer < shape.hash_layers {
|
||
expect(
|
||
model,
|
||
&name("ffn_gate_tid2eid.weight"),
|
||
&[I32],
|
||
&[shape.experts_used, shape.vocab],
|
||
)?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn validate_glm_tensors(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||
let q_dim = shape.heads * shape.key_mla;
|
||
let q_nope = shape.key_mla - shape.rot;
|
||
let index_q = shape.indexer_heads * shape.indexer_head_dim;
|
||
expect(
|
||
model,
|
||
"token_embd.weight",
|
||
DENSE,
|
||
&[shape.embd, shape.vocab],
|
||
)?;
|
||
expect(model, "output_norm.weight", &[F32], &[shape.embd])?;
|
||
expect(model, "output.weight", DENSE, &[shape.embd, shape.vocab])?;
|
||
for layer in 0..shape.layers {
|
||
let name = |suffix: &str| format!("blk.{layer}.{suffix}");
|
||
expect(model, &name("attn_norm.weight"), &[F32], &[shape.embd])?;
|
||
expect(
|
||
model,
|
||
&name("attn_q_a.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.lora_q],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_q_a_norm.weight"),
|
||
&[F32],
|
||
&[shape.lora_q],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_q_b.weight"),
|
||
DENSE,
|
||
&[shape.lora_q, q_dim],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_kv_a_mqa.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.head_dim],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_kv_a_norm.weight"),
|
||
&[F32],
|
||
&[shape.kv_lora],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_k_b.weight"),
|
||
DENSE,
|
||
&[q_nope, shape.kv_lora, shape.heads],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_v_b.weight"),
|
||
DENSE,
|
||
&[shape.kv_lora, shape.value_mla, shape.heads],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("attn_output.weight"),
|
||
DENSE,
|
||
&[shape.heads * shape.value_mla, shape.embd],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer.attn_k.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.indexer_head_dim],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer.attn_q_b.weight"),
|
||
DENSE,
|
||
&[shape.lora_q, index_q],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer.k_norm.weight"),
|
||
&[F32],
|
||
&[shape.indexer_head_dim],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer.k_norm.bias"),
|
||
&[F32],
|
||
&[shape.indexer_head_dim],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("indexer.proj.weight"),
|
||
&[F32],
|
||
&[shape.embd, shape.indexer_heads],
|
||
)?;
|
||
expect(model, &name("ffn_norm.weight"), &[F32], &[shape.embd])?;
|
||
if layer < shape.leading_dense {
|
||
expect(
|
||
model,
|
||
&name("ffn_gate.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.ff_dense],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_up.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.ff_dense],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_down.weight"),
|
||
DENSE,
|
||
&[shape.ff_dense, shape.embd],
|
||
)?;
|
||
} else {
|
||
expect(
|
||
model,
|
||
&name("ffn_gate_inp.weight"),
|
||
&[F32],
|
||
&[shape.embd, shape.experts],
|
||
)?;
|
||
expect(model, &name("exp_probs_b.bias"), &[F32], &[shape.experts])?;
|
||
expect(
|
||
model,
|
||
&name("ffn_gate_exps.weight"),
|
||
ROUTED,
|
||
&[shape.embd, shape.ff_expert, shape.experts],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_up_exps.weight"),
|
||
ROUTED,
|
||
&[shape.embd, shape.ff_expert, shape.experts],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_down_exps.weight"),
|
||
ROUTED,
|
||
&[shape.ff_expert, shape.embd, shape.experts],
|
||
)?;
|
||
same_type(
|
||
model,
|
||
&name("ffn_gate_exps.weight"),
|
||
&name("ffn_up_exps.weight"),
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_gate_shexp.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.ff_expert],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_up_shexp.weight"),
|
||
DENSE,
|
||
&[shape.embd, shape.ff_expert],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&name("ffn_down_shexp.weight"),
|
||
DENSE,
|
||
&[shape.ff_expert, shape.embd],
|
||
)?;
|
||
}
|
||
if layer + shape.nextn >= shape.layers {
|
||
expect(
|
||
model,
|
||
&name("nextn.eh_proj.weight"),
|
||
DENSE,
|
||
&[2 * shape.embd, shape.embd],
|
||
)?;
|
||
expect(model, &name("nextn.enorm.weight"), &[F32], &[shape.embd])?;
|
||
expect(model, &name("nextn.hnorm.weight"), &[F32], &[shape.embd])?;
|
||
expect(
|
||
model,
|
||
&name("nextn.shared_head_norm.weight"),
|
||
&[F32],
|
||
&[shape.embd],
|
||
)?;
|
||
}
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
pub(super) fn validate_dspark(model: &Gguf, shape: &Shape) -> Result<(), String> {
|
||
if shape.model != ModelChoice::DeepSeekV4Flash {
|
||
return Err("DSpark support is available only for DeepSeek V4 Flash".into());
|
||
}
|
||
let block_size = first_u32(
|
||
model,
|
||
&[
|
||
"deepseek4.dspark.block_size",
|
||
"deepseek4.dspark_block_size",
|
||
"dspark.block_size",
|
||
],
|
||
)?;
|
||
let markov_rank = first_u32(
|
||
model,
|
||
&[
|
||
"deepseek4.dspark.markov_rank",
|
||
"deepseek4.dspark_markov_rank",
|
||
"dspark.markov_rank",
|
||
],
|
||
)?;
|
||
let noise_token = first_u32(
|
||
model,
|
||
&[
|
||
"deepseek4.dspark.noise_token_id",
|
||
"deepseek4.dspark_noise_token_id",
|
||
"dspark.noise_token_id",
|
||
],
|
||
)?;
|
||
let targets = first_u32s(
|
||
model,
|
||
&[
|
||
"deepseek4.dspark.target_layer_ids",
|
||
"deepseek4.dspark_target_layer_ids",
|
||
"dspark.target_layer_ids",
|
||
],
|
||
)?;
|
||
if !(1..=16).contains(&block_size) || markov_rank == 0 || noise_token >= shape.vocab as u32 {
|
||
return Err("invalid DSpark block, Markov, or noise-token metadata".into());
|
||
}
|
||
if targets.is_empty()
|
||
|| targets.len() > 8
|
||
|| targets.windows(2).any(|pair| pair[0] >= pair[1])
|
||
|| targets.iter().any(|layer| *layer >= shape.layers)
|
||
{
|
||
return Err("invalid DSpark target-layer metadata".into());
|
||
}
|
||
let stages = model
|
||
.tensors
|
||
.keys()
|
||
.filter_map(|name| {
|
||
name.strip_prefix("mtp.")?
|
||
.split('.')
|
||
.next()?
|
||
.parse::<u32>()
|
||
.ok()
|
||
})
|
||
.max()
|
||
.map_or(0, |stage| stage + 1);
|
||
if !(1..=8).contains(&stages) {
|
||
return Err(format!("invalid DSpark stage count: {stages}"));
|
||
}
|
||
for stage in 0..stages {
|
||
validate_dspark_block(model, shape, stage)?;
|
||
if stage == 0 {
|
||
expect(
|
||
model,
|
||
&format!("mtp.{stage}.main_proj.weight"),
|
||
DSPARK_DENSE,
|
||
&[targets.len() as u64 * shape.embd, shape.embd],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&format!("mtp.{stage}.main_norm.weight"),
|
||
&[F32],
|
||
&[shape.embd],
|
||
)?;
|
||
}
|
||
}
|
||
let prefix = format!("mtp.{}", stages - 1);
|
||
expect(
|
||
model,
|
||
&format!("{prefix}.norm.weight"),
|
||
&[F32],
|
||
&[shape.embd],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&format!("{prefix}.hc_head_base.weight"),
|
||
&[F32],
|
||
&[shape.hc],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&format!("{prefix}.hc_head_fn.weight"),
|
||
PLAIN,
|
||
&[shape.embd * shape.hc, shape.hc],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&format!("{prefix}.hc_head_scale.weight"),
|
||
&[F32],
|
||
&[1],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&format!("{prefix}.markov_head.markov_w1.weight"),
|
||
DSPARK_DENSE,
|
||
&[u64::from(markov_rank), shape.vocab],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&format!("{prefix}.markov_head.markov_w2.weight"),
|
||
DSPARK_DENSE,
|
||
&[u64::from(markov_rank), shape.vocab],
|
||
)?;
|
||
expect(
|
||
model,
|
||
&format!("{prefix}.confidence_head.proj.weight"),
|
||
DSPARK_DENSE,
|
||
&[shape.embd + u64::from(markov_rank), 1],
|
||
)?;
|
||
Ok(())
|
||
}
|
||
|
||
fn validate_dspark_block(model: &Gguf, shape: &Shape, stage: u32) -> Result<(), String> {
|
||
let hc_dim = shape.embd * shape.hc;
|
||
let hc_mix = 2 * shape.hc + shape.hc * shape.hc;
|
||
let q_dim = shape.heads * shape.head_dim;
|
||
let output_low = shape.out_groups * shape.lora_o;
|
||
let name = |suffix: &str| format!("mtp.{stage}.{suffix}");
|
||
for (suffix, types, dims) in [
|
||
("hc_attn_fn.weight", PLAIN, vec![hc_dim, hc_mix]),
|
||
("hc_attn_scale.weight", &[F32][..], vec![3]),
|
||
("hc_attn_base.weight", &[F32][..], vec![hc_mix]),
|
||
("attn_norm.weight", &[F32][..], vec![shape.embd]),
|
||
(
|
||
"attn_q_a.weight",
|
||
DSPARK_DENSE,
|
||
vec![shape.embd, shape.lora_q],
|
||
),
|
||
("attn_q_a_norm.weight", &[F32][..], vec![shape.lora_q]),
|
||
("attn_q_b.weight", DSPARK_DENSE, vec![shape.lora_q, q_dim]),
|
||
(
|
||
"attn_kv.weight",
|
||
DSPARK_DENSE,
|
||
vec![shape.embd, shape.head_dim],
|
||
),
|
||
("attn_kv_a_norm.weight", &[F32][..], vec![shape.head_dim]),
|
||
("attn_sinks.weight", &[F32][..], vec![shape.heads]),
|
||
(
|
||
"attn_output_a.weight",
|
||
DSPARK_DENSE,
|
||
vec![
|
||
shape.head_dim * (shape.heads / shape.out_groups),
|
||
output_low,
|
||
],
|
||
),
|
||
(
|
||
"attn_output_b.weight",
|
||
DSPARK_DENSE,
|
||
vec![output_low, shape.embd],
|
||
),
|
||
("hc_ffn_fn.weight", PLAIN, vec![hc_dim, hc_mix]),
|
||
("hc_ffn_scale.weight", &[F32][..], vec![3]),
|
||
("hc_ffn_base.weight", &[F32][..], vec![hc_mix]),
|
||
("ffn_norm.weight", &[F32][..], vec![shape.embd]),
|
||
(
|
||
"ffn_gate_inp.weight",
|
||
DSPARK_DENSE,
|
||
vec![shape.embd, shape.experts],
|
||
),
|
||
("exp_probs_b.bias", &[F32][..], vec![shape.experts]),
|
||
(
|
||
"ffn_gate_exps.weight",
|
||
ROUTED,
|
||
vec![shape.embd, shape.ff_expert, shape.experts],
|
||
),
|
||
(
|
||
"ffn_up_exps.weight",
|
||
ROUTED,
|
||
vec![shape.embd, shape.ff_expert, shape.experts],
|
||
),
|
||
(
|
||
"ffn_down_exps.weight",
|
||
ROUTED,
|
||
vec![shape.ff_expert, shape.embd, shape.experts],
|
||
),
|
||
(
|
||
"ffn_gate_shexp.weight",
|
||
DSPARK_DENSE,
|
||
vec![shape.embd, shape.ff_expert],
|
||
),
|
||
(
|
||
"ffn_up_shexp.weight",
|
||
DSPARK_DENSE,
|
||
vec![shape.embd, shape.ff_expert],
|
||
),
|
||
(
|
||
"ffn_down_shexp.weight",
|
||
DSPARK_DENSE,
|
||
vec![shape.ff_expert, shape.embd],
|
||
),
|
||
] {
|
||
expect(model, &name(suffix), types, &dims)?;
|
||
}
|
||
same_type(
|
||
model,
|
||
&name("ffn_gate_exps.weight"),
|
||
&name("ffn_up_exps.weight"),
|
||
)
|
||
}
|
||
|
||
fn expect(model: &Gguf, name: &str, types: &[u32], dims: &[u64]) -> Result<(), String> {
|
||
validate_tensor(name, model.tensor(name)?, types, dims)
|
||
}
|
||
|
||
fn expect_optional(model: &Gguf, name: &str, types: &[u32], dims: &[u64]) -> Result<(), String> {
|
||
match model.tensors.get(name) {
|
||
Some(tensor) => validate_tensor(name, tensor, types, dims),
|
||
None => Ok(()),
|
||
}
|
||
}
|
||
|
||
fn validate_tensor(name: &str, tensor: &Tensor, types: &[u32], dims: &[u64]) -> Result<(), String> {
|
||
if !types.contains(&tensor.kind) {
|
||
return Err(format!(
|
||
"tensor {name} has unsupported type {}",
|
||
tensor.kind
|
||
));
|
||
}
|
||
if tensor.dims != dims {
|
||
return Err(format!(
|
||
"tensor {name} has dimensions {:?}, expected {dims:?}",
|
||
tensor.dims
|
||
));
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
fn same_type(model: &Gguf, first: &str, second: &str) -> Result<(), String> {
|
||
if model.tensor(first)?.kind != model.tensor(second)?.kind {
|
||
Err(format!(
|
||
"tensors {first} and {second} use different quantizations"
|
||
))
|
||
} else {
|
||
Ok(())
|
||
}
|
||
}
|
||
|
||
fn expect_u64(model: &Gguf, key: &str, expected: u64) -> Result<(), String> {
|
||
let actual = model.u64(key)?;
|
||
if actual == expected {
|
||
Ok(())
|
||
} else {
|
||
Err(format!("{key} is {actual}, expected {expected}"))
|
||
}
|
||
}
|
||
|
||
fn expect_float(model: &Gguf, key: &str, expected: f32) -> Result<(), String> {
|
||
let actual = model.f32(key)?;
|
||
if float_eq(actual, expected) {
|
||
Ok(())
|
||
} else {
|
||
Err(format!("{key} is {actual}, expected {expected}"))
|
||
}
|
||
}
|
||
|
||
fn float_eq(actual: f32, expected: f32) -> bool {
|
||
actual.is_finite() && (actual - expected).abs() <= expected.abs().max(1.0) * 1.0e-6
|
||
}
|
||
|
||
fn compression_ratio(shape: &Shape, layer: u32) -> u32 {
|
||
match shape.model {
|
||
ModelChoice::DeepSeekV4Flash if layer < 2 => 0,
|
||
ModelChoice::DeepSeekV4Pro if layer < 2 => 128,
|
||
ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Pro if layer.is_multiple_of(2) => 4,
|
||
ModelChoice::DeepSeekV4Flash | ModelChoice::DeepSeekV4Pro => 128,
|
||
ModelChoice::Glm52 => 0,
|
||
}
|
||
}
|
||
|
||
fn first_u32(model: &Gguf, keys: &[&str]) -> Result<u32, String> {
|
||
keys.iter()
|
||
.find_map(|key| model.u32(key).ok())
|
||
.ok_or_else(|| format!("required DSpark metadata is missing: {}", keys[0]))
|
||
}
|
||
|
||
fn first_u32s<'a>(model: &'a Gguf, keys: &[&str]) -> Result<&'a [u32], String> {
|
||
keys.iter()
|
||
.find_map(|key| model.u32s(key).ok())
|
||
.ok_or_else(|| format!("required DSpark metadata is missing: {}", keys[0]))
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn installed_ds4_fixture_opens_and_renders_a_prompt() {
|
||
let path = Path::new(
|
||
"../ds4/gguf/DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf",
|
||
);
|
||
if !path.exists() {
|
||
return;
|
||
}
|
||
let model = Model::open_main(path, ModelChoice::DeepSeekV4Flash).unwrap();
|
||
let summary = model.summary();
|
||
assert_eq!(summary.model, ModelChoice::DeepSeekV4Flash);
|
||
assert_eq!(summary.vocabulary_size, 129_280);
|
||
assert_eq!(
|
||
model.tokenize("Hello, world! 1234\nint café = 7;\n中文テスト"),
|
||
[
|
||
19_923, 14, 2_058, 3, 223, 6_895, 22, 201, 650, 57_664, 438, 223, 25, 510, 21_134,
|
||
109_288,
|
||
]
|
||
);
|
||
assert_eq!(
|
||
model.render_prompt("", "Hello", ReasoningMode::Direct),
|
||
[0, 128_803, 19_923, 128_804, 128_822]
|
||
);
|
||
assert_eq!(model.tokenizer.tokenize_rendered("｜DSML｜").len(), 1);
|
||
assert!(model.is_stop_token_for_reasoning(128_822, ReasoningMode::Direct));
|
||
assert!(!model.is_stop_token_for_reasoning(128_822, ReasoningMode::High));
|
||
let think_start = *model
|
||
.render_prompt("", "Hello", ReasoningMode::High)
|
||
.last()
|
||
.unwrap();
|
||
assert_eq!(
|
||
model.render_conversation(
|
||
"",
|
||
&[
|
||
ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Hello".into(),
|
||
},
|
||
ChatTurn {
|
||
user: false,
|
||
tool: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Hello".into(),
|
||
},
|
||
ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Hello".into(),
|
||
},
|
||
],
|
||
ReasoningMode::Direct,
|
||
),
|
||
[
|
||
0, 128_803, 19_923, 128_804, 128_822, 19_923, 1, 128_803, 19_923, 128_804, 128_822,
|
||
]
|
||
);
|
||
assert_eq!(
|
||
model.render_continuation("Hello", ReasoningMode::Direct, false),
|
||
[1, 128_803, 19_923, 128_804, 128_822]
|
||
);
|
||
assert_eq!(
|
||
model.render_continuation("Hello", ReasoningMode::Direct, true),
|
||
[128_803, 19_923, 128_804, 128_822]
|
||
);
|
||
let tool_turn = ChatTurn {
|
||
user: false,
|
||
tool: true,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Hello".into(),
|
||
};
|
||
let wrapped_user = ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "<tool_result>Hello</tool_result>".into(),
|
||
};
|
||
assert_eq!(
|
||
model.render_conversation("", &[tool_turn], ReasoningMode::Direct),
|
||
model.render_conversation("", &[wrapped_user], ReasoningMode::Direct)
|
||
);
|
||
assert_eq!(
|
||
model.render_conversation(
|
||
"",
|
||
&[
|
||
ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Hello".into(),
|
||
},
|
||
ChatTurn {
|
||
user: false,
|
||
tool: false,
|
||
skip_previous_eos: false,
|
||
reasoning: Some("Hello".into()),
|
||
reasoning_complete: true,
|
||
content: "Hello".into(),
|
||
},
|
||
ChatTurn {
|
||
user: true,
|
||
tool: false,
|
||
skip_previous_eos: false,
|
||
reasoning: None,
|
||
reasoning_complete: true,
|
||
content: "Hello".into(),
|
||
},
|
||
],
|
||
ReasoningMode::Direct,
|
||
),
|
||
[
|
||
0,
|
||
128_803,
|
||
19_923,
|
||
128_804,
|
||
think_start,
|
||
19_923,
|
||
128_822,
|
||
19_923,
|
||
1,
|
||
128_803,
|
||
19_923,
|
||
128_804,
|
||
128_822,
|
||
]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn installed_dspark_fixture_passes_the_target_layout() {
|
||
let path = Path::new("../ds4/gguf/DeepSeek-V4-Flash-DSpark-support.gguf");
|
||
if path.exists() {
|
||
validate_model_artifact(path, ModelChoice::DeepSeekV4Flash, true).unwrap();
|
||
}
|
||
}
|
||
}
|