diff --git a/base-convert/Cargo.lock b/base-convert/Cargo.lock index e6c3d7d..67186cd 100644 --- a/base-convert/Cargo.lock +++ b/base-convert/Cargo.lock @@ -151,7 +151,7 @@ dependencies = [ [[package]] name = "base-arch" -version = "0.2.3" +version = "0.2.4" dependencies = [ "anyhow", "base-format", @@ -162,7 +162,7 @@ dependencies = [ [[package]] name = "base-awq" -version = "0.2.3" +version = "0.2.4" dependencies = [ "anyhow", "base-format", @@ -176,7 +176,7 @@ dependencies = [ [[package]] name = "base-convert" -version = "0.2.3" +version = "0.2.4" dependencies = [ "anyhow", "base-arch", @@ -199,7 +199,7 @@ dependencies = [ [[package]] name = "base-format" -version = "0.2.3" +version = "0.2.4" dependencies = [ "anyhow", "bitflags", @@ -215,13 +215,15 @@ dependencies = [ [[package]] name = "base-hub" -version = "0.2.3" +version = "0.2.4" dependencies = [ "anyhow", "base-format", "dirs 5.0.1", "hf-hub", "indicatif", + "libc", + "percent-encoding", "reqwest", "serde", "serde_json", @@ -232,7 +234,7 @@ dependencies = [ [[package]] name = "base-quant" -version = "0.2.3" +version = "0.2.4" dependencies = [ "anyhow", "base-format", @@ -244,7 +246,7 @@ dependencies = [ [[package]] name = "base-readers" -version = "0.2.3" +version = "0.2.4" dependencies = [ "anyhow", "base-format", @@ -258,7 +260,7 @@ dependencies = [ [[package]] name = "base-sign" -version = "0.2.3" +version = "0.2.4" dependencies = [ "anyhow", "base-format", diff --git a/base-convert/Cargo.toml b/base-convert/Cargo.toml index 95ccd30..cc00d55 100644 --- a/base-convert/Cargo.toml +++ b/base-convert/Cargo.toml @@ -12,7 +12,7 @@ members = [ ] [workspace.package] -version = "0.2.3" +version = "0.2.4" edition = "2021" license = "Apache-2.0" repository = "https://github.com/basecompute/baseRT" diff --git a/base-convert/crates/base-arch/src/bert.rs b/base-convert/crates/base-arch/src/bert.rs index 65c1c1c..0e153b5 100644 --- a/base-convert/crates/base-arch/src/bert.rs +++ b/base-convert/crates/base-arch/src/bert.rs @@ -97,8 +97,8 @@ impl GgufMapper for NomicBertMapper { let hidden_size = u32_key(&format!("{prefix}.embedding_length"))?; let num_hidden_layers = u32_key(&format!("{prefix}.block_count"))?; let num_attention_heads = u32_key(&format!("{prefix}.attention.head_count"))?; - let num_kv_heads = u32_key(&format!("{prefix}.attention.head_count_kv")) - .unwrap_or(num_attention_heads); + let num_kv_heads = + u32_key(&format!("{prefix}.attention.head_count_kv")).unwrap_or(num_attention_heads); let intermediate_size = u32_key(&format!("{prefix}.feed_forward_length"))?; let head_dim = u32_key(&format!("{prefix}.attention.key_length")) .unwrap_or(hidden_size / num_attention_heads); @@ -112,8 +112,8 @@ impl GgufMapper for NomicBertMapper { let rope_theta = f32_key(&format!("{prefix}.rope.freq_base")).unwrap_or(10_000.0); // nomic-bert's epsilon key uses `_epsilon` (LayerNorm) rather // than `_rms_epsilon` (RMSNorm) — keep the same struct field. - let rms_norm_eps = f32_key(&format!("{prefix}.attention.layer_norm_epsilon")) - .unwrap_or(1e-12); + let rms_norm_eps = + f32_key(&format!("{prefix}.attention.layer_norm_epsilon")).unwrap_or(1e-12); let max_position_embeddings = u32_key(&format!("{prefix}.context_length")).unwrap_or(0); let bos_token_id = u32_key("tokenizer.ggml.bos_token_id").unwrap_or(0); let eos_token_id = u32_key("tokenizer.ggml.eos_token_id").unwrap_or(0); diff --git a/base-convert/crates/base-arch/src/gemma.rs b/base-convert/crates/base-arch/src/gemma.rs index b33f2e6..fae169e 100644 --- a/base-convert/crates/base-arch/src/gemma.rs +++ b/base-convert/crates/base-arch/src/gemma.rs @@ -129,14 +129,14 @@ impl crate::HfMapper for Gemma4HfMapper { // `head_dim` (HF) = SWA head dim (256 on E2B/E4B); // `global_head_dim` (HF) = Global head dim (512 on E2B/E4B). config.head_dim_swa = u32_key("head_dim").unwrap_or(config.head_dim); - config.head_dim_global = - u32_key("global_head_dim").unwrap_or(config.head_dim_swa); + config.head_dim_global = u32_key("global_head_dim").unwrap_or(config.head_dim_swa); let num_kv_shared = u32_key("num_kv_shared_layers").unwrap_or(0); - config.n_layer_kv_from_start = if num_kv_shared > 0 && num_kv_shared < config.num_hidden_layers { - config.num_hidden_layers - num_kv_shared - } else { - config.num_hidden_layers - }; + config.n_layer_kv_from_start = + if num_kv_shared > 0 && num_kv_shared < config.num_hidden_layers { + config.num_hidden_layers - num_kv_shared + } else { + config.num_hidden_layers + }; config.logit_softcap = f32_key("final_logit_softcapping").unwrap_or(0.0); config.sliding_window = u32_key("sliding_window").unwrap_or(0); @@ -150,9 +150,7 @@ impl crate::HfMapper for Gemma4HfMapper { if let Some(theta) = full.get("rope_theta").and_then(|v| v.as_f64()) { config.rope_theta = theta as f32; } - if let Some(prf) = - full.get("partial_rotary_factor").and_then(|v| v.as_f64()) - { + if let Some(prf) = full.get("partial_rotary_factor").and_then(|v| v.as_f64()) { config.global_rope_partial_factor = prf as f32; } } @@ -172,7 +170,11 @@ impl crate::HfMapper for Gemma4HfMapper { if let Some(arr) = source.get("layer_types").and_then(|v| v.as_array()) { config.swa_layers = arr .iter() - .map(|v| v.as_str().map(|s| s == "sliding_attention").unwrap_or(false)) + .map(|v| { + v.as_str() + .map(|s| s == "sliding_attention") + .unwrap_or(false) + }) .collect(); config.per_layer_attn = arr .iter() @@ -188,10 +190,7 @@ impl crate::HfMapper for Gemma4HfMapper { .get("sliding_window_pattern") .and_then(|v| v.as_array()) { - config.swa_layers = arr - .iter() - .map(|v| v.as_bool().unwrap_or(false)) - .collect(); + config.swa_layers = arr.iter().map(|v| v.as_bool().unwrap_or(false)).collect(); config.per_layer_attn = config .swa_layers .iter() @@ -261,13 +260,14 @@ fn extract_config(m: &BTreeMap, prefix: &str) -> Result ( - u32_key(&kv_key).unwrap_or(num_attention_heads), - Vec::new(), - ), + Some(_) => (u32_key(&kv_key).unwrap_or(num_attention_heads), Vec::new()), None => (num_attention_heads, Vec::new()), }; @@ -499,11 +499,9 @@ impl GgufMapper for Gemma4Mapper { // 1/sqrt(head_dim) like other archs). config.attention_scale = 1.0; - config.n_embd_per_layer = - u32_key("gemma4.embedding_length_per_layer_input").unwrap_or(0); + config.n_embd_per_layer = u32_key("gemma4.embedding_length_per_layer_input").unwrap_or(0); - config.head_dim_global = - u32_key("gemma4.attention.key_length").unwrap_or(config.head_dim); + config.head_dim_global = u32_key("gemma4.attention.key_length").unwrap_or(config.head_dim); config.head_dim_swa = u32_key("gemma4.attention.key_length_swa").unwrap_or(config.head_dim_global); @@ -511,17 +509,15 @@ impl GgufMapper for Gemma4Mapper { // layer's KV cache. n_layer_kv_from_start is therefore // (n_layers - shared_kv_layers). let n_shared_kv = u32_key("gemma4.attention.shared_kv_layers").unwrap_or(0); - config.n_layer_kv_from_start = - if n_shared_kv > 0 && n_shared_kv < config.num_hidden_layers { - config.num_hidden_layers - n_shared_kv - } else { - config.num_hidden_layers - }; + config.n_layer_kv_from_start = if n_shared_kv > 0 && n_shared_kv < config.num_hidden_layers + { + config.num_hidden_layers - n_shared_kv + } else { + config.num_hidden_layers + }; - config.logit_softcap = - f32_key("gemma4.final_logit_softcapping").unwrap_or(0.0); - config.sliding_window = - u32_key("gemma4.attention.sliding_window").unwrap_or(0); + config.logit_softcap = f32_key("gemma4.final_logit_softcapping").unwrap_or(0.0); + config.sliding_window = u32_key("gemma4.attention.sliding_window").unwrap_or(0); config.rope_local_theta = f32_key("gemma4.rope.local.freq_base") .or_else(|| f32_key("gemma4.rope.freq_base_swa")) .unwrap_or(0.0); @@ -535,8 +531,7 @@ impl GgufMapper for Gemma4Mapper { // absent (MLX-source bundles skip that tensor). if let Some(rope_dim) = u32_key("gemma4.rope.dimension_count") { if config.head_dim_global > 0 { - config.global_rope_partial_factor = - rope_dim as f32 / config.head_dim_global as f32; + config.global_rope_partial_factor = rope_dim as f32 / config.head_dim_global as f32; } } @@ -589,11 +584,21 @@ pub fn map_gemma4_mmproj_name(n: &str) -> Option { // Patch embedder + factorized positional embedding. match rest { "patch_embedder.input_proj.weight" => return Some("vision.patch_embed.weight".into()), - "patch_embedder.input_proj.input_max" => return Some("vision.patch_embed.input_max".into()), - "patch_embedder.input_proj.input_min" => return Some("vision.patch_embed.input_min".into()), - "patch_embedder.input_proj.output_max" => return Some("vision.patch_embed.output_max".into()), - "patch_embedder.input_proj.output_min" => return Some("vision.patch_embed.output_min".into()), - "patch_embedder.position_embedding_table" => return Some("vision.pos_embed.weight".into()), + "patch_embedder.input_proj.input_max" => { + return Some("vision.patch_embed.input_max".into()) + } + "patch_embedder.input_proj.input_min" => { + return Some("vision.patch_embed.input_min".into()) + } + "patch_embedder.input_proj.output_max" => { + return Some("vision.patch_embed.output_max".into()) + } + "patch_embedder.input_proj.output_min" => { + return Some("vision.patch_embed.output_min".into()) + } + "patch_embedder.position_embedding_table" => { + return Some("vision.pos_embed.weight".into()) + } _ => {} } // encoder.layers.{n}. @@ -616,11 +621,21 @@ pub fn map_gemma4_mmproj_name(n: &str) -> Option { if let Some(rest) = n.strip_prefix("audio_tower.") { // SubSampleConvProjection (front-end). match rest { - "subsample_conv_projection.layer0.conv.weight" => return Some("audio.sscp.layer0.conv.weight".into()), - "subsample_conv_projection.layer0.norm.weight" => return Some("audio.sscp.layer0.norm.weight".into()), - "subsample_conv_projection.layer1.conv.weight" => return Some("audio.sscp.layer1.conv.weight".into()), - "subsample_conv_projection.layer1.norm.weight" => return Some("audio.sscp.layer1.norm.weight".into()), - "subsample_conv_projection.input_proj_linear.weight" => return Some("audio.sscp.proj.weight".into()), + "subsample_conv_projection.layer0.conv.weight" => { + return Some("audio.sscp.layer0.conv.weight".into()) + } + "subsample_conv_projection.layer0.norm.weight" => { + return Some("audio.sscp.layer0.norm.weight".into()) + } + "subsample_conv_projection.layer1.conv.weight" => { + return Some("audio.sscp.layer1.conv.weight".into()) + } + "subsample_conv_projection.layer1.norm.weight" => { + return Some("audio.sscp.layer1.norm.weight".into()) + } + "subsample_conv_projection.input_proj_linear.weight" => { + return Some("audio.sscp.proj.weight".into()) + } "output_proj.weight" => return Some("audio.output_proj.weight".into()), "output_proj.bias" => return Some("audio.output_proj.bias".into()), _ => {} @@ -657,12 +672,16 @@ fn vision_layer_suffix(s: &str) -> Option { _ => {} } // Self-attention: q/k/v/o proj and qk-norm. - if let Some(out) = strip_clipped_proj(s, "self_attn", &[ - ("q_proj", "attention.q"), - ("k_proj", "attention.k"), - ("v_proj", "attention.v"), - ("o_proj", "attention.output"), - ]) { + if let Some(out) = strip_clipped_proj( + s, + "self_attn", + &[ + ("q_proj", "attention.q"), + ("k_proj", "attention.k"), + ("v_proj", "attention.v"), + ("o_proj", "attention.output"), + ], + ) { return Some(out); } if s == "self_attn.q_norm.weight" { @@ -672,11 +691,15 @@ fn vision_layer_suffix(s: &str) -> Option { return Some("attention.k_norm.weight".into()); } // MLP (GeGLU): gate / up / down with clipped-linear bounds. - if let Some(out) = strip_clipped_proj(s, "mlp", &[ - ("gate_proj", "ffn.gate"), - ("up_proj", "ffn.up"), - ("down_proj", "ffn.down"), - ]) { + if let Some(out) = strip_clipped_proj( + s, + "mlp", + &[ + ("gate_proj", "ffn.gate"), + ("up_proj", "ffn.up"), + ("down_proj", "ffn.down"), + ], + ) { return Some(out); } None @@ -846,14 +869,17 @@ mod tests { g.insert( "gemma4.attention.head_count_kv".into(), KvValue::Array( - std::iter::repeat_n([ - KvValue::U32(8), - KvValue::U32(8), - KvValue::U32(8), - KvValue::U32(8), - KvValue::U32(8), - KvValue::U32(2), - ], 5) + std::iter::repeat_n( + [ + KvValue::U32(8), + KvValue::U32(8), + KvValue::U32(8), + KvValue::U32(8), + KvValue::U32(8), + KvValue::U32(2), + ], + 5, + ) .flatten() .collect(), ), @@ -866,7 +892,10 @@ mod tests { g.insert("gemma4.attention.key_length".into(), KvValue::U32(512)); g.insert("gemma4.attention.value_length".into(), KvValue::U32(512)); g.insert("gemma4.attention.key_length_swa".into(), KvValue::U32(256)); - g.insert("gemma4.attention.value_length_swa".into(), KvValue::U32(256)); + g.insert( + "gemma4.attention.value_length_swa".into(), + KvValue::U32(256), + ); g.insert("gemma4.vocab_size".into(), KvValue::U32(262144)); g.insert("gemma4.rope.freq_base".into(), KvValue::F32(1_000_000.0)); g.insert("gemma4.rope.freq_base_swa".into(), KvValue::F32(10_000.0)); @@ -878,23 +907,23 @@ mod tests { ); g.insert("gemma4.expert_count".into(), KvValue::U32(128)); g.insert("gemma4.expert_used_count".into(), KvValue::U32(8)); - g.insert( - "gemma4.final_logit_softcapping".into(), - KvValue::F32(30.0), - ); + g.insert("gemma4.final_logit_softcapping".into(), KvValue::F32(30.0)); g.insert("gemma4.attention.sliding_window".into(), KvValue::U32(1024)); g.insert("gemma4.attention.shared_kv_layers".into(), KvValue::U32(0)); g.insert( "gemma4.attention.sliding_window_pattern".into(), KvValue::Array( - std::iter::repeat_n([ - KvValue::Bool(true), - KvValue::Bool(true), - KvValue::Bool(true), - KvValue::Bool(true), - KvValue::Bool(true), - KvValue::Bool(false), - ], 5) + std::iter::repeat_n( + [ + KvValue::Bool(true), + KvValue::Bool(true), + KvValue::Bool(true), + KvValue::Bool(true), + KvValue::Bool(true), + KvValue::Bool(false), + ], + 5, + ) .flatten() .collect(), ), @@ -968,7 +997,10 @@ mod tests { g.insert("gemma3.attention.key_length".into(), KvValue::U32(256)); g.insert("gemma3.attention.value_length".into(), KvValue::U32(256)); g.insert("gemma3.feed_forward_length".into(), KvValue::U32(6912)); - g.insert("gemma3.attention.layer_norm_rms_epsilon".into(), KvValue::F32(1e-6)); + g.insert( + "gemma3.attention.layer_norm_rms_epsilon".into(), + KvValue::F32(1e-6), + ); g.insert("gemma3.rope.freq_base".into(), KvValue::F32(1_000_000.0)); g.insert("gemma3.vocab_size".into(), KvValue::U32(262144)); g.insert("gemma3.attention.sliding_window".into(), KvValue::U32(512)); @@ -976,12 +1008,18 @@ mod tests { let c = Gemma3Mapper.config_from_gguf(&g).unwrap(); assert_eq!(c.sliding_window, 512, "read from the GGUF, not assumed"); assert_eq!(c.sliding_window_pattern, 6, "5 SWA : 1 global"); - assert_eq!(c.rope_local_theta, 10_000.0, "SWA layers use their own theta"); + assert_eq!( + c.rope_local_theta, 10_000.0, + "SWA layers use their own theta" + ); assert_eq!(c.rope_theta, 1_000_000.0, "global layers keep freq_base"); // 4B/12B/27B ship a 1024 window — the value must track the file. g.insert("gemma3.attention.sliding_window".into(), KvValue::U32(1024)); - assert_eq!(Gemma3Mapper.config_from_gguf(&g).unwrap().sliding_window, 1024); + assert_eq!( + Gemma3Mapper.config_from_gguf(&g).unwrap().sliding_window, + 1024 + ); } /// The gemma3 defaults must NOT leak onto the gemma2/gemma fallback @@ -995,7 +1033,10 @@ mod tests { g.insert("gemma2.attention.head_count".into(), KvValue::U32(8)); g.insert("gemma2.attention.head_count_kv".into(), KvValue::U32(4)); g.insert("gemma2.feed_forward_length".into(), KvValue::U32(9216)); - g.insert("gemma2.attention.layer_norm_rms_epsilon".into(), KvValue::F32(1e-6)); + g.insert( + "gemma2.attention.layer_norm_rms_epsilon".into(), + KvValue::F32(1e-6), + ); g.insert("gemma2.vocab_size".into(), KvValue::U32(256000)); g.insert("gemma2.attention.sliding_window".into(), KvValue::U32(4096)); @@ -1025,7 +1066,7 @@ mod tests { "lm_head.weight", "layers.0.self_attn.q_proj.weight", "layers.0.mlp.gate_proj.weight", - "layers.0.input_layernorm.bias", // hypothetical; norms have no bias in gemma3 + "layers.0.input_layernorm.bias", // hypothetical; norms have no bias in gemma3 ] { assert_eq!(m.norm_shift(s), 0.0, "should not shift: {s}"); } @@ -1078,7 +1119,12 @@ mod tests { }); let c = Gemma3HfMapper.config_from_hf(&cfg).unwrap(); let want = 1.0_f32 / (168.0_f32).sqrt(); - assert!((c.attention_scale - want).abs() < 1e-6, "want {}, got {}", want, c.attention_scale); + assert!( + (c.attention_scale - want).abs() < 1e-6, + "want {}, got {}", + want, + c.attention_scale + ); } /// Defensive: HF sometimes stores numeric scalars as floats. Make @@ -1095,6 +1141,11 @@ mod tests { }); let c = Gemma3HfMapper.config_from_hf(&cfg).unwrap(); let want = 1.0_f32 / (168.0_f32).sqrt(); - assert!((c.attention_scale - want).abs() < 1e-6, "want {}, got {}", want, c.attention_scale); + assert!( + (c.attention_scale - want).abs() < 1e-6, + "want {}, got {}", + want, + c.attention_scale + ); } } diff --git a/base-convert/crates/base-arch/src/glm.rs b/base-convert/crates/base-arch/src/glm.rs new file mode 100644 index 0000000..4cce2d8 --- /dev/null +++ b/base-convert/crates/base-arch/src/glm.rs @@ -0,0 +1,596 @@ +//! GLM 5.2 (`glm-dsa`) GGUF → canonical `.base` mapping. +//! +//! GLM 5.2 is a DeepSeek-V3.2-style decoder: +//! * Multi-head Latent Attention (MLA): compressed query (`attn_q_a` → +//! `attn_q_a_norm` → `attn_q_b`) and compressed KV +//! (`attn_kv_a_mqa` → `attn_kv_a_norm`, up-projected by `attn_k_b` / +//! `attn_v_b`), with a decoupled `qk_rope_head_dim`-wide RoPE. +//! * DeepSeek Sparse Attention (DSA) "lightning indexer" +//! (`indexer.*`) selecting the top-`indexer_top_k` keys per query. +//! * Sigmoid-gated MoE (256 experts, top-8) with a bias-corrected +//! selection (`exp_probs_b`), routed-weight normalization + scaling, +//! and one always-on shared expert. The first +//! `leading_dense_block_count` layers use a dense SwiGLU FFN. +//! * A Multi-Token-Prediction (`nextn`) head — dropped at convert time. +//! +//! Tensor names for the shared pieces (norms, o_proj, router, expert +//! stacks, shared expert, dense FFN, embeddings, output) follow the +//! Llama/Qwen convention, so we delegate those to `map_llama_style` and +//! only special-case the MLA + indexer tensors here. + +use crate::llama::map_llama_style; +use crate::{ArchConfig, GgufMapper, HfMapper}; +use anyhow::{bail, Context, Result}; +use base_readers::gguf::KvValue; +use std::collections::BTreeMap; + +pub struct GlmDsaMapper; + +impl GgufMapper for GlmDsaMapper { + fn canonical_arch(&self) -> &'static str { + // Must contain "glm" — the runtime's arch_from_config() matches + // on `strstr(arch, "glm")` to select the MLA + DSA encoder. + "glm_dsa" + } + + fn config_from_gguf(&self, m: &BTreeMap) -> Result { + let prefix = "glm-dsa"; + + let u32_key = |k: &str| { + m.get(k) + .and_then(|v| v.as_u64()) + .map(|n| n as u32) + .with_context(|| format!("missing metadata key: {k}")) + }; + let u32_opt = |k: &str| m.get(k).and_then(|v| v.as_u64()).map(|n| n as u32); + let f32_key = |k: &str| m.get(k).and_then(|v| v.as_f32()); + let bool_opt = |k: &str| m.get(k).and_then(|v| v.as_bool()); + + let hidden_size = u32_key(&format!("{prefix}.embedding_length"))?; + // GGUF `block_count` includes the trailing Multi-Token-Prediction + // (MTP / nextn) block, which llama.cpp loads but never runs. The + // real transformer stack is `block_count - nextn_predict_layers`. + let block_count = u32_key(&format!("{prefix}.block_count"))?; + let nextn_predict_layers = u32_opt(&format!("{prefix}.nextn_predict_layers")).unwrap_or(0); + let num_hidden_layers = block_count.saturating_sub(nextn_predict_layers); + let num_attention_heads = u32_key(&format!("{prefix}.attention.head_count"))?; + let num_kv_heads = + u32_opt(&format!("{prefix}.attention.head_count_kv")).unwrap_or(num_attention_heads); + // Dense-FFN width (used by the leading dense layers). + let intermediate_size = u32_key(&format!("{prefix}.feed_forward_length"))?; + let vocab_size = u32_key(&format!("{prefix}.vocab_size")).or_else(|_| { + m.get("tokenizer.ggml.tokens") + .and_then(|v| match v { + KvValue::Array(a) => Some(a.len() as u32), + _ => None, + }) + .context("no vocab_size and no tokenizer.ggml.tokens") + })?; + + let rope_theta = f32_key(&format!("{prefix}.rope.freq_base")).unwrap_or(10_000.0); + let rope_scale = f32_key(&format!("{prefix}.rope.scaling.factor")).unwrap_or(1.0); + let rms_norm_eps = + f32_key(&format!("{prefix}.attention.layer_norm_rms_epsilon")).unwrap_or(1e-6); + + // ── MLA geometry ──────────────────────────────────────────── + let q_lora_rank = u32_opt(&format!("{prefix}.attention.q_lora_rank")).unwrap_or(0); + let kv_lora_rank = u32_opt(&format!("{prefix}.attention.kv_lora_rank")).unwrap_or(0); + // Decoupled-RoPE width (the only rotated part of Q/K). + let qk_rope_head_dim = u32_opt(&format!("{prefix}.rope.dimension_count")).unwrap_or(0); + // Per-head expanded Q/K dim (`key_length_mla`) = nope + rope. + let key_length_mla = u32_opt(&format!("{prefix}.attention.key_length_mla")).unwrap_or(0); + let qk_nope_head_dim = key_length_mla.saturating_sub(qk_rope_head_dim); + // Per-head value dim after v_b up-projection. + let v_head_dim = u32_opt(&format!("{prefix}.attention.value_length_mla")).unwrap_or(0); + // Generic head_dim slot: the full per-head Q/K dim. The MLA + // encoder reads the specific fields above; this keeps the + // header's derived q_dim/kv_dim sane for any generic consumer. + let head_dim = if key_length_mla > 0 { + key_length_mla + } else { + hidden_size / num_attention_heads + }; + + // ── MoE topology ──────────────────────────────────────────── + let num_experts = u32_opt(&format!("{prefix}.expert_count")).unwrap_or(0); + let num_experts_per_tok = u32_opt(&format!("{prefix}.expert_used_count")).unwrap_or(0); + let moe_intermediate_size = + u32_opt(&format!("{prefix}.expert_feed_forward_length")).unwrap_or(0); + let num_shared_experts = u32_opt(&format!("{prefix}.expert_shared_count")).unwrap_or(0); + // GGUF `expert_gating_func`: 1 = softmax, 2 = sigmoid (DeepSeek/GLM). + // Runtime `expert_gating`: 0 = softmax, 1 = sigmoid. + let expert_gating_func = u32_opt(&format!("{prefix}.expert_gating_func")).unwrap_or(1); + let expert_gating = if expert_gating_func == 2 { 1 } else { 0 }; + let routed_scaling_factor = + f32_key(&format!("{prefix}.expert_weights_scale")).unwrap_or(0.0); + let norm_topk_prob = bool_opt(&format!("{prefix}.expert_weights_norm")).unwrap_or(true); + let first_k_dense_replace = + u32_opt(&format!("{prefix}.leading_dense_block_count")).unwrap_or(0); + + // ── DSA lightning indexer ─────────────────────────────────── + let indexer_head_count = + u32_opt(&format!("{prefix}.attention.indexer.head_count")).unwrap_or(0); + let indexer_key_length = + u32_opt(&format!("{prefix}.attention.indexer.key_length")).unwrap_or(0); + let indexer_top_k = u32_opt(&format!("{prefix}.attention.indexer.top_k")).unwrap_or(0); + + // Token ids (best-effort; runtime also harvests from the tokenizer). + let bos_token_id = u32_opt("tokenizer.ggml.bos_token_id").unwrap_or(0); + let eos_token_id = u32_opt("tokenizer.ggml.eos_token_id").unwrap_or(0); + let max_position_embeddings = u32_opt(&format!("{prefix}.context_length")).unwrap_or(0); + + Ok(ArchConfig { + hidden_size, + num_hidden_layers, + num_attention_heads, + num_kv_heads, + head_dim, + intermediate_size, + vocab_size, + rope_theta, + rope_scale, + rms_norm_eps, + tie_word_embeddings: false, + num_experts, + num_experts_per_tok, + moe_intermediate_size, + norm_topk_prob, + num_shared_experts, + max_position_embeddings, + bos_token_id, + eos_token_id, + // MLA + sparse-attention + GLM-MoE extras. + q_lora_rank, + kv_lora_rank, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + routed_scaling_factor, + expert_gating, + first_k_dense_replace, + nextn_predict_layers, + indexer_head_count, + indexer_key_length, + indexer_top_k, + ..ArchConfig::default() + }) + } + + fn map_tensor_name(&self, n: &str) -> Option { + // Drop the Multi-Token-Prediction (nextn) head — not used for + // standard next-token decoding. These tensors live on the trailing + // MTP block; dropping the `nextn.*` projections is enough to skip + // the head (the block's ordinary attn/ffn tensors are ignored by + // the runtime, which only iterates the first `num_hidden_layers`). + if n.contains(".nextn.") { + return None; + } + + // MLA + DSA-indexer tensors that map_llama_style doesn't know. + if let Some(rest) = n.strip_prefix("blk.") { + if let Some((layer_str, suffix)) = rest.split_once('.') { + if let Ok(layer) = layer_str.parse::() { + let canonical_suffix = match suffix { + // ── MLA attention ── + "attn_q_a.weight" => Some("self_attn.q_a_proj.weight"), + "attn_q_a_norm.weight" => Some("self_attn.q_a_layernorm.weight"), + "attn_q_b.weight" => Some("self_attn.q_b_proj.weight"), + "attn_kv_a_mqa.weight" => Some("self_attn.kv_a_proj_with_mqa.weight"), + "attn_kv_a_norm.weight" => Some("self_attn.kv_a_layernorm.weight"), + "attn_k_b.weight" => Some("self_attn.k_b_proj.weight"), + "attn_v_b.weight" => Some("self_attn.v_b_proj.weight"), + // ── DSA lightning indexer ── + "indexer.attn_q_b.weight" => Some("self_attn.indexer.q_b_proj.weight"), + "indexer.attn_k.weight" => Some("self_attn.indexer.k_proj.weight"), + "indexer.k_norm.weight" => Some("self_attn.indexer.k_norm.weight"), + "indexer.k_norm.bias" => Some("self_attn.indexer.k_norm.bias"), + "indexer.proj.weight" => Some("self_attn.indexer.weights_proj.weight"), + // ── MoE bias-correction (sigmoid top-k selection) ── + "exp_probs_b.bias" => Some("mlp.gate.e_score_correction_bias"), + _ => None, + }; + if let Some(s) = canonical_suffix { + return Some(format!("layers.{layer}.{s}")); + } + } + } + } + + // Everything else (norms, o_proj, router, expert stacks, shared + // expert, dense FFN, embeddings, output) is Llama-shaped. + map_llama_style(n) + } +} + +/// HF/MLX `config.json` (`model_type: glm_moe_dsa`) → ArchConfig. +/// +/// Mirrors [`GlmDsaMapper::config_from_gguf`] so a GGUF-sourced and an +/// MLX-sourced `.base` carry equivalent headers. Differences from the +/// GGUF path, on purpose: +/// * `num_hidden_layers` is already ex-MTP in HF configs (78; the +/// MLX export additionally sets `num_nextn_predict_layers: 0` +/// because it drops the MTP weights entirely). +/// * `num_kv_heads` is forced to 1 (the MQA latent). The HF config +/// says `num_key_value_heads: 64`, which describes the *expanded* +/// per-head view, but the GGUF path (and the runtime's MLA KV +/// sizing) use 1 — keep the headers consistent. +/// * The DSA full/shared layer pattern (`indexer_types` + friends) +/// only exists here; GGUF metadata doesn't carry it. +impl HfMapper for GlmDsaMapper { + fn canonical_arch(&self) -> &'static str { + // Same string as the GGUF side — the runtime dispatches on + // `strstr(arch, "glm")`. + "glm_dsa" + } + + fn config_from_hf(&self, c: &serde_json::Value) -> Result { + let u32_key = |k: &str| { + c.get(k) + .and_then(|v| v.as_u64()) + .map(|n| n as u32) + .with_context(|| format!("config.json missing key: {k}")) + }; + let u32_opt = |k: &str| c.get(k).and_then(|v| v.as_u64()).map(|n| n as u32); + let f32_opt = |k: &str| c.get(k).and_then(|v| v.as_f64()).map(|f| f as f32); + let bool_opt = |k: &str| c.get(k).and_then(|v| v.as_bool()); + + // The runtime's sigmoid+bias router kernel has no expert-group + // masking; GLM 5.2 ships n_group = 1 so the MLX reference's + // group step is a no-op. Refuse anything else loudly. + let n_group = u32_opt("n_group").unwrap_or(1); + if n_group > 1 { + bail!( + "glm_moe_dsa with n_group = {n_group} (grouped expert routing) is not \ + supported — the runtime router has no group masking" + ); + } + if let Some(func) = c.get("scoring_func").and_then(|v| v.as_str()) { + if func != "sigmoid" { + bail!("glm_moe_dsa scoring_func {func:?} not supported (expected \"sigmoid\")"); + } + } + + let qk_nope_head_dim = u32_key("qk_nope_head_dim")?; + let qk_rope_head_dim = u32_key("qk_rope_head_dim")?; + + // `rope_theta` lives under `rope_parameters` on GLM 5.2. + let rope_theta = c + .get("rope_parameters") + .and_then(|rp| rp.get("rope_theta")) + .and_then(|v| v.as_f64()) + .map(|f| f as f32) + .or_else(|| f32_opt("rope_theta")) + .unwrap_or(10_000.0); + + // `eos_token_id` is a list on GLM 5.2; first entry is primary, + // the rest register as additional stop ids. + let (eos_token_id, eos_token_ids) = match c.get("eos_token_id") { + Some(serde_json::Value::Array(a)) => { + let ids: Vec = a + .iter() + .filter_map(|v| v.as_u64().map(|n| n as u32)) + .collect(); + ( + ids.first().copied().unwrap_or(0), + ids.get(1..).unwrap_or(&[]).to_vec(), + ) + } + Some(v) => (v.as_u64().unwrap_or(0) as u32, vec![]), + None => (0, vec![]), + }; + + let indexer_layer_types: Vec = c + .get("indexer_types") + .and_then(|v| v.as_array()) + .map(|a| { + a.iter() + .filter_map(|x| x.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + + Ok(ArchConfig { + hidden_size: u32_key("hidden_size")?, + num_hidden_layers: u32_key("num_hidden_layers")?, + num_attention_heads: u32_key("num_attention_heads")?, + num_kv_heads: 1, + head_dim: qk_nope_head_dim + qk_rope_head_dim, + intermediate_size: u32_key("intermediate_size")?, + vocab_size: u32_key("vocab_size")?, + rope_theta, + rope_scale: 1.0, + rms_norm_eps: f32_opt("rms_norm_eps").unwrap_or(1e-5), + tie_word_embeddings: bool_opt("tie_word_embeddings").unwrap_or(false), + num_experts: u32_opt("n_routed_experts").unwrap_or(0), + num_experts_per_tok: u32_opt("num_experts_per_tok").unwrap_or(0), + moe_intermediate_size: u32_opt("moe_intermediate_size").unwrap_or(0), + norm_topk_prob: bool_opt("norm_topk_prob").unwrap_or(true), + num_shared_experts: u32_opt("n_shared_experts").unwrap_or(0), + max_position_embeddings: u32_opt("max_position_embeddings").unwrap_or(0), + bos_token_id: 0, + eos_token_id, + eos_token_ids, + // MLA geometry. + q_lora_rank: u32_opt("q_lora_rank").unwrap_or(0), + kv_lora_rank: u32_opt("kv_lora_rank").unwrap_or(0), + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim: u32_key("v_head_dim")?, + routed_scaling_factor: f32_opt("routed_scaling_factor").unwrap_or(0.0), + expert_gating: 1, // sigmoid (validated above) + first_k_dense_replace: u32_opt("first_k_dense_replace").unwrap_or(0), + nextn_predict_layers: u32_opt("num_nextn_predict_layers").unwrap_or(0), + // DSA lightning indexer. + indexer_head_count: u32_opt("index_n_heads").unwrap_or(0), + indexer_key_length: u32_opt("index_head_dim").unwrap_or(0), + indexer_top_k: u32_opt("index_topk").unwrap_or(0), + indexer_layer_types, + index_topk_freq: u32_opt("index_topk_freq").unwrap_or(0), + index_skip_topk_offset: u32_opt("index_skip_topk_offset").unwrap_or(0), + indexer_rope_interleave: bool_opt("indexer_rope_interleave").unwrap_or(false), + index_share_for_mtp_iteration: bool_opt("index_share_for_mtp_iteration") + .unwrap_or(false), + ..ArchConfig::default() + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn glm_metadata() -> BTreeMap { + let mut m = BTreeMap::new(); + let u = |v: u32| KvValue::U32(v); + let f = |v: f32| KvValue::F32(v); + m.insert("glm-dsa.embedding_length".into(), u(6144)); + m.insert("glm-dsa.block_count".into(), u(79)); + m.insert("glm-dsa.attention.head_count".into(), u(64)); + m.insert("glm-dsa.attention.head_count_kv".into(), u(1)); + m.insert("glm-dsa.feed_forward_length".into(), u(12288)); + m.insert("glm-dsa.vocab_size".into(), u(154880)); + m.insert("glm-dsa.attention.layer_norm_rms_epsilon".into(), f(1e-6)); + m.insert("glm-dsa.rope.freq_base".into(), f(8_000_000.0)); + m.insert("glm-dsa.rope.dimension_count".into(), u(64)); + m.insert("glm-dsa.attention.q_lora_rank".into(), u(2048)); + m.insert("glm-dsa.attention.kv_lora_rank".into(), u(512)); + m.insert("glm-dsa.attention.key_length".into(), u(576)); + m.insert("glm-dsa.attention.value_length".into(), u(512)); + m.insert("glm-dsa.attention.key_length_mla".into(), u(256)); + m.insert("glm-dsa.attention.value_length_mla".into(), u(256)); + m.insert("glm-dsa.expert_count".into(), u(256)); + m.insert("glm-dsa.expert_used_count".into(), u(8)); + m.insert("glm-dsa.expert_shared_count".into(), u(1)); + m.insert("glm-dsa.expert_feed_forward_length".into(), u(2048)); + m.insert("glm-dsa.expert_gating_func".into(), u(2)); + m.insert("glm-dsa.expert_weights_scale".into(), f(2.5)); + m.insert("glm-dsa.expert_weights_norm".into(), KvValue::Bool(true)); + m.insert("glm-dsa.leading_dense_block_count".into(), u(3)); + m.insert("glm-dsa.nextn_predict_layers".into(), u(1)); + m.insert("glm-dsa.attention.indexer.head_count".into(), u(32)); + m.insert("glm-dsa.attention.indexer.key_length".into(), u(128)); + m.insert("glm-dsa.attention.indexer.top_k".into(), u(2048)); + m + } + + #[test] + fn glm_dsa_config_from_gguf() { + let c = GlmDsaMapper.config_from_gguf(&glm_metadata()).unwrap(); + assert_eq!(c.hidden_size, 6144); + // block_count(79) - nextn_predict_layers(1) = 78 real layers. + assert_eq!(c.num_hidden_layers, 78); + assert_eq!(c.num_attention_heads, 64); + assert_eq!(c.num_kv_heads, 1); + assert_eq!(c.intermediate_size, 12288); + assert_eq!(c.vocab_size, 154880); + assert_eq!(c.rope_theta, 8_000_000.0); + // MLA geometry. + assert_eq!(c.q_lora_rank, 2048); + assert_eq!(c.kv_lora_rank, 512); + assert_eq!(c.qk_rope_head_dim, 64); + assert_eq!(c.qk_nope_head_dim, 192); // 256 - 64 + assert_eq!(c.v_head_dim, 256); + assert_eq!(c.head_dim, 256); + // MoE. + assert_eq!(c.num_experts, 256); + assert_eq!(c.num_experts_per_tok, 8); + assert_eq!(c.num_shared_experts, 1); + assert_eq!(c.moe_intermediate_size, 2048); + assert_eq!(c.expert_gating, 1); // sigmoid + assert_eq!(c.routed_scaling_factor, 2.5); + assert!(c.norm_topk_prob); + assert_eq!(c.first_k_dense_replace, 3); + assert_eq!(c.nextn_predict_layers, 1); + // Indexer. + assert_eq!(c.indexer_head_count, 32); + assert_eq!(c.indexer_key_length, 128); + assert_eq!(c.indexer_top_k, 2048); + } + + #[test] + fn glm_dsa_config_round_trips_through_header() { + let c = GlmDsaMapper.config_from_gguf(&glm_metadata()).unwrap(); + let m = c.to_config_map(); + use serde_json::json; + assert_eq!(m["q_lora_rank"], json!(2048)); + assert_eq!(m["kv_lora_rank"], json!(512)); + assert_eq!(m["qk_nope_head_dim"], json!(192)); + assert_eq!(m["qk_rope_head_dim"], json!(64)); + assert_eq!(m["v_head_dim"], json!(256)); + assert_eq!(m["routed_scaling_factor"], json!(2.5)); + assert_eq!(m["expert_gating"], json!(1)); + assert_eq!(m["first_k_dense_replace"], json!(3)); + assert_eq!(m["num_experts"], json!(256)); + assert_eq!(m["num_shared_experts"], json!(1)); + assert_eq!(m["moe_intermediate_size"], json!(2048)); + assert_eq!(m["indexer_head_count"], json!(32)); + assert_eq!(m["indexer_top_k"], json!(2048)); + } + + #[test] + fn glm_dsa_maps_mla_and_indexer_tensors() { + let map = |n: &str| GlmDsaMapper.map_tensor_name(n); + // MLA. + assert_eq!( + map("blk.5.attn_q_a.weight").as_deref(), + Some("layers.5.self_attn.q_a_proj.weight") + ); + assert_eq!( + map("blk.5.attn_q_a_norm.weight").as_deref(), + Some("layers.5.self_attn.q_a_layernorm.weight") + ); + assert_eq!( + map("blk.5.attn_kv_a_mqa.weight").as_deref(), + Some("layers.5.self_attn.kv_a_proj_with_mqa.weight") + ); + assert_eq!( + map("blk.5.attn_k_b.weight").as_deref(), + Some("layers.5.self_attn.k_b_proj.weight") + ); + assert_eq!( + map("blk.5.attn_v_b.weight").as_deref(), + Some("layers.5.self_attn.v_b_proj.weight") + ); + // Indexer. + assert_eq!( + map("blk.5.indexer.attn_q_b.weight").as_deref(), + Some("layers.5.self_attn.indexer.q_b_proj.weight") + ); + assert_eq!( + map("blk.5.indexer.k_norm.bias").as_deref(), + Some("layers.5.self_attn.indexer.k_norm.bias") + ); + assert_eq!( + map("blk.5.indexer.proj.weight").as_deref(), + Some("layers.5.self_attn.indexer.weights_proj.weight") + ); + // MoE bias. + assert_eq!( + map("blk.5.exp_probs_b.bias").as_deref(), + Some("layers.5.mlp.gate.e_score_correction_bias") + ); + } + + #[test] + fn glm_dsa_config_from_hf_mlx_checkpoint() { + // Mirrors mlx-community/GLM-5.2-4bit's config.json (trimmed). + let c = serde_json::json!({ + "model_type": "glm_moe_dsa", + "hidden_size": 6144, + "num_hidden_layers": 78, + "num_attention_heads": 64, + "num_key_value_heads": 64, + "head_dim": 192, + "qk_head_dim": 256, + "intermediate_size": 12288, + "vocab_size": 154880, + "rms_norm_eps": 1e-5, + "rope_parameters": {"rope_theta": 8000000, "rope_type": "default"}, + "q_lora_rank": 2048, + "kv_lora_rank": 512, + "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, + "v_head_dim": 256, + "n_routed_experts": 256, + "num_experts_per_tok": 8, + "moe_intermediate_size": 2048, + "n_shared_experts": 1, + "routed_scaling_factor": 2.5, + "norm_topk_prob": true, + "first_k_dense_replace": 3, + "n_group": 1, + "topk_group": 1, + "topk_method": "noaux_tc", + "scoring_func": "sigmoid", + "max_position_embeddings": 1048576, + "eos_token_id": [154820, 154827, 154829], + "tie_word_embeddings": false, + "num_nextn_predict_layers": 0, + "index_n_heads": 32, + "index_head_dim": 128, + "index_topk": 2048, + "index_topk_freq": 4, + "index_skip_topk_offset": 3, + "indexer_rope_interleave": true, + "index_share_for_mtp_iteration": true, + "indexer_types": ["full", "full", "full", "shared"], + }); + let cfg = HfMapper::config_from_hf(&GlmDsaMapper, &c).unwrap(); + assert_eq!(cfg.hidden_size, 6144); + assert_eq!(cfg.num_hidden_layers, 78); + assert_eq!(cfg.num_kv_heads, 1); // forced MQA-latent, not the HF 64 + assert_eq!(cfg.head_dim, 256); // nope 192 + rope 64 + assert_eq!(cfg.rope_theta, 8_000_000.0); + assert_eq!(cfg.rms_norm_eps, 1e-5); + assert_eq!(cfg.q_lora_rank, 2048); + assert_eq!(cfg.kv_lora_rank, 512); + assert_eq!(cfg.v_head_dim, 256); + assert_eq!(cfg.num_experts, 256); + assert_eq!(cfg.num_experts_per_tok, 8); + assert_eq!(cfg.num_shared_experts, 1); + assert_eq!(cfg.expert_gating, 1); + assert_eq!(cfg.routed_scaling_factor, 2.5); + assert_eq!(cfg.first_k_dense_replace, 3); + assert_eq!(cfg.nextn_predict_layers, 0); + assert_eq!(cfg.eos_token_id, 154820); + assert_eq!(cfg.eos_token_ids, vec![154827, 154829]); + assert_eq!(cfg.indexer_head_count, 32); + assert_eq!(cfg.indexer_key_length, 128); + assert_eq!(cfg.indexer_top_k, 2048); + assert_eq!(cfg.indexer_layer_types.len(), 4); + assert_eq!(cfg.index_topk_freq, 4); + assert_eq!(cfg.index_skip_topk_offset, 3); + assert!(cfg.indexer_rope_interleave); + assert!(cfg.index_share_for_mtp_iteration); + // The pattern fields round-trip through the header map. + let m = cfg.to_config_map(); + assert_eq!(m["indexer_layer_types"][3], serde_json::json!("shared")); + assert_eq!(m["index_topk_freq"], serde_json::json!(4)); + assert_eq!(m["indexer_rope_interleave"], serde_json::json!(true)); + } + + #[test] + fn glm_dsa_config_from_hf_rejects_grouped_routing() { + let c = serde_json::json!({ + "hidden_size": 6144, "num_hidden_layers": 78, + "num_attention_heads": 64, "intermediate_size": 12288, + "vocab_size": 154880, "qk_nope_head_dim": 192, + "qk_rope_head_dim": 64, "v_head_dim": 256, + "n_group": 8, "scoring_func": "sigmoid", + }); + assert!(HfMapper::config_from_hf(&GlmDsaMapper, &c).is_err()); + } + + #[test] + fn glm_dsa_delegates_shared_tensors_and_drops_mtp() { + let map = |n: &str| GlmDsaMapper.map_tensor_name(n); + // Shared (Llama-shaped) tensors delegate to map_llama_style. + assert_eq!( + map("blk.5.attn_output.weight").as_deref(), + Some("layers.5.self_attn.o_proj.weight") + ); + assert_eq!( + map("blk.5.ffn_gate_inp.weight").as_deref(), + Some("layers.5.mlp.router.weight") + ); + assert_eq!( + map("blk.5.ffn_down_exps.weight").as_deref(), + Some("layers.5.mlp.experts.down_proj.weight") + ); + assert_eq!( + map("blk.5.ffn_up_shexp.weight").as_deref(), + Some("layers.5.mlp.shared_expert.up_proj.weight") + ); + assert_eq!( + map("blk.1.ffn_gate.weight").as_deref(), + Some("layers.1.mlp.gate_proj.weight") + ); + assert_eq!( + map("token_embd.weight").as_deref(), + Some("embed_tokens.weight") + ); + // MTP / nextn head dropped. + assert_eq!(map("blk.79.nextn.eh_proj.weight"), None); + assert_eq!(map("blk.79.nextn.enorm.weight"), None); + } +} diff --git a/base-convert/crates/base-arch/src/gpt_oss.rs b/base-convert/crates/base-arch/src/gpt_oss.rs new file mode 100644 index 0000000..def8ae3 --- /dev/null +++ b/base-convert/crates/base-arch/src/gpt_oss.rs @@ -0,0 +1,155 @@ +//! OpenAI gpt-oss (gpt-oss-20b / gpt-oss-120b) — HF `model_type = "gpt_oss"`. +//! +//! Decoder-only MoE transformer with a few departures from the Llama / +//! Qwen family the runtime has to know about (all carried in the header +//! `config`, never inferred from the arch string): +//! +//! - alternating sliding-window (128) / full attention layers +//! (`layer_types`), mirrored into `swa_layers` + `sliding_window`; +//! - learned per-head attention **sinks** (an extra softmax logit per +//! head that absorbs probability mass; tensor +//! `self_attn.sinks`), biased q/k/v/o projections; +//! - YaRN RoPE (factor 32 over 4096 original positions, beta_fast 32, +//! beta_slow 1) — the betas are emitted so the runtime can derive +//! the per-pair divisors and the attention-temperature `mscale`; +//! - MoE: 32 routed experts, top-4, biased linear router, softmax over +//! the selected experts (== `norm_topk_prob`), experts with biases on +//! every projection and a clamped SwiGLU +//! `(up + 1) * gate * sigmoid(1.702 * gate)` with +//! `gate <= limit`, `|up| <= limit` (`swiglu_limit`, 7.0); +//! - expert weights shipped as MXFP4 (E2M1 + E8M0 block scales), which +//! the converter transplants verbatim (see `convert_gpt_oss` in +//! base-convert). +//! +//! Tensor names are already canonical HF names; the converter's gpt-oss +//! path maps them itself (the fused `gate_up_proj` stays fused and +//! row-interleaved exactly as the checkpoint stores it — the runtime's +//! gpt-oss expert kernel reads gate = row 2j, up = row 2j+1). + +use anyhow::Result; + +use crate::{ArchConfig, HfMapper}; + +pub struct GptOssHfMapper; + +impl HfMapper for GptOssHfMapper { + fn canonical_arch(&self) -> &'static str { + "gpt_oss" + } + + fn config_from_hf(&self, c: &serde_json::Value) -> Result { + let mut config = crate::llama::hf_generic_config(c)?; + let u32_v = |k: &str| c.get(k).and_then(|v| v.as_u64()).map(|n| n as u32); + let f32_v = |k: &str| c.get(k).and_then(|v| v.as_f64()).map(|f| f as f32); + + // MoE topology. `num_local_experts` is the HF key; `experts_per_token` + // duplicates `num_experts_per_tok` on gpt-oss configs. + config.num_experts = u32_v("num_local_experts").unwrap_or(0); + config.num_experts_per_tok = u32_v("num_experts_per_tok") + .or_else(|| u32_v("experts_per_token")) + .unwrap_or(4); + // Per-expert FFN width == intermediate_size (there is no dense FFN). + config.moe_intermediate_size = config.intermediate_size; + // Router: top-k over the raw logits, then softmax over the k winners — + // algebraically the renormalized full softmax. + config.norm_topk_prob = true; + config.num_shared_experts = 0; + + // Attention schedule: `layer_types` lists "sliding_attention" | + // "full_attention" per layer (gpt-oss-20b: alternating, sliding first). + config.sliding_window = u32_v("sliding_window").unwrap_or(128); + if let Some(arr) = c.get("layer_types").and_then(|v| v.as_array()) { + config.swa_layers = arr + .iter() + .map(|v| { + v.as_str() + .map(|s| s == "sliding_attention") + .unwrap_or(false) + }) + .collect(); + config.per_layer_attn = config + .swa_layers + .iter() + .map(|&b| if b { "sliding" } else { "global" }.to_string()) + .collect(); + } else { + // Default schedule from the reference implementation: even layers + // sliding, odd layers full. + config.swa_layers = (0..config.num_hidden_layers).map(|i| i % 2 == 0).collect(); + config.per_layer_attn = config + .swa_layers + .iter() + .map(|&b| if b { "sliding" } else { "global" }.to_string()) + .collect(); + } + + // YaRN betas (HF `rope_scaling.beta_fast` / `beta_slow`); the factor and + // original_max_position_embeddings are already parsed generically. + if let Some(rs) = c.get("rope_scaling") { + config.rope_yarn_beta_fast = rs + .get("beta_fast") + .and_then(|v| v.as_f64()) + .map(|f| f as f32) + .unwrap_or(32.0); + config.rope_yarn_beta_slow = rs + .get("beta_slow") + .and_then(|v| v.as_f64()) + .map(|f| f as f32) + .unwrap_or(1.0); + config.rope_yarn_truncate = + rs.get("truncate").and_then(|v| v.as_bool()).unwrap_or(true); + if config.rope_original_max_pos == 0 { + config.rope_original_max_pos = u32_v("initial_context_length").unwrap_or(4096); + } + } + + // Clamped SwiGLU constants. + config.swiglu_limit = f32_v("swiglu_limit").unwrap_or(7.0); + config.swiglu_alpha = 1.702; + config.attention_sinks = true; + config.attention_bias = c + .get("attention_bias") + .and_then(|v| v.as_bool()) + .unwrap_or(true); + Ok(config) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_gpt_oss_20b_config() { + let c: serde_json::Value = serde_json::json!({ + "model_type": "gpt_oss", + "hidden_size": 2880, "num_hidden_layers": 4, "num_attention_heads": 64, + "num_key_value_heads": 8, "head_dim": 64, "intermediate_size": 2880, + "vocab_size": 201088, "rope_theta": 150000, "rms_norm_eps": 1e-5, + "num_local_experts": 32, "num_experts_per_tok": 4, "sliding_window": 128, + "layer_types": ["sliding_attention", "full_attention", "sliding_attention", "full_attention"], + "rope_scaling": {"rope_type": "yarn", "factor": 32.0, "beta_fast": 32.0, "beta_slow": 1.0, + "original_max_position_embeddings": 4096, "truncate": true}, + "eos_token_id": 200002, "max_position_embeddings": 131072 + }); + let cfg = GptOssHfMapper.config_from_hf(&c).unwrap(); + assert_eq!(cfg.num_experts, 32); + assert_eq!(cfg.num_experts_per_tok, 4); + assert_eq!(cfg.moe_intermediate_size, 2880); + assert!(cfg.norm_topk_prob); + assert_eq!(cfg.swa_layers, vec![true, false, true, false]); + assert_eq!(cfg.sliding_window, 128); + assert_eq!(cfg.rope_scaling_type, "yarn"); + assert_eq!(cfg.rope_scale, 32.0); + assert_eq!(cfg.rope_original_max_pos, 4096); + assert_eq!(cfg.rope_yarn_beta_fast, 32.0); + assert_eq!(cfg.rope_yarn_beta_slow, 1.0); + assert!(cfg.rope_yarn_truncate); + assert_eq!(cfg.swiglu_limit, 7.0); + let m = cfg.to_config_map(); + assert_eq!(m["num_experts"], 32); + assert_eq!(m["rope_yarn_beta_fast"], 32.0); + assert_eq!(m["swiglu_limit"], 7.0); + assert_eq!(m["attention_sinks"], true); + } +} diff --git a/base-convert/crates/base-arch/src/lib.rs b/base-convert/crates/base-arch/src/lib.rs index a2a1bc0..d5fafa1 100644 --- a/base-convert/crates/base-arch/src/lib.rs +++ b/base-convert/crates/base-arch/src/lib.rs @@ -8,8 +8,11 @@ pub mod bert; pub mod gemma; +pub mod glm; +pub mod gpt_oss; pub mod llama; pub mod muse_glimmer; +pub mod nemotron; pub mod qwen; pub mod tokenizer; pub mod whisper; @@ -38,7 +41,10 @@ pub fn source_mapper_for_gguf(arch: &str) -> Option<&'static dyn GgufMapper> { "qwen2moe" | "qwen3moe" | "qwen35moe" | "qwen36moe" => Some(&qwen::QwenMoeMapper), "gemma" | "gemma2" | "gemma3" => Some(&gemma::Gemma3Mapper), "gemma4" => Some(&gemma::Gemma4Mapper), + "nemotron_h" | "nemotron_h_moe" => Some(&nemotron::NemotronHMapper), "nomic-bert" => Some(&bert::NomicBertMapper), + // GLM 5.2 — DeepSeek-V3.2-style MLA + sparse-attention MoE. + "glm-dsa" => Some(&glm::GlmDsaMapper), // Muse Glimmer. `general.architecture` in the llama.cpp-produced // GGUF is the HYPHENATED "muse-glimmer" (llama.cpp's arch registry // spells multi-word archs with hyphens: "nomic-bert", @@ -54,6 +60,28 @@ pub fn source_mapper_for_gguf(arch: &str) -> Option<&'static dyn GgufMapper> { /// HF config.json model_type → mapper. HF tensor names already follow /// canonical convention, so the mapper only needs to extract ArchConfig /// from config.json — no tensor renaming. +/// An element-wise reparameterization undone at convert time. +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +pub enum ValueTransform { + /// `x -> -exp(x)`. Mamba-2 checkpoints store the state-transition + /// matrix as `A_log` and materialize `A = -exp(A_log)` in the model + /// code; the value the scan kernel wants is `A`, which must be + /// negative for the recurrence to decay. + NegExp, +} + +impl ValueTransform { + pub fn apply(self, values: &mut [f32]) { + match self { + ValueTransform::NegExp => { + for v in values.iter_mut() { + *v = -v.exp(); + } + } + } + } +} + pub trait HfMapper: Sync { fn canonical_arch(&self) -> &'static str; fn config_from_hf(&self, config: &serde_json::Value) -> anyhow::Result; @@ -83,6 +111,30 @@ pub trait HfMapper: Sync { None } + /// Element-wise transform to apply to a tensor's values on the way + /// from HF to `.base`, or `None` to copy them through. + /// + /// Some architectures store a *reparameterized* weight that the + /// reference implementation undoes at load or at run time, while + /// the GGUF conversion bakes the undo in. Baking it at convert time + /// keeps the runtime kernel simple and keeps bundles from the two + /// sources interchangeable. + fn value_transform(&self, _canonical: &str) -> Option { + None + } + + /// Write tensor `shape` fastest-varying dim first (GGUF's `ne` + /// order, `[in, out]`) instead of HF's C order (`[out, in]`). + /// + /// The two describe the *same bytes* — an HF `[out, in]` matrix is + /// stored with `in` contiguous, exactly like a GGUF `[in, out]` one + /// — so this changes how the header reports a tensor, not the + /// buffer behind it. Set it on architectures whose GGUF-converted + /// bundles are already in circulation, so a bundle built from either + /// source describes itself identically. Default: keep HF order. + fn shape_fastest_first(&self) -> bool { + false + } /// RMS-normalize each ROW of a 2-D tensor at HF→.base conversion /// time, returning the epsilon to use, or None to leave it alone. /// @@ -151,6 +203,13 @@ pub fn hf_mapper_for_model_type(model_type: &str) -> Option<&'static dyn HfMappe // text conversion skips like any other non-text tower. Config is // gemma4-shaped (uniform head_dim, rope_parameters, layer_types). "gemma4" | "gemma4_text" | "gemma4_unified" => Some(&gemma::Gemma4HfMapper), + // Nemotron-H hybrid (Mamba-2 + attention + MoE). The HF + // `model_type` is `nemotron_h` for both the dense and MoE + // builds — the block schedule comes from + // `hybrid_override_pattern`, and the MoE keys are simply absent + // on a dense checkpoint — so one mapper covers both. Canonical + // arch stays `nemotron_h_moe` to match the GGUF path. + "nemotron_h" | "nemotron_h_moe" => Some(&nemotron::NemotronHHfMapper), // Whisper encoder-decoder speech models (openai/whisper-*). HF // safetensors only — whisper GGML files are not GGUF, so the GGUF // dispatch table stays untouched. Tensor renaming is @@ -158,6 +217,11 @@ pub fn hf_mapper_for_model_type(model_type: &str) -> Option<&'static dyn HfMappe // path emits everything as f16 (the engine's fused whisper // kernels are f16-only in v1). "whisper" => Some(&whisper::WhisperHfMapper), + // GLM 5.2 — DeepSeek-V3.2-style MLA + DSA MoE. HF/MLX + // checkpoints declare `model_type: glm_moe_dsa`. + "glm_moe_dsa" => Some(&glm::GlmDsaMapper), + // OpenAI gpt-oss (MXFP4 MoE, attention sinks, YaRN, alternating SWA). + "gpt_oss" => Some(&gpt_oss::GptOssHfMapper), _ => None, } } @@ -190,6 +254,8 @@ pub const SUPPORTED_HF_MODEL_TYPES: &[&str] = &[ "muse_glimmer", "muse_glimmer_text", "whisper", + "glm_moe_dsa", + "gpt_oss", ]; pub trait GgufMapper: Sync { @@ -393,6 +459,43 @@ pub struct ArchConfig { /// concatenating them (Qwen3.5 = true). pub mrope_interleaved: bool, + // ── gpt-oss fields (zero/false for other archs) ────────────────── + /// YaRN correction-range betas (HF `rope_scaling.beta_fast` / + /// `beta_slow`). Only meaningful when `rope_scaling_type == "yarn"`. + pub rope_yarn_beta_fast: f32, + pub rope_yarn_beta_slow: f32, + /// YaRN `truncate` (HF `rope_scaling.truncate`, default true): floor/ceil + /// the correction range. gpt-oss ships `false` (continuous ramp bounds). + pub rope_yarn_truncate: bool, + /// Clamped-SwiGLU limit (`gate <= limit`, `|up| <= limit`) and the + /// swish alpha (`gate * sigmoid(alpha * gate)`). 0 = plain SwiGLU. + pub swiglu_limit: f32, + pub swiglu_alpha: f32, + /// Learned per-head attention sinks (`self_attn.sinks`). + pub attention_sinks: bool, + /// q/k/v/o projections carry biases. + pub attention_bias: bool, + + // ── Nemotron-H / Mamba-2 SSM fields (zero for non-SSM archs) ───── + // Nemotron-H interleaves Mamba-2 SSM blocks, GQA attention blocks + // and (MoE) FFN blocks — the schedule rides in `layer_types` + // ("mamba" | "attention" | "moe" | "mlp"). + /// Mamba-2 SSM state size per head (`d_state`). 0 = no SSM. + pub ssm_state_size: u32, + /// Causal depthwise conv kernel width in the SSM mixer (`d_conv`). + pub ssm_conv_kernel: u32, + /// Number of B/C groups (`n_groups`). + pub ssm_num_groups: u32, + /// SSM inner width (`d_inner` = heads × head dim). + pub ssm_inner_size: u32, + /// Number of SSM heads (GGUF stores this in `ssm.time_step_rank` + /// for Mamba-2 checkpoints). + pub ssm_num_heads: u32, + + // ── MoE routing extensions (DeepSeek-style routers) ────────────── + /// Routed-expert scaling factor applied after top-k + /// renormalization (Nemotron 3 Nano: 2.5). 0 = none. + pub expert_weights_scale: f32, // ── Whisper encoder-decoder fields (zero/empty for other archs) ── // The decoder half reuses the standard fields above (hidden_size / // num_hidden_layers / num_attention_heads / intermediate_size / @@ -446,6 +549,66 @@ pub struct ArchConfig { pub max_source_positions: u32, /// Decoder positional-embedding length (`max_target_positions`, 448). pub max_target_positions: u32, + + // ── GLM-DSA / DeepSeek-V3.2-style MLA + sparse-attention fields ─── + // (all zero for other archs). GLM 5.2 uses Multi-head Latent + // Attention (compressed q/kv latents + decoupled RoPE) plus a + // DeepSeek Sparse Attention "lightning indexer", with a sigmoid-gated + // MoE (bias-corrected top-k, routed weight scaling, shared expert) + // and the first `first_k_dense_replace` layers dense. + /// MLA query compression rank (`attn_q_a` output width). 0 = not MLA. + pub q_lora_rank: u32, + /// MLA key/value compression rank (`attn_kv_a_mqa` kv part). 0 = not MLA. + pub kv_lora_rank: u32, + /// Per-head non-positional Q/K dim (the part that attends in latent + /// space via k_b). GLM 5.2 = 192. + pub qk_nope_head_dim: u32, + /// Per-head decoupled-RoPE Q/K dim (the only rotated part). GLM 5.2 = 64. + pub qk_rope_head_dim: u32, + /// Per-head value dim after v_b up-projection. GLM 5.2 = 256. + pub v_head_dim: u32, + /// Routed-expert output scaling (DeepSeek `routed_scaling_factor`). + /// GLM 5.2 = 2.5. 0 = no scaling. + pub routed_scaling_factor: f32, + /// Expert gating function: 0 = softmax (default), 1 = sigmoid + /// (GLM/DeepSeek; GGUF `expert_gating_func = 2`). + pub expert_gating: u32, + /// First N layers use a dense SwiGLU FFN instead of MoE + /// (`leading_dense_block_count` / HF `first_k_dense_replace`). + /// GLM 5.2 = 3. 0 = all MoE layers. + pub first_k_dense_replace: u32, + /// Multi-Token-Prediction (nextn) head layer count. Tensors are + /// dropped at convert time; kept for the header record. GLM 5.2 = 1. + pub nextn_predict_layers: u32, + /// DSA lightning-indexer head count. GLM 5.2 = 32. 0 = no indexer. + pub indexer_head_count: u32, + /// DSA indexer per-head key dim. GLM 5.2 = 128. + pub indexer_key_length: u32, + /// DSA indexer top-k keys selected per query. GLM 5.2 = 2048. + pub indexer_top_k: u32, + /// Per-layer indexer kind: "full" (layer computes its own top-k + /// selection and carries indexer weights) or "shared" (layer reuses + /// the most recent full layer's selection; NO indexer weights). + /// GLM 5.2: 21 full / 57 shared. Empty = every layer is full (the + /// GGUF metadata doesn't carry the pattern). HF `indexer_types`. + pub indexer_layer_types: Vec, + /// Selection-reuse period for shared indexer layers (HF + /// `index_topk_freq`). GLM 5.2 = 4. 0 = unset. + pub index_topk_freq: u32, + /// Offset into the reuse period (HF `index_skip_topk_offset`). + /// GLM 5.2 = 3. Only meaningful when `index_topk_freq > 0`. + pub index_skip_topk_offset: u32, + /// Indexer RoPE layout: true = interleaved/traditional (GPT-J + /// adjacent-pair — what mlx-lm applies via `traditional=True`), + /// false = NeoX half-split. HF `indexer_rope_interleave`. + /// GLM 5.2 = true. NOTE: llama.cpp's deepseek32 reference uses NeoX + /// for its indexer; GLM's config + the MLX implementation say + /// interleaved. Trust this flag, not the deepseek32 source. + pub indexer_rope_interleave: bool, + /// Whether MTP iterations reuse the same top-k selection (HF + /// `index_share_for_mtp_iteration`). Recorded for the header; only + /// relevant once MTP lands. + pub index_share_for_mtp_iteration: bool, } impl ArchConfig { @@ -506,14 +669,17 @@ impl ArchConfig { } if self.num_experts > 0 { m.insert("num_experts".into(), json!(self.num_experts)); - m.insert("num_experts_per_tok".into(), json!(self.num_experts_per_tok)); - m.insert("moe_intermediate_size".into(), json!(self.moe_intermediate_size)); + m.insert( + "num_experts_per_tok".into(), + json!(self.num_experts_per_tok), + ); + m.insert( + "moe_intermediate_size".into(), + json!(self.moe_intermediate_size), + ); m.insert("norm_topk_prob".into(), json!(self.norm_topk_prob)); if self.num_shared_experts > 0 { - m.insert( - "num_shared_experts".into(), - json!(self.num_shared_experts), - ); + m.insert("num_shared_experts".into(), json!(self.num_shared_experts)); } } if self.max_position_embeddings > 0 { @@ -595,7 +761,10 @@ impl ArchConfig { "linear_num_value_heads".into(), json!(self.linear_num_value_heads), ); - m.insert("linear_key_head_dim".into(), json!(self.linear_key_head_dim)); + m.insert( + "linear_key_head_dim".into(), + json!(self.linear_key_head_dim), + ); m.insert( "linear_value_head_dim".into(), json!(self.linear_value_head_dim), @@ -635,6 +804,45 @@ impl ArchConfig { m.insert("mrope_section".into(), json!(self.mrope_section)); m.insert("mrope_interleaved".into(), json!(self.mrope_interleaved)); } + // gpt-oss fields — emitted only when set. + if self.rope_yarn_beta_fast > 0.0 { + m.insert( + "rope_yarn_beta_fast".into(), + json!(self.rope_yarn_beta_fast), + ); + m.insert( + "rope_yarn_beta_slow".into(), + json!(self.rope_yarn_beta_slow), + ); + m.insert("rope_yarn_truncate".into(), json!(self.rope_yarn_truncate)); + } + if self.swiglu_limit > 0.0 { + m.insert("swiglu_limit".into(), json!(self.swiglu_limit)); + m.insert("swiglu_alpha".into(), json!(self.swiglu_alpha)); + } + if self.attention_sinks { + m.insert("attention_sinks".into(), json!(true)); + } + if self.attention_bias { + m.insert("attention_bias".into(), json!(true)); + } + // Nemotron-H / Mamba-2 SSM fields — only emit when set. + if self.ssm_inner_size > 0 { + m.insert("ssm_state_size".into(), json!(self.ssm_state_size)); + m.insert("ssm_conv_kernel".into(), json!(self.ssm_conv_kernel)); + m.insert("ssm_num_groups".into(), json!(self.ssm_num_groups)); + m.insert("ssm_inner_size".into(), json!(self.ssm_inner_size)); + m.insert("ssm_num_heads".into(), json!(self.ssm_num_heads)); + } + if self.expert_gating > 0 { + m.insert("expert_gating".into(), json!(self.expert_gating)); + } + if self.expert_weights_scale > 0.0 { + m.insert( + "expert_weights_scale".into(), + json!(self.expert_weights_scale), + ); + } // Whisper encoder-decoder fields — only emit when the encoder // half is populated so decoder-only archs' headers stay tidy. if self.encoder_layers > 0 { @@ -663,6 +871,73 @@ impl ArchConfig { // struct-level uniformity). m.insert("norm_eps".into(), json!(self.rms_norm_eps)); } + // GLM-DSA / MLA + sparse-attention fields — only emit when set so + // other archs' headers stay tidy. The runtime reads these back in + // BaseWeightStore::extract_config. + if self.q_lora_rank > 0 { + m.insert("q_lora_rank".into(), json!(self.q_lora_rank)); + } + if self.kv_lora_rank > 0 { + m.insert("kv_lora_rank".into(), json!(self.kv_lora_rank)); + } + if self.qk_nope_head_dim > 0 { + m.insert("qk_nope_head_dim".into(), json!(self.qk_nope_head_dim)); + } + if self.qk_rope_head_dim > 0 { + m.insert("qk_rope_head_dim".into(), json!(self.qk_rope_head_dim)); + } + if self.v_head_dim > 0 { + m.insert("v_head_dim".into(), json!(self.v_head_dim)); + } + if self.routed_scaling_factor > 0.0 { + m.insert( + "routed_scaling_factor".into(), + json!(self.routed_scaling_factor), + ); + } + // expert_gating is meaningful even when 0 (softmax) for MoE models, + // but we only emit non-default sigmoid to keep other headers tidy. + if self.expert_gating > 0 { + m.insert("expert_gating".into(), json!(self.expert_gating)); + } + if self.first_k_dense_replace > 0 { + m.insert( + "first_k_dense_replace".into(), + json!(self.first_k_dense_replace), + ); + } + if self.nextn_predict_layers > 0 { + m.insert( + "nextn_predict_layers".into(), + json!(self.nextn_predict_layers), + ); + } + if self.indexer_head_count > 0 { + m.insert("indexer_head_count".into(), json!(self.indexer_head_count)); + m.insert("indexer_key_length".into(), json!(self.indexer_key_length)); + m.insert("indexer_top_k".into(), json!(self.indexer_top_k)); + } + // DSA full/shared layer pattern — present only on sources that + // declare it (the MLX/HF config.json; GGUF metadata doesn't). + if !self.indexer_layer_types.is_empty() { + m.insert( + "indexer_layer_types".into(), + json!(self.indexer_layer_types), + ); + m.insert("index_topk_freq".into(), json!(self.index_topk_freq)); + m.insert( + "index_skip_topk_offset".into(), + json!(self.index_skip_topk_offset), + ); + m.insert( + "indexer_rope_interleave".into(), + json!(self.indexer_rope_interleave), + ); + m.insert( + "index_share_for_mtp_iteration".into(), + json!(self.index_share_for_mtp_iteration), + ); + } m } } diff --git a/base-convert/crates/base-arch/src/llama.rs b/base-convert/crates/base-arch/src/llama.rs index c9f2db3..eb7c229 100644 --- a/base-convert/crates/base-arch/src/llama.rs +++ b/base-convert/crates/base-arch/src/llama.rs @@ -57,8 +57,7 @@ pub(crate) fn hf_generic_config(c: &serde_json::Value) -> Result Result anyhow::bail!( - "config.json `intermediate_size` must be u32 or array, got {other:?}" - ), + Some(other) => { + anyhow::bail!("config.json `intermediate_size` must be u32 or array, got {other:?}") + } None => anyhow::bail!("config.json missing intermediate_size"), }; let vocab_size = u32_key("vocab_size")?; @@ -123,14 +122,22 @@ pub(crate) fn hf_generic_config(c: &serde_json::Value) -> Result = if eos_ids.len() > 1 { eos_ids[1..].to_vec() } else { Vec::new() }; + let eos_token_ids: Vec = if eos_ids.len() > 1 { + eos_ids[1..].to_vec() + } else { + Vec::new() + }; Ok(crate::ArchConfig { hidden_size, @@ -175,18 +182,17 @@ impl GgufMapper for LlamaMapper { let num_attention_heads = u32_key("llama.attention.head_count")?; let num_kv_heads = u32_key("llama.attention.head_count_kv").unwrap_or(num_attention_heads); let intermediate_size = u32_key("llama.feed_forward_length")?; - let vocab_size = u32_key("llama.vocab_size") - .or_else(|_| { - // Some GGUFs don't store it explicitly; derive from tokenizer. - m.get("tokenizer.ggml.tokens") - .and_then(|v| match v { - KvValue::Array(a) => Some(a.len() as u32), - _ => None, - }) - .context("no vocab_size and no tokenizer.ggml.tokens") - })?; - let head_dim = u32_key("llama.attention.key_length") - .unwrap_or(hidden_size / num_attention_heads); + let vocab_size = u32_key("llama.vocab_size").or_else(|_| { + // Some GGUFs don't store it explicitly; derive from tokenizer. + m.get("tokenizer.ggml.tokens") + .and_then(|v| match v { + KvValue::Array(a) => Some(a.len() as u32), + _ => None, + }) + .context("no vocab_size and no tokenizer.ggml.tokens") + })?; + let head_dim = + u32_key("llama.attention.key_length").unwrap_or(hidden_size / num_attention_heads); let rope_theta = f32_key("llama.rope.freq_base").unwrap_or(10_000.0); let rope_scale = f32_key("llama.rope.scaling.factor").unwrap_or(1.0); @@ -346,7 +352,11 @@ mod tests { let c = hf_generic_config(&cfg).unwrap(); assert_eq!(c.bos_token_id, 2, "bos_token_id scalar form still works"); assert_eq!(c.eos_token_id, 1, "eos_token_id array → first element"); - assert_eq!(c.eos_token_ids, vec![106u32], "trailing eos ids land in eos_token_ids"); + assert_eq!( + c.eos_token_ids, + vec![106u32], + "trailing eos ids land in eos_token_ids" + ); } /// Llama-3 instruct: `eos_token_id: [128001, 128008, 128009]`. Primary diff --git a/base-convert/crates/base-arch/src/muse_glimmer.rs b/base-convert/crates/base-arch/src/muse_glimmer.rs index de5901b..0e83f13 100644 --- a/base-convert/crates/base-arch/src/muse_glimmer.rs +++ b/base-convert/crates/base-arch/src/muse_glimmer.rs @@ -224,8 +224,8 @@ impl GgufMapper for MuseGlimmerGgufMapper { let hidden_size = u32_req(&format!("{prefix}.embedding_length"))?; let num_hidden_layers = u32_req(&format!("{prefix}.block_count"))?; let num_attention_heads = u32_req(&format!("{prefix}.attention.head_count"))?; - let num_kv_heads = u32_key(&format!("{prefix}.attention.head_count_kv")) - .unwrap_or(num_attention_heads); + let num_kv_heads = + u32_key(&format!("{prefix}.attention.head_count_kv")).unwrap_or(num_attention_heads); let intermediate_size = u32_req(&format!("{prefix}.feed_forward_length"))?; let vocab_size = u32_key(&format!("{prefix}.vocab_size")) .or_else(|| match m.get("tokenizer.ggml.tokens") { @@ -280,8 +280,8 @@ impl GgufMapper for MuseGlimmerGgufMapper { // exported key; otherwise fall back to the published constant. config.qk_scale_factor = f32_any(&["attention.qk_scale_factor", "qk_scale_factor"]) .unwrap_or(REF_QK_SCALE_FACTOR); - config.logit_softcap = f32_any(&["final_logit_softcapping", "logit_softcap"]) - .unwrap_or(REF_LOGIT_SOFTCAP); + config.logit_softcap = + f32_any(&["final_logit_softcapping", "logit_softcap"]).unwrap_or(REF_LOGIT_SOFTCAP); config.post_norm_eps = f32_any(&["attention.post_norm_epsilon", "post_norm_eps"]).unwrap_or(REF_POST_NORM_EPS); // llama.cpp exports HF's `output_multiplier` under its own generic @@ -293,17 +293,19 @@ impl GgufMapper for MuseGlimmerGgufMapper { // The derivation below is the fallback: the value is exactly // 1/sqrt(hidden_size/256) on the released checkpoint (6656/256 = 26 // → 0.19611613…), so it beats hardcoding a width-specific number. - config.output_multiplier = f32_any(&["logit_scale", "output_multiplier"]).unwrap_or_else(|| { - let d = hidden_size as f32 / 256.0; - if d > 0.0 { - 1.0 / d.sqrt() - } else { - 0.0 - } - }); + config.output_multiplier = + f32_any(&["logit_scale", "output_multiplier"]).unwrap_or_else(|| { + let d = hidden_size as f32 / 256.0; + if d > 0.0 { + 1.0 / d.sqrt() + } else { + 0.0 + } + }); // ── Local/global + NoPE schedules ─────────────────────────── - config.sliding_window = u32_any(&["attention.sliding_window"]).unwrap_or(REF_SLIDING_WINDOW); + config.sliding_window = + u32_any(&["attention.sliding_window"]).unwrap_or(REF_SLIDING_WINDOW); let n_layers = num_hidden_layers as usize; // A bool array under `attention.sliding_window_pattern` is the // shape Gemma 4 uses and the natural one for this arch; fall back @@ -320,8 +322,12 @@ impl GgufMapper for MuseGlimmerGgufMapper { _ => None, }; let layer_types = default_layer_types(n_layers); - config.swa_layers = swa_from_kv - .unwrap_or_else(|| layer_types.iter().map(|t| t == "sliding_attention").collect()); + config.swa_layers = swa_from_kv.unwrap_or_else(|| { + layer_types + .iter() + .map(|t| t == "sliding_attention") + .collect() + }); // NoPE mask. HF encodes it as `layer_rope_theta[i] == 0`; if a // per-layer theta array ever lands in GGUF metadata, honour it. @@ -798,13 +804,25 @@ mod tests { ("v.post_ln.weight", "vision.ln_post.weight"), ("v.post_ln.bias", "vision.ln_post.bias"), // Encoder blocks — every stem, both tails. - ("v.blk.0.ln1.weight", "vision.layers.0.attention_norm.weight"), + ( + "v.blk.0.ln1.weight", + "vision.layers.0.attention_norm.weight", + ), ("v.blk.0.ln1.bias", "vision.layers.0.attention_norm.bias"), ("v.blk.2.ln2.weight", "vision.layers.2.ffn_norm.weight"), - ("v.blk.3.attn_q.weight", "vision.layers.3.attention.q.weight"), + ( + "v.blk.3.attn_q.weight", + "vision.layers.3.attention.q.weight", + ), ("v.blk.3.attn_k.bias", "vision.layers.3.attention.k.bias"), - ("v.blk.3.attn_v.weight", "vision.layers.3.attention.v.weight"), - ("v.blk.3.attn_out.weight", "vision.layers.3.attention.output.weight"), + ( + "v.blk.3.attn_v.weight", + "vision.layers.3.attention.v.weight", + ), + ( + "v.blk.3.attn_out.weight", + "vision.layers.3.attention.output.weight", + ), ("v.blk.7.ffn_up.weight", "vision.layers.7.ffn.up.weight"), ("v.blk.7.ffn_down.bias", "vision.layers.7.ffn.down.bias"), // Projector. @@ -841,12 +859,30 @@ mod tests { let pairs = [ ("vision_tower.layers.5.norm1.weight", "v.blk.5.ln1.weight"), ("vision_tower.layers.5.norm2.bias", "v.blk.5.ln2.bias"), - ("vision_tower.layers.5.attn.q_proj.weight", "v.blk.5.attn_q.weight"), - ("vision_tower.layers.5.attn.k_proj.bias", "v.blk.5.attn_k.bias"), - ("vision_tower.layers.5.attn.v_proj.weight", "v.blk.5.attn_v.weight"), - ("vision_tower.layers.5.attn.proj.weight", "v.blk.5.attn_out.weight"), - ("vision_tower.layers.5.mlp.fc1.weight", "v.blk.5.ffn_up.weight"), - ("vision_tower.layers.5.mlp.fc2.bias", "v.blk.5.ffn_down.bias"), + ( + "vision_tower.layers.5.attn.q_proj.weight", + "v.blk.5.attn_q.weight", + ), + ( + "vision_tower.layers.5.attn.k_proj.bias", + "v.blk.5.attn_k.bias", + ), + ( + "vision_tower.layers.5.attn.v_proj.weight", + "v.blk.5.attn_v.weight", + ), + ( + "vision_tower.layers.5.attn.proj.weight", + "v.blk.5.attn_out.weight", + ), + ( + "vision_tower.layers.5.mlp.fc1.weight", + "v.blk.5.ffn_up.weight", + ), + ( + "vision_tower.layers.5.mlp.fc2.bias", + "v.blk.5.ffn_down.bias", + ), ("vision_tower.ln_pre.weight", "v.pre_ln.weight"), ("vision_tower.ln_post.bias", "v.post_ln.bias"), ( @@ -879,11 +915,23 @@ mod tests { m.insert("muse-glimmer.embedding_length".into(), KvValue::U32(6656)); m.insert("muse-glimmer.block_count".into(), KvValue::U32(4)); m.insert("muse-glimmer.attention.head_count".into(), KvValue::U32(32)); - m.insert("muse-glimmer.attention.head_count_kv".into(), KvValue::U32(2)); - m.insert("muse-glimmer.feed_forward_length".into(), KvValue::U32(19968)); + m.insert( + "muse-glimmer.attention.head_count_kv".into(), + KvValue::U32(2), + ); + m.insert( + "muse-glimmer.feed_forward_length".into(), + KvValue::U32(19968), + ); m.insert("muse-glimmer.vocab_size".into(), KvValue::U32(202048)); - m.insert("muse-glimmer.attention.key_length".into(), KvValue::U32(128)); - m.insert("muse-glimmer.rope.freq_base".into(), KvValue::F32(500_000.0)); + m.insert( + "muse-glimmer.attention.key_length".into(), + KvValue::U32(128), + ); + m.insert( + "muse-glimmer.rope.freq_base".into(), + KvValue::F32(500_000.0), + ); m.insert("muse-glimmer.context_length".into(), KvValue::U32(131072)); m.insert( "muse-glimmer.attention.layer_norm_rms_epsilon".into(), @@ -903,8 +951,7 @@ mod tests { .expect("hyphenated `muse-glimmer` must resolve"); assert_eq!(m.canonical_arch(), "muse_glimmer"); assert_eq!( - crate::source_mapper_for_gguf("muse_glimmer") - .map(|m| m.canonical_arch()), + crate::source_mapper_for_gguf("muse_glimmer").map(|m| m.canonical_arch()), Some("muse_glimmer"), "underscored spelling accepted too" ); @@ -912,7 +959,9 @@ mod tests { #[test] fn gguf_config_matches_hf_config() { - let g = MuseGlimmerGgufMapper.config_from_gguf(&gguf_meta()).unwrap(); + let g = MuseGlimmerGgufMapper + .config_from_gguf(&gguf_meta()) + .unwrap(); let h = MuseGlimmerHfMapper.config_from_hf(&cfg()).unwrap(); assert_eq!(g.hidden_size, h.hidden_size); assert_eq!(g.num_hidden_layers, h.num_hidden_layers); @@ -935,7 +984,9 @@ mod tests { /// hidden_size rather than hardcoded. #[test] fn gguf_falls_back_to_reference_scales() { - let g = MuseGlimmerGgufMapper.config_from_gguf(&gguf_meta()).unwrap(); + let g = MuseGlimmerGgufMapper + .config_from_gguf(&gguf_meta()) + .unwrap(); assert_eq!(g.qk_scale_factor, 3.87); assert_eq!(g.logit_softcap, 20.0); assert_eq!(g.post_norm_eps, 1e-8); @@ -956,7 +1007,10 @@ mod tests { "muse-glimmer.final_logit_softcapping".into(), KvValue::F32(30.0), ); - m.insert("muse-glimmer.attention.sliding_window".into(), KvValue::U32(1024)); + m.insert( + "muse-glimmer.attention.sliding_window".into(), + KvValue::U32(1024), + ); let g = MuseGlimmerGgufMapper.config_from_gguf(&m).unwrap(); assert_eq!(g.qk_scale_factor, 2.5); assert_eq!(g.output_multiplier, 0.5); @@ -980,7 +1034,10 @@ mod tests { "blk.7.post_attention_norm.weight", "layers.7.post_attention_norm.weight", ), - ("blk.7.post_ffw_norm.weight", "layers.7.post_ffw_norm.weight"), + ( + "blk.7.post_ffw_norm.weight", + "layers.7.post_ffw_norm.weight", + ), ("blk.1.attn_q.weight", "layers.1.self_attn.q_proj.weight"), ("blk.1.attn_k.weight", "layers.1.self_attn.k_proj.weight"), ("blk.1.attn_v.weight", "layers.1.self_attn.v_proj.weight"), @@ -989,9 +1046,18 @@ mod tests { "layers.1.self_attn.o_proj.weight", ), // NOT `self_attn.gate.weight` — see map_gguf_name's docs. - ("blk.1.attn_gate.weight", "layers.1.self_attn.gate_proj.weight"), - ("blk.1.attn_q_norm.weight", "layers.1.self_attn.q_norm.weight"), - ("blk.1.attn_k_norm.weight", "layers.1.self_attn.k_norm.weight"), + ( + "blk.1.attn_gate.weight", + "layers.1.self_attn.gate_proj.weight", + ), + ( + "blk.1.attn_q_norm.weight", + "layers.1.self_attn.q_norm.weight", + ), + ( + "blk.1.attn_k_norm.weight", + "layers.1.self_attn.k_norm.weight", + ), ("blk.51.ffn_gate.weight", "layers.51.mlp.gate_proj.weight"), ("blk.51.ffn_up.weight", "layers.51.mlp.up_proj.weight"), ("blk.51.ffn_down.weight", "layers.51.mlp.down_proj.weight"), @@ -1029,7 +1095,9 @@ mod tests { /// — v, o, the attention gate, the FFN — must be left alone. #[test] fn rope_unpermute_targets_q_and_k_only() { - let c = MuseGlimmerGgufMapper.config_from_gguf(&gguf_meta()).unwrap(); + let c = MuseGlimmerGgufMapper + .config_from_gguf(&gguf_meta()) + .unwrap(); let m = MuseGlimmerGgufMapper; assert_eq!( m.rope_unpermute_heads("layers.0.self_attn.q_proj.weight", &c), diff --git a/base-convert/crates/base-arch/src/nemotron.rs b/base-convert/crates/base-arch/src/nemotron.rs new file mode 100644 index 0000000..7c4ab1a --- /dev/null +++ b/base-convert/crates/base-arch/src/nemotron.rs @@ -0,0 +1,600 @@ +//! Nemotron-H GGUF → canonical `.base` mapping. +//! +//! Covers the hybrid Mamba-2 + attention decoder in both its MoE form +//! (`nemotron_h_moe`, e.g. Nemotron 3 Nano 30B-A3B) and the dense form +//! (`nemotron_h`). Unlike a standard transformer, each block has exactly +//! ONE mixer: a Mamba-2 SSM, a GQA attention, or an (MoE) FFN. The GGUF +//! metadata encodes the schedule via per-layer arrays: +//! +//! `{arch}.attention.head_count_kv` — nonzero ⇒ attention block +//! `{arch}.feed_forward_length` — nonzero ⇒ FFN (MoE) block +//! both zero ⇒ Mamba-2 block +//! +//! The schedule is emitted as `layer_types` ("mamba" | "attention" | +//! "moe" | "mlp") in the `.base` header config, following the Qwen3.5 +//! hybrid precedent of flat config scalars + a `layer_types` array. +//! +//! MoE FFN blocks use DeepSeek-style routing: sigmoid scores + a +//! selection-only correction bias (`exp_probs_b.bias`), top-k with +//! renormalization (`expert_weights_norm`) and a routed scaling factor +//! (`expert_weights_scale`), plus one always-on shared expert. The FFN +//! itself is squared-ReLU up/down — there is no `ffn_gate` tensor. + +use crate::{ArchConfig, GgufMapper}; +use anyhow::{Context, Result}; +use base_readers::gguf::KvValue; +use std::collections::BTreeMap; + +pub struct NemotronHMapper; + +fn u32_array(m: &BTreeMap, k: &str) -> Option> { + match m.get(k)? { + KvValue::Array(a) => Some( + a.iter() + .filter_map(|v| v.as_u64()) + .map(|n| n as u32) + .collect(), + ), + _ => None, + } +} + +impl GgufMapper for NemotronHMapper { + fn canonical_arch(&self) -> &'static str { + "nemotron_h_moe" + } + + fn config_from_gguf(&self, m: &BTreeMap) -> Result { + let prefix = if m.keys().any(|k| k.starts_with("nemotron_h_moe.")) { + "nemotron_h_moe" + } else { + "nemotron_h" + }; + + let u32_key = |k: &str| { + m.get(&format!("{prefix}.{k}")) + .and_then(|v| v.as_u64()) + .map(|n| n as u32) + .with_context(|| format!("missing metadata key: {prefix}.{k}")) + }; + let f32_key = |k: &str| m.get(&format!("{prefix}.{k}")).and_then(|v| v.as_f32()); + let bool_key = |k: &str| { + m.get(&format!("{prefix}.{k}")).and_then(|v| match v { + KvValue::Bool(b) => Some(*b), + _ => None, + }) + }; + + let hidden_size = u32_key("embedding_length")?; + let num_hidden_layers = u32_key("block_count")?; + let num_attention_heads = u32_key("attention.head_count")?; + + // Per-layer arrays encode the block schedule (see module doc). + // `head_count_kv` is an array on hybrid checkpoints; tolerate a + // scalar for hypothetical homogeneous ones. + let kv_per_layer = u32_array(m, &format!("{prefix}.attention.head_count_kv")) + .unwrap_or_else(|| { + let scalar = u32_key("attention.head_count_kv").unwrap_or(num_attention_heads); + vec![scalar; num_hidden_layers as usize] + }); + let ffn_per_layer = + u32_array(m, &format!("{prefix}.feed_forward_length")).unwrap_or_else(|| { + let scalar = u32_key("feed_forward_length").unwrap_or(0); + vec![scalar; num_hidden_layers as usize] + }); + let num_kv_heads = kv_per_layer.iter().copied().max().unwrap_or(0); + + // MoE topology. + let num_experts = u32_key("expert_count").unwrap_or(0); + let num_experts_per_tok = u32_key("expert_used_count").unwrap_or(0); + let moe_intermediate_size = u32_key("expert_feed_forward_length").unwrap_or(0); + let num_shared_experts = u32_key("expert_shared_count").unwrap_or(0); + + let layer_types: Vec = (0..num_hidden_layers as usize) + .map(|i| { + if kv_per_layer.get(i).copied().unwrap_or(0) > 0 { + "attention" + } else if ffn_per_layer.get(i).copied().unwrap_or(0) > 0 { + if num_experts > 0 { + "moe" + } else { + "mlp" + } + } else { + "mamba" + } + .to_string() + }) + .collect(); + + // Dense-slot FFN width: the shared expert on MoE checkpoints + // (Qwen3.5-MoE precedent: `intermediate_size` = shared/dense + // width, `moe_intermediate_size` = per-routed-expert width); + // the widest per-layer FFN otherwise. + let intermediate_size = u32_key("expert_shared_feed_forward_length") + .ok() + .filter(|&v| v > 0) + .or_else(|| ffn_per_layer.iter().copied().max().filter(|&v| v > 0)) + .context("neither expert_shared_feed_forward_length nor feed_forward_length set")?; + + let vocab_size = u32_key("vocab_size").or_else(|_| { + m.get("tokenizer.ggml.tokens") + .and_then(|v| match v { + KvValue::Array(a) => Some(a.len() as u32), + _ => None, + }) + .context("no vocab_size and no tokenizer.ggml.tokens") + })?; + let head_dim = u32_key("attention.key_length").unwrap_or(hidden_size / num_attention_heads); + + let rope_theta = f32_key("rope.freq_base").unwrap_or(10_000.0); + let rope_scale = f32_key("rope.scaling.factor").unwrap_or(1.0); + let rms_norm_eps = f32_key("attention.layer_norm_rms_epsilon").unwrap_or(1e-6); + // Partial RoPE on the attention blocks: only the first + // `rope.dimension_count` of `head_dim` dims rotate (Nemotron 3 + // Nano: 84 / 128 = 0.65625). + let partial_rotary_factor = match u32_key("rope.dimension_count") { + Ok(d) if d > 0 && d < head_dim => d as f32 / head_dim as f32, + _ => 0.0, + }; + let max_position_embeddings = u32_key("context_length").unwrap_or(0); + + // Mamba-2 mixer geometry. `time_step_rank` carries the SSM head + // count on Mamba-2 GGUFs (llama.cpp convention). + let ssm_state_size = u32_key("ssm.state_size").unwrap_or(0); + let ssm_conv_kernel = u32_key("ssm.conv_kernel").unwrap_or(0); + let ssm_num_groups = u32_key("ssm.group_count").unwrap_or(0); + let ssm_inner_size = u32_key("ssm.inner_size").unwrap_or(0); + let ssm_num_heads = u32_key("ssm.time_step_rank").unwrap_or(0); + + Ok(ArchConfig { + hidden_size, + num_hidden_layers, + num_attention_heads, + num_kv_heads, + head_dim, + intermediate_size, + vocab_size, + rope_theta, + rope_scale, + rms_norm_eps, + partial_rotary_factor, + max_position_embeddings, + tie_word_embeddings: false, + n_kv_heads_per_layer: kv_per_layer, + layer_types, + num_experts, + num_experts_per_tok, + moe_intermediate_size, + num_shared_experts, + // Sigmoid routing with renormalized top-k and routed scaling + // (DeepSeek-style). `expert_weights_norm` defaults true for + // this family. + norm_topk_prob: bool_key("expert_weights_norm").unwrap_or(true), + expert_gating: if num_experts > 0 { 1 } else { 0 }, + expert_weights_scale: f32_key("expert_weights_scale").unwrap_or(0.0), + ssm_state_size, + ssm_conv_kernel, + ssm_num_groups, + ssm_inner_size, + ssm_num_heads, + ..ArchConfig::default() + }) + } + + fn map_tensor_name(&self, n: &str) -> Option { + // Router selection-bias (DeepSeek-style `e_score_correction_bias`) + // — not part of the shared llama-style table. + if let Some(rest) = n.strip_prefix("blk.") { + if let Some((layer_str, suffix)) = rest.split_once('.') { + if suffix == "exp_probs_b.bias" { + let layer: u32 = layer_str.parse().ok()?; + return Some(format!("layers.{layer}.mlp.router.e_score_correction_bias")); + } + } + } + crate::llama::map_llama_style(n) + } +} + +/// Nemotron-H HF / MLX-safetensors → canonical `.base` mapping. +/// +/// The HF checkpoint names everything under `backbone.` and gives every +/// block the same `mixer.` prefix regardless of what the mixer *is* — +/// the block schedule lives in `hybrid_override_pattern`, not in the +/// tensor names. So `mixer.in_proj` (Mamba), `mixer.q_proj` (attention) +/// and `mixer.gate` (MoE router) are siblings, and the rename table +/// below is what splits them back into the canonical `ssm.*`, +/// `self_attn.*` and `mlp.*` families the runtime expects. +/// +/// Targets exactly the names the GGUF mapper emits, so a bundle +/// converted from an MLX checkpoint and one converted from the GGUF are +/// interchangeable as far as the runtime is concerned. +pub struct NemotronHHfMapper; + +/// HF/MLX tensor name → canonical `.base` name. `None` drops the tensor. +pub fn nemotron_hf_rename(name: &str) -> Option { + match name { + "backbone.embeddings.weight" => return Some("embed_tokens.weight".into()), + "backbone.norm_f.weight" => return Some("final_norm.weight".into()), + "lm_head.weight" => return Some("lm_head.weight".into()), + _ => {} + } + + let rest = name.strip_prefix("backbone.layers.")?; + let (layer_str, tail) = rest.split_once('.')?; + let layer: u32 = layer_str.parse().ok()?; + + // Every block's pre-mixer norm. Named `norm.weight` directly under + // the layer — distinct from `mixer.norm.weight`, which is the + // Mamba-2 grouped gated norm *inside* the SSM mixer. + if tail == "norm.weight" { + return Some(format!("layers.{layer}.input_norm.weight")); + } + + let mixer = tail.strip_prefix("mixer.")?; + let canonical_tail = match mixer { + // Attention blocks. + "q_proj.weight" => "self_attn.q_proj.weight", + "k_proj.weight" => "self_attn.k_proj.weight", + "v_proj.weight" => "self_attn.v_proj.weight", + "o_proj.weight" => "self_attn.o_proj.weight", + + // Mamba-2 blocks. + "in_proj.weight" => "ssm.in_proj.weight", + "out_proj.weight" => "ssm.out_proj.weight", + "conv1d.weight" => "ssm.conv1d.weight", + "conv1d.bias" => "ssm.conv1d.bias", + "A_log" => "ssm.a_log", + "D" => "ssm.d", + "dt_bias" => "ssm.dt_bias", + // The grouped RMSNorm applied to the scan output. + "norm.weight" => "ssm.norm.weight", + + // MoE blocks. HF calls the router `gate`; the runtime's `mlp.router` + // is the same matrix. `switch_mlp.fc1/fc2` are MLX's stacked + // per-expert projections — this family has no gate projection + // (squared-ReLU, not SwiGLU), so fc1/fc2 are up/down. + "gate.weight" => "mlp.router.weight", + "gate.e_score_correction_bias" => "mlp.router.e_score_correction_bias", + "switch_mlp.fc1.weight" => "mlp.experts.up_proj.weight", + "switch_mlp.fc2.weight" => "mlp.experts.down_proj.weight", + "shared_experts.up_proj.weight" => "mlp.shared_expert.up_proj.weight", + "shared_experts.down_proj.weight" => "mlp.shared_expert.down_proj.weight", + + // NVFP4 checkpoints (per-expert tensors arrive pre-stacked as the + // virtual `experts.` names). The fp4 code bytes + e4m3 block + // scales are transplanted into the `.weight` tensor; the global + // scale and the calibrated activation scale ride along as f32 + // sidecar tensors. + "experts.up_proj.weight" => "mlp.experts.up_proj.weight", + "experts.down_proj.weight" => "mlp.experts.down_proj.weight", + "experts.up_proj.weight_scale_2" => "mlp.experts.up_proj.weight_scale_2", + "experts.down_proj.weight_scale_2" => "mlp.experts.down_proj.weight_scale_2", + "experts.up_proj.input_scale" => "mlp.experts.up_proj.input_scale", + "experts.down_proj.input_scale" => "mlp.experts.down_proj.input_scale", + "in_proj.weight_scale_2" => "ssm.in_proj.weight_scale_2", + "in_proj.input_scale" => "ssm.in_proj.input_scale", + "out_proj.weight_scale_2" => "ssm.out_proj.weight_scale_2", + "out_proj.input_scale" => "ssm.out_proj.input_scale", + "shared_experts.up_proj.weight_scale_2" => "mlp.shared_expert.up_proj.weight_scale_2", + "shared_experts.down_proj.weight_scale_2" => "mlp.shared_expert.down_proj.weight_scale_2", + "shared_experts.up_proj.input_scale" => "mlp.shared_expert.up_proj.input_scale", + "shared_experts.down_proj.input_scale" => "mlp.shared_expert.down_proj.input_scale", + + _ => return None, + }; + Some(format!("layers.{layer}.{canonical_tail}")) +} + +impl crate::HfMapper for NemotronHHfMapper { + fn canonical_arch(&self) -> &'static str { + "nemotron_h_moe" + } + + fn value_transform(&self, canonical: &str) -> Option { + // The HF checkpoint stores the Mamba-2 state-transition matrix + // as `A_log`; `modeling_nemotron_h.py` computes + // `A = -exp(A_log)` on every forward. The GGUF conversion bakes + // that in, and the runtime's scan kernel consumes `A` directly + // — it expects negative values, since `dA = exp(dt * A)` has to + // decay. Copying `A_log` through verbatim leaves the scan with + // positive transitions and the model emits one token forever. + if canonical.ends_with(".ssm.a_log") { + Some(crate::ValueTransform::NegExp) + } else { + None + } + } + + fn shape_fastest_first(&self) -> bool { + // Match the GGUF mapper's `ne`-style shape notation ([in, out]) + // so both conversion paths describe this architecture the same + // way. The bytes are identical either way — HF C-order + // [out, in] and GGUF [in, out] are the same buffer, in + // contiguous along the reduction dim — this only fixes how the + // header *reports* it. + true + } + + fn config_from_hf(&self, c: &serde_json::Value) -> Result { + let u32_key = |k: &str| c.get(k).and_then(|v| v.as_u64()).map(|n| n as u32); + let f32_key = |k: &str| c.get(k).and_then(|v| v.as_f64()).map(|n| n as f32); + let req = |k: &str| u32_key(k).with_context(|| format!("config.json missing {k}")); + + let hidden_size = req("hidden_size")?; + let num_hidden_layers = req("num_hidden_layers")?; + let num_attention_heads = req("num_attention_heads")?; + let num_kv_heads = u32_key("num_key_value_heads").unwrap_or(num_attention_heads); + let head_dim = u32_key("head_dim").unwrap_or(hidden_size / num_attention_heads); + + // The block schedule. `hybrid_override_pattern` is one character + // per layer: M = Mamba-2, * = attention, E/- = the FFN slot. + let pattern = c + .get("hybrid_override_pattern") + .and_then(|v| v.as_str()) + .context("config.json missing hybrid_override_pattern (the block schedule)")?; + if pattern.chars().count() != num_hidden_layers as usize { + anyhow::bail!( + "hybrid_override_pattern has {} entries but num_hidden_layers is {}", + pattern.chars().count(), + num_hidden_layers + ); + } + let num_experts = u32_key("n_routed_experts").unwrap_or(0); + let layer_types: Vec = pattern + .chars() + .map(|ch| { + match ch { + 'M' => "mamba", + '*' => "attention", + _ if num_experts > 0 => "moe", + _ => "mlp", + } + .to_string() + }) + .collect(); + let n_kv_heads_per_layer: Vec = layer_types + .iter() + .map(|t| if t == "attention" { num_kv_heads } else { 0 }) + .collect(); + + // `intermediate_size` is the *routed* expert width on this + // checkpoint; the dense-slot width the runtime wants is the + // shared expert's (Qwen3.5-MoE precedent, and what the GGUF + // mapper reads out of expert_shared_feed_forward_length). + let moe_intermediate_size = u32_key("moe_intermediate_size") + .or_else(|| u32_key("intermediate_size")) + .unwrap_or(0); + let intermediate_size = u32_key("moe_shared_expert_intermediate_size") + .or_else(|| u32_key("intermediate_size")) + .context("neither moe_shared_expert_intermediate_size nor intermediate_size set")?; + + // Mamba-2 mixer geometry. `n_groups` is the SSM group count; + // `n_group` (singular) is the DeepSeek routing group count and + // is a different thing entirely — reading the wrong one gives 1 + // group instead of 8 and silently mis-shapes the scan. + let ssm_num_heads = u32_key("mamba_num_heads").unwrap_or(0); + let mamba_head_dim = u32_key("mamba_head_dim").unwrap_or(0); + + Ok(ArchConfig { + hidden_size, + num_hidden_layers, + num_attention_heads, + num_kv_heads, + head_dim, + intermediate_size, + vocab_size: req("vocab_size")?, + rope_theta: f32_key("rope_theta").unwrap_or(10_000.0), + rope_scale: 1.0, + rms_norm_eps: f32_key("norm_eps") + .or_else(|| f32_key("layer_norm_epsilon")) + .unwrap_or(1e-5), + // Nemotron-H attention is NoPE — llama.cpp maps this + // architecture to rope_type NONE regardless of the rope keys + // the config carries, and the reference generations match + // that. Recorded for completeness; the runtime does not + // rotate these blocks. + partial_rotary_factor: f32_key("partial_rotary_factor").unwrap_or(0.0), + max_position_embeddings: u32_key("max_position_embeddings").unwrap_or(0), + tie_word_embeddings: c + .get("tie_word_embeddings") + .and_then(|v| v.as_bool()) + .unwrap_or(false), + n_kv_heads_per_layer, + layer_types, + num_experts, + num_experts_per_tok: u32_key("num_experts_per_tok").unwrap_or(0), + moe_intermediate_size, + num_shared_experts: u32_key("n_shared_experts").unwrap_or(0), + norm_topk_prob: c + .get("norm_topk_prob") + .and_then(|v| v.as_bool()) + .unwrap_or(true), + expert_gating: if num_experts > 0 { 1 } else { 0 }, + expert_weights_scale: f32_key("routed_scaling_factor").unwrap_or(0.0), + ssm_state_size: u32_key("ssm_state_size").unwrap_or(0), + ssm_conv_kernel: u32_key("conv_kernel").unwrap_or(0), + ssm_num_groups: u32_key("n_groups").unwrap_or(0), + ssm_inner_size: ssm_num_heads * mamba_head_dim, + ssm_num_heads, + ..ArchConfig::default() + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Synthetic metadata mirroring the real Nemotron-3-Nano-30B-A3B + /// GGUF (52 blocks: 23 Mamba-2, 23 MoE FFN, 6 attention; per-layer + /// arrays carry the schedule). + fn nano_metadata() -> BTreeMap { + let mut m = BTreeMap::new(); + let p = "nemotron_h_moe"; + let u = KvValue::U32; + m.insert("general.architecture".into(), KvValue::String(p.into())); + m.insert(format!("{p}.block_count"), u(8)); + m.insert(format!("{p}.embedding_length"), u(2688)); + m.insert(format!("{p}.attention.head_count"), u(32)); + m.insert(format!("{p}.attention.key_length"), u(128)); + m.insert(format!("{p}.attention.value_length"), u(128)); + m.insert( + format!("{p}.attention.head_count_kv"), + KvValue::Array(vec![u(0), u(0), u(2), u(0), u(0), u(0), u(2), u(0)]), + ); + m.insert( + format!("{p}.feed_forward_length"), + KvValue::Array(vec![ + u(0), + u(1856), + u(0), + u(1856), + u(0), + u(1856), + u(0), + u(1856), + ]), + ); + m.insert(format!("{p}.vocab_size"), u(131072)); + m.insert(format!("{p}.context_length"), u(1048576)); + m.insert(format!("{p}.rope.freq_base"), KvValue::F32(10000.0)); + m.insert(format!("{p}.rope.dimension_count"), u(84)); + m.insert( + format!("{p}.attention.layer_norm_rms_epsilon"), + KvValue::F32(1e-5), + ); + m.insert(format!("{p}.expert_count"), u(128)); + m.insert(format!("{p}.expert_used_count"), u(6)); + m.insert(format!("{p}.expert_feed_forward_length"), u(1856)); + m.insert(format!("{p}.expert_shared_feed_forward_length"), u(3712)); + m.insert(format!("{p}.expert_shared_count"), u(1)); + m.insert(format!("{p}.expert_weights_norm"), KvValue::Bool(true)); + m.insert(format!("{p}.expert_weights_scale"), KvValue::F32(2.5)); + m.insert(format!("{p}.ssm.conv_kernel"), u(4)); + m.insert(format!("{p}.ssm.state_size"), u(128)); + m.insert(format!("{p}.ssm.group_count"), u(8)); + m.insert(format!("{p}.ssm.inner_size"), u(4096)); + m.insert(format!("{p}.ssm.time_step_rank"), u(64)); + m + } + + #[test] + fn nano_config_from_gguf() { + let c = NemotronHMapper.config_from_gguf(&nano_metadata()).unwrap(); + assert_eq!(c.hidden_size, 2688); + assert_eq!(c.num_hidden_layers, 8); + assert_eq!(c.num_attention_heads, 32); + assert_eq!(c.num_kv_heads, 2, "max of the per-layer kv-head array"); + assert_eq!(c.head_dim, 128); + assert_eq!( + c.intermediate_size, 3712, + "dense slot = shared-expert width on MoE checkpoints" + ); + assert_eq!(c.moe_intermediate_size, 1856); + assert_eq!(c.num_experts, 128); + assert_eq!(c.num_experts_per_tok, 6); + assert_eq!(c.num_shared_experts, 1); + assert!(c.norm_topk_prob); + assert_eq!(c.expert_gating, 1, "sigmoid routing"); + assert_eq!(c.expert_weights_scale, 2.5); + assert_eq!( + c.layer_types, + vec![ + "mamba", + "moe", + "attention", + "moe", + "mamba", + "moe", + "attention", + "moe" + ] + ); + assert_eq!( + c.n_kv_heads_per_layer, + vec![0, 0, 2, 0, 0, 0, 2, 0], + "schedule array preserved verbatim" + ); + assert_eq!(c.ssm_state_size, 128); + assert_eq!(c.ssm_conv_kernel, 4); + assert_eq!(c.ssm_num_groups, 8); + assert_eq!(c.ssm_inner_size, 4096); + assert_eq!(c.ssm_num_heads, 64); + assert!((c.partial_rotary_factor - 0.65625).abs() < 1e-6); + assert_eq!(c.max_position_embeddings, 1048576); + } + + #[test] + fn config_map_carries_ssm_and_routing_keys() { + let c = NemotronHMapper.config_from_gguf(&nano_metadata()).unwrap(); + let map = c.to_config_map(); + assert_eq!(map["ssm_state_size"], serde_json::json!(128)); + assert_eq!(map["ssm_conv_kernel"], serde_json::json!(4)); + assert_eq!(map["ssm_num_groups"], serde_json::json!(8)); + assert_eq!(map["ssm_inner_size"], serde_json::json!(4096)); + assert_eq!(map["ssm_num_heads"], serde_json::json!(64)); + assert_eq!(map["expert_gating"], serde_json::json!(1)); + assert_eq!(map["expert_weights_scale"], serde_json::json!(2.5)); + assert_eq!(map["num_shared_experts"], serde_json::json!(1)); + assert!(map.contains_key("layer_types")); + assert!(map.contains_key("n_kv_heads_per_layer")); + } + + /// Every tensor pattern present in the real Nano GGUF must map — an + /// unmapped name is silently dropped by convert_gguf, so this list is + /// the converter-side completeness gate. + #[test] + fn nano_tensor_names_all_map() { + let cases = [ + ("token_embd.weight", "embed_tokens.weight"), + ("output.weight", "lm_head.weight"), + ("output_norm.weight", "final_norm.weight"), + ("blk.0.attn_norm.weight", "layers.0.input_norm.weight"), + ("blk.5.attn_q.weight", "layers.5.self_attn.q_proj.weight"), + ("blk.5.attn_k.weight", "layers.5.self_attn.k_proj.weight"), + ("blk.5.attn_v.weight", "layers.5.self_attn.v_proj.weight"), + ( + "blk.5.attn_output.weight", + "layers.5.self_attn.o_proj.weight", + ), + ( + "blk.1.ffn_up_exps.weight", + "layers.1.mlp.experts.up_proj.weight", + ), + ( + "blk.1.ffn_down_exps.weight", + "layers.1.mlp.experts.down_proj.weight", + ), + ("blk.1.ffn_gate_inp.weight", "layers.1.mlp.router.weight"), + ( + "blk.1.exp_probs_b.bias", + "layers.1.mlp.router.e_score_correction_bias", + ), + ( + "blk.1.ffn_up_shexp.weight", + "layers.1.mlp.shared_expert.up_proj.weight", + ), + ( + "blk.1.ffn_down_shexp.weight", + "layers.1.mlp.shared_expert.down_proj.weight", + ), + ("blk.0.ssm_in.weight", "layers.0.ssm.in_proj.weight"), + ("blk.0.ssm_out.weight", "layers.0.ssm.out_proj.weight"), + ("blk.0.ssm_conv1d.weight", "layers.0.ssm.conv1d.weight"), + ("blk.0.ssm_conv1d.bias", "layers.0.ssm.conv1d.bias"), + ("blk.0.ssm_dt.bias", "layers.0.ssm.dt_bias"), + ("blk.0.ssm_a", "layers.0.ssm.a_log"), + ("blk.0.ssm_d", "layers.0.ssm.d"), + ("blk.0.ssm_norm.weight", "layers.0.ssm.norm.weight"), + ]; + for (gguf, canonical) in cases { + assert_eq!( + NemotronHMapper.map_tensor_name(gguf).as_deref(), + Some(canonical), + "mapping for {gguf}" + ); + } + } +} diff --git a/base-convert/crates/base-arch/src/qwen.rs b/base-convert/crates/base-arch/src/qwen.rs index 7cafc7a..0679c1d 100644 --- a/base-convert/crates/base-arch/src/qwen.rs +++ b/base-convert/crates/base-arch/src/qwen.rs @@ -234,24 +234,22 @@ impl GgufMapper for QwenMapper { let hidden_size = u32_key(&format!("{prefix}.embedding_length"))?; let num_hidden_layers = u32_key(&format!("{prefix}.block_count"))?; let num_attention_heads = u32_key(&format!("{prefix}.attention.head_count"))?; - let num_kv_heads = u32_key(&format!("{prefix}.attention.head_count_kv")) - .unwrap_or(num_attention_heads); + let num_kv_heads = + u32_key(&format!("{prefix}.attention.head_count_kv")).unwrap_or(num_attention_heads); let intermediate_size = u32_key(&format!("{prefix}.feed_forward_length"))?; - let vocab_size = u32_key(&format!("{prefix}.vocab_size")) - .or_else(|_| { - m.get("tokenizer.ggml.tokens") - .and_then(|v| match v { - KvValue::Array(a) => Some(a.len() as u32), - _ => None, - }) - .context("no vocab_size and no tokenizer.ggml.tokens") - })?; + let vocab_size = u32_key(&format!("{prefix}.vocab_size")).or_else(|_| { + m.get("tokenizer.ggml.tokens") + .and_then(|v| match v { + KvValue::Array(a) => Some(a.len() as u32), + _ => None, + }) + .context("no vocab_size and no tokenizer.ggml.tokens") + })?; let head_dim = u32_key(&format!("{prefix}.attention.key_length")) .unwrap_or(hidden_size / num_attention_heads); let rope_theta = f32_key(&format!("{prefix}.rope.freq_base")).unwrap_or(10_000.0); - let rope_scale = - f32_key(&format!("{prefix}.rope.scaling.factor")).unwrap_or(1.0); + let rope_scale = f32_key(&format!("{prefix}.rope.scaling.factor")).unwrap_or(1.0); let rms_norm_eps = f32_key(&format!("{prefix}.attention.layer_norm_rms_epsilon")).unwrap_or(1e-6); @@ -308,8 +306,8 @@ impl GgufMapper for QwenMoeMapper { let hidden_size = u32_key(&format!("{prefix}.embedding_length"))?; let num_hidden_layers = u32_key(&format!("{prefix}.block_count"))?; let num_attention_heads = u32_key(&format!("{prefix}.attention.head_count"))?; - let num_kv_heads = u32_key(&format!("{prefix}.attention.head_count_kv")) - .unwrap_or(num_attention_heads); + let num_kv_heads = + u32_key(&format!("{prefix}.attention.head_count_kv")).unwrap_or(num_attention_heads); // For MoE, GGUF carries two FFN widths: // `feed_forward_length` = nominal/dense width (HF `intermediate_size`) // `expert_feed_forward_length` = per-expert width (HF `moe_intermediate_size`) diff --git a/base-convert/crates/base-arch/src/whisper.rs b/base-convert/crates/base-arch/src/whisper.rs index db49dd9..97c7ea9 100644 --- a/base-convert/crates/base-arch/src/whisper.rs +++ b/base-convert/crates/base-arch/src/whisper.rs @@ -826,9 +826,15 @@ mod tests { num_hidden_layers: dec, ..ArchConfig::default() }; - assert!(!supports_translate(&mk(32, 4)), "large-v3-turbo must drop translate"); + assert!( + !supports_translate(&mk(32, 4)), + "large-v3-turbo must drop translate" + ); for (enc, dec) in [(4u32, 4u32), (6, 6), (12, 12), (24, 24), (32, 32)] { - assert!(supports_translate(&mk(enc, dec)), "{enc}/{dec} keeps translate"); + assert!( + supports_translate(&mk(enc, dec)), + "{enc}/{dec} keeps translate" + ); } } diff --git a/base-convert/crates/base-awq/src/lib.rs b/base-convert/crates/base-awq/src/lib.rs index c900712..389b0ed 100644 --- a/base-convert/crates/base-awq/src/lib.rs +++ b/base-convert/crates/base-awq/src/lib.rs @@ -230,8 +230,8 @@ fn rtn_reconstruction_mse( group_size: u32, symmetric: bool, ) -> f32 { - use base_quant::{pack_rtn, unpack_rtn, RtnConfig}; use base_format::ScaleDtype; + use base_quant::{pack_rtn, unpack_rtn, RtnConfig}; let cfg = RtnConfig { bits, @@ -281,9 +281,8 @@ mod tests { let plan = AwqConfig::default().search(&weights, n_in, &absmax, 4, 64, false); let identity: Vec = vec![1.0; n_in]; - let plain_mse = rtn_reconstruction_mse( - &weights, &weights, n_in, n_out, &identity, 4, 64, false, - ); + let plain_mse = + rtn_reconstruction_mse(&weights, &weights, n_in, n_out, &identity, 4, 64, false); // Tiny tolerance for floating-point order-of-ops. assert!( plan.mse <= plain_mse * 1.0001, @@ -314,9 +313,8 @@ mod tests { let plan = AwqConfig::default().search(&weights, n_in, &absmax, 2, 32, false); let identity: Vec = vec![1.0; n_in]; - let plain_mse = rtn_reconstruction_mse( - &weights, &weights, n_in, n_out, &identity, 2, 32, false, - ); + let plain_mse = + rtn_reconstruction_mse(&weights, &weights, n_in, n_out, &identity, 2, 32, false); assert!( plan.mse < plain_mse, "AWQ should help: awq={} plain={}", diff --git a/base-convert/crates/base-awq/src/sidecar.rs b/base-convert/crates/base-awq/src/sidecar.rs index 2a35266..3f3885b 100644 --- a/base-convert/crates/base-awq/src/sidecar.rs +++ b/base-convert/crates/base-awq/src/sidecar.rs @@ -20,15 +20,13 @@ impl AwqProfile { pub fn load(path: &Path) -> Result { let bytes = std::fs::read(path) .with_context(|| format!("reading AWQ profile {}", path.display()))?; - let p: AwqProfile = serde_json::from_slice(&bytes) - .context("parsing AWQ profile JSON")?; + let p: AwqProfile = serde_json::from_slice(&bytes).context("parsing AWQ profile JSON")?; Ok(p) } /// Write the profile to `path` (canonical JSON, sorted keys). pub fn save(&self, path: &Path) -> Result<()> { - let json = serde_json::to_vec_pretty(self) - .context("serializing AWQ profile")?; + let json = serde_json::to_vec_pretty(self).context("serializing AWQ profile")?; std::fs::write(path, json) .with_context(|| format!("writing AWQ profile {}", path.display()))?; Ok(()) @@ -38,7 +36,9 @@ impl AwqProfile { /// Returns None if the profile lacks an entry — callers fall back /// to plain RTN for that tensor. pub fn absmax(&self, tensor_name: &str) -> Option<&[f32]> { - self.per_tensor_absmax.get(tensor_name).map(|v| v.as_slice()) + self.per_tensor_absmax + .get(tensor_name) + .map(|v| v.as_slice()) } /// Validate the profile against an expected source fingerprint diff --git a/base-convert/crates/base-awq/src/wikitext.rs b/base-convert/crates/base-awq/src/wikitext.rs index 3f6a0dd..2a2e9cd 100644 --- a/base-convert/crates/base-awq/src/wikitext.rs +++ b/base-convert/crates/base-awq/src/wikitext.rs @@ -120,9 +120,8 @@ mod tests { #[test] fn skips_empty_lines_and_headings() { - let f = write_temp( - " \n = Title = \n \nFirst paragraph.\n \n = = Subsection = = \nSecond.\n", - ); + let f = + write_temp(" \n = Title = \n \nFirst paragraph.\n \n = = Subsection = = \nSecond.\n"); let mut r = WikiTextReader::open(f.path()).unwrap(); let text = r.read_n_lines(5).unwrap(); assert!(text.contains("First paragraph.")); @@ -149,7 +148,11 @@ mod tests { let mut r = WikiTextReader::open(f.path()).unwrap(); let got = r.read_chars(50).unwrap(); assert!(got.len() >= 50, "got {} chars", got.len()); - assert!(got.len() < 100, "should stop near target, got {}", got.len()); + assert!( + got.len() < 100, + "should stop near target, got {}", + got.len() + ); } #[test] diff --git a/base-convert/crates/base-awq/tests/awq_pipeline.rs b/base-convert/crates/base-awq/tests/awq_pipeline.rs index 8ea5946..c26133a 100644 --- a/base-convert/crates/base-awq/tests/awq_pipeline.rs +++ b/base-convert/crates/base-awq/tests/awq_pipeline.rs @@ -61,9 +61,7 @@ fn awq_then_rtn_mse( // `weights` here is post-AWQ rotation; inverse_scales undo it. let original_unrotated: Vec = (0..out_features) .flat_map(|i| { - (0..in_features).map(move |j| { - weights[i * in_features + j] * inverse_scales[j] - }) + (0..in_features).map(move |j| weights[i * in_features + j] * inverse_scales[j]) }) .collect(); for i in 0..out_features { @@ -101,15 +99,7 @@ fn awq_plus_rtn_beats_plain_rtn_at_q2() { let rotated = awq_apply(&weights, n_in, &plan.scales); let plain = rtn_mse(&weights, n_in, n_out, 2, 32, false); - let awq = awq_then_rtn_mse( - &rotated, - n_in, - n_out, - &plan.inverse_scales, - 2, - 32, - false, - ); + let awq = awq_then_rtn_mse(&rotated, n_in, n_out, &plan.inverse_scales, 2, 32, false); assert!( awq < plain, @@ -137,15 +127,7 @@ fn awq_plus_rtn_at_q4_does_not_regress() { let rotated = awq_apply(&weights, n_in, &plan.scales); let plain = rtn_mse(&weights, n_in, n_out, 4, 64, false); - let awq = awq_then_rtn_mse( - &rotated, - n_in, - n_out, - &plan.inverse_scales, - 4, - 64, - false, - ); + let awq = awq_then_rtn_mse(&rotated, n_in, n_out, &plan.inverse_scales, 4, 64, false); // Non-regression invariant: AWQ search includes α=0, so the // optimum is by construction never worse than plain. @@ -176,15 +158,7 @@ fn awq_lite_is_a_safe_fallback() { let rotated = awq_apply(&weights, n_in, &plan.scales); let plain = rtn_mse(&weights, n_in, n_out, 4, 64, false); - let lite = awq_then_rtn_mse( - &rotated, - n_in, - n_out, - &plan.inverse_scales, - 4, - 64, - false, - ); + let lite = awq_then_rtn_mse(&rotated, n_in, n_out, &plan.inverse_scales, 4, 64, false); // Uniform absmax → AWQ-lite is identity; lite ≈ plain. assert!( diff --git a/base-convert/crates/base-convert/src/gpt_oss.rs b/base-convert/crates/base-convert/src/gpt_oss.rs new file mode 100644 index 0000000..d42fadf --- /dev/null +++ b/base-convert/crates/base-convert/src/gpt_oss.rs @@ -0,0 +1,672 @@ +//! gpt-oss conversion: HF `model_type = "gpt_oss"` MXFP4 checkpoints. +//! +//! This is a **mirror-policy transplant** path, the same class as the MLX +//! quantized-checkpoint path: every tensor the checkpoint stores quantized +//! (the MoE expert stacks, MXFP4 `*_blocks` + `*_scales`) is copied into +//! the bundle byte-for-byte — packed FP4 codes and E8M0 block scales +//! verbatim, never dequantized or requantized — and every tensor the +//! checkpoint keeps unquantized is carried at a dtype that represents it +//! losslessly (bf16 matrices stay bf16; 1-D / bias tensors are widened to +//! f32). The bundle therefore holds exactly the checkpoint's numbers, and +//! the header's `provenance` block records, per bundle tensor, which +//! checkpoint tensors it came from and how, so an independent gate can +//! check the claim. +//! +//! The HF generic path is not used because it funnels everything through +//! f32 + the target quantizer (which would requantize the experts) and its +//! safetensors reader has no route for the U8 block/scale tensors. +//! +//! Canonical bundle names (HF convention, `model.` stripped): +//! +//! | checkpoint | bundle | dtype | +//! |----------------------------------------------|---------------------------------------------|--------| +//! | `model.embed_tokens.weight` | `embed_tokens.weight` | bf16 | +//! | `lm_head.weight` | `lm_head.weight` | bf16 | +//! | `model.norm.weight` | `final_norm.weight` | f32 | +//! | `…input_layernorm.weight` | `layers.N.input_norm.weight` | f32 | +//! | `…post_attention_layernorm.weight` | `layers.N.post_attn_norm.weight` | f32 | +//! | `…self_attn.{q,k,v,o}_proj.weight` | same | bf16 | +//! | `…self_attn.{q,k,v,o}_proj.bias` | same | f32 | +//! | `…self_attn.sinks` | same | f32 | +//! | `…mlp.router.weight` | same | bf16 | +//! | `…mlp.router.bias` | same | f32 | +//! | `…mlp.experts.gate_up_proj_blocks/_scales` | `layers.N.mlp.experts.gate_up_proj.weight` | mxfp4 | +//! | `…mlp.experts.gate_up_proj_bias` | `layers.N.mlp.experts.gate_up_proj.bias` | f32 | +//! | `…mlp.experts.down_proj_blocks/_scales` | `layers.N.mlp.experts.down_proj.weight` | mxfp4 | +//! | `…mlp.experts.down_proj_bias` | `layers.N.mlp.experts.down_proj.bias` | f32 | +//! +//! `--target base-q2..base-q8` narrows the bundle without touching the +//! checkpoint's quantized numbers: the attention projections are +//! RTN-quantized to the requested scheme — quant from full precision, +//! never quant-from-quant — while the MXFP4 expert stacks are still +//! transplanted verbatim. Embeddings and lm_head stay bf16: gpt-oss +//! embedding rows carry per-group outliers past ±17, and measured on the +//! PPL anchor even base-q8 groups cost +8% there. The router, norms and +//! biases stay full precision too (the router once set a whole run's +//! accuracy floor). Group size is the canonical one for the bit width, +//! dropped to the largest of 128/64/32 that divides the tensor's +//! in-features; a matrix nothing divides is carried bf16 with a note. +//! +//! The fused expert `gate_up_proj` keeps the checkpoint's row order +//! (gate = even rows, up = odd rows); the runtime's gpt-oss expert kernel +//! reads it that way, so nothing is permuted. +//! +//! MXFP4 payload layout (per tensor, `[n_experts, out, in]`): the packed +//! FP4 nibbles (`in/2` bytes per row, low nibble first — exactly the HF +//! `_blocks` bytes, `[E, out, in/32, 16]` flattened) followed at +//! `scale_offset` by one E8M0 byte per 32-value group (exactly the HF +//! `_scales` bytes). `group_size = 32`, `scale_dtype = e8m0`, no biases. + +use anyhow::{bail, Context, Result}; +use base_format::{ + AlignmentConfig, ComputeRegion, Header, HeaderFlags, LayerDescriptor, LayerKind, + LayerPrecision, ModelConfig, QuantScheme, ResidencyHint, ScaleDtype, SourceInfo, TargetBackend, + TensorDtype, TensorEntry, TensorFlags, TokenizerBlob, +}; +use base_format::{BaseWriter, TensorPayload}; +use base_readers::hf::HfDir; +use base_readers::safetensors::StDtype; +use serde_json::json; + +use crate::QuantContext; + +/// True when `config.json` describes a gpt-oss checkpoint this path handles. +pub(crate) fn is_gpt_oss(config: &serde_json::Value) -> bool { + config.get("model_type").and_then(|v| v.as_str()) == Some("gpt_oss") +} + +fn quant_method(config: &serde_json::Value) -> Option { + config + .get("quantization_config") + .and_then(|q| q.get("quant_method")) + .and_then(|v| v.as_str()) + .map(|s| s.to_ascii_lowercase()) +} + +fn entry(name: &str, dtype: TensorDtype, shape: Vec, len: u64, hot: bool) -> TensorEntry { + TensorEntry { + name: name.to_string(), + dtype, + shape, + offset: 0, + length: len, + scale_offset: None, + scale_length: None, + bias_offset: None, + bias_length: None, + awq_scale_offset: None, + awq_scale_length: None, + group_size: None, + layout: None, + residency: Some(if hot { + ResidencyHint::Hot + } else { + ResidencyHint::Warm + }), + compute_region: if hot { + ComputeRegion::Gpu + } else { + ComputeRegion::Accelerator + }, + scale_dtype: None, + symmetric: false, + flags: TensorFlags::empty(), + checksum_xxh64: None, + source_ggml_type: None, + } +} + +/// bf16 source bytes → f32 bytes (lossless widening). +fn bf16_to_f32_bytes(bytes: &[u8]) -> Vec { + let mut out = Vec::with_capacity(bytes.len() * 2); + for c in bytes.chunks_exact(2) { + let bits = (u16::from_le_bytes([c[0], c[1]]) as u32) << 16; + out.extend_from_slice(&bits.to_le_bytes()); + } + out +} + +fn f16_to_f32_bytes(bytes: &[u8]) -> Vec { + let mut out = Vec::with_capacity(bytes.len() * 2); + for c in bytes.chunks_exact(2) { + let v = half::f16::from_le_bytes([c[0], c[1]]).to_f32(); + out.extend_from_slice(&v.to_le_bytes()); + } + out +} + +/// bf16/f16 source bytes → f32 values (for the RTN packer). +fn half_bytes_to_f32(bytes: &[u8], dtype: StDtype) -> Vec { + let mut out = Vec::with_capacity(bytes.len() / 2); + for c in bytes.chunks_exact(2) { + out.push(match dtype { + StDtype::Bf16 => f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16), + _ => half::f16::from_le_bytes([c[0], c[1]]).to_f32(), + }); + } + out +} + +/// The dense-quant request carried by `--target`: bit width + bundle +/// dtype + header scheme for base-q2..q8; None for the mirror targets. +/// Exact id of an added/special token from the checkpoint's `tokenizer.json`, +/// by literal content. Recovers the Harmony stop set for a checkpoint that +/// ships no `generation_config.json`; `None` when the file or the token is +/// absent, which the caller surfaces rather than silently accepting. +fn added_token_id(input: &std::path::Path, content: &str) -> Option { + let bytes = std::fs::read(input.join("tokenizer.json")).ok()?; + let tok: serde_json::Value = serde_json::from_slice(&bytes).ok()?; + for t in tok.get("added_tokens")?.as_array()? { + if t.get("content").and_then(|c| c.as_str()) == Some(content) { + return t.get("id").and_then(|i| i.as_u64()).map(|v| v as u32); + } + } + None +} + +fn dense_quant(target: crate::TargetScheme) -> Option<(u32, TensorDtype, QuantScheme)> { + use crate::TargetScheme as T; + match target { + T::BaseQ2 => Some((2, TensorDtype::BaseQ2, QuantScheme::BaseQ2)), + T::BaseQ3 => Some((3, TensorDtype::BaseQ3, QuantScheme::BaseQ3)), + T::BaseQ4 => Some((4, TensorDtype::BaseQ4, QuantScheme::BaseQ4)), + T::BaseQ5 => Some((5, TensorDtype::BaseQ5, QuantScheme::BaseQ5)), + T::BaseQ6 => Some((6, TensorDtype::BaseQ6, QuantScheme::BaseQ6)), + T::BaseQ8 => Some((8, TensorDtype::BaseQ8, QuantScheme::BaseQ8)), + _ => None, + } +} + +/// The canonical group size for the bit width, dropped to the largest of +/// 128/64/32 that divides `in_features` (groups must not straddle rows — +/// gpt-oss's hidden of 2880 rules out q8's canonical 128). +fn dense_group_size(bits: u32, in_features: u64) -> Option { + let canonical = base_quant::rtn::RtnConfig::canonical(bits).group_size; + [canonical, 64, 32] + .into_iter() + .find(|gs| *gs <= canonical && in_features % *gs as u64 == 0) +} + +pub(crate) fn convert_gpt_oss( + input: &std::path::Path, + output: &std::path::Path, + ctx: &QuantContext, +) -> Result<()> { + use base_arch::hf_mapper_for_model_type; + + let hf = HfDir::open(input)?; + let qm = quant_method(&hf.config); + match qm.as_deref() { + Some("mxfp4") => {} + other => bail!( + "gpt_oss: expected an MXFP4 checkpoint (quantization_config.quant_method = \"mxfp4\"), \ + found {:?} — only the MXFP4 transplant path is implemented", + other + ), + } + if ctx.profile.is_some() { + bail!("gpt_oss: --profile is not applicable — the checkpoint's MXFP4 experts are transplanted verbatim (mirror policy)"); + } + let dense = dense_quant(ctx.target); + match (ctx.target, &dense) { + (crate::TargetScheme::Mxfp4, _) | (crate::TargetScheme::Bf16, _) => {} + (_, Some((bits, _, _))) => eprintln!( + " note: --target {:?} quantizes the attention projections to {bits}-bit RTN; the MXFP4 \ + experts are transplanted verbatim and embed/lm_head/router stay full precision (outliers)", + ctx.target + ), + (other, None) => eprintln!( + " note: --target {:?} ignored for gpt_oss — the bundle mirrors the checkpoint (mxfp4 experts, bf16 dense)", + other + ), + } + + let mapper = hf_mapper_for_model_type("gpt_oss").expect("gpt_oss mapper registered"); + let mut config = mapper.config_from_hf(&hf.config)?; + // generation_config.json carries the full harmony stop set + // (`<|return|>`, `<|endoftext|>`, `<|call|>`); config.json only names + // the first. Merge the rest into `eos_token_ids` so the runtime stops + // on every end-of-turn marker. + let gc_path = input.join("generation_config.json"); + match std::fs::read(&gc_path) { + Ok(bytes) => { + // A file that exists but will not parse is a broken checkpoint, not + // an absent one: silently dropping it produces a bundle that looks + // fine and runs past `<|call|>`. + let gc: serde_json::Value = serde_json::from_slice(&bytes) + .with_context(|| format!("parsing {}", gc_path.display()))?; + let ids: Vec = match gc.get("eos_token_id") { + Some(serde_json::Value::Number(n)) => { + n.as_u64().map(|x| vec![x as u32]).unwrap_or_default() + } + Some(serde_json::Value::Array(a)) => a + .iter() + .filter_map(|v| v.as_u64().map(|x| x as u32)) + .collect(), + _ => Vec::new(), + }; + for id in ids { + if id != config.eos_token_id && !config.eos_token_ids.contains(&id) { + config.eos_token_ids.push(id); + } + } + } + Err(e) if e.kind() == std::io::ErrorKind::NotFound => { + // No generation_config.json: config.json names only `<|return|>`, + // so derive the rest of the Harmony stop set from the tokenizer. + // Without `<|call|>` the runtime keeps generating after a finished + // tool call until some other stop or the token limit. + // Every marker must resolve. A missing tokenizer.json, a missing + // added_tokens array, or an absent `<|call|>` would otherwise be + // accepted silently and the converter would report success while + // writing the same incomplete stop set this branch exists to + // repair — the bundle then runs past finished tool calls. + let mut missing: Vec<&str> = Vec::new(); + for marker in ["<|return|>", "<|call|>", "<|endoftext|>"] { + match added_token_id(input, marker) { + Some(id) => { + if id != config.eos_token_id && !config.eos_token_ids.contains(&id) { + config.eos_token_ids.push(id); + } + } + None => missing.push(marker), + } + } + if !missing.is_empty() { + bail!( + "gpt-oss checkpoint has no generation_config.json and its tokenizer.json does not \ + define the Harmony stop token(s) {missing:?}; the bundle would generate past a \ + finished tool call. Supply generation_config.json or a complete tokenizer.json." + ); + } + eprintln!( + " note: no generation_config.json — Harmony stops derived from the tokenizer ({:?})", + config.eos_token_ids + ); + } + Err(e) => return Err(e).with_context(|| format!("reading {}", gc_path.display())), + } + let n_layers = config.num_hidden_layers as usize; + let n_experts = config.num_experts as u64; + let hidden = config.hidden_size as u64; + let ffn = config.moe_intermediate_size as u64; + eprintln!( + " arch: gpt_oss (MXFP4 transplant) hidden={} layers={} heads={}/{} experts={} top_k={} ffn={} vocab={}", + hidden, + n_layers, + config.num_attention_heads, + config.num_kv_heads, + n_experts, + config.num_experts_per_tok, + ffn, + config.vocab_size + ); + + // ── header ─────────────────────────────────────────────────────── + let mut config_map = config.to_config_map(); + config_map.insert("model_type".into(), json!("gpt_oss")); + let header = Header { + schema: 1, + arch: "gpt_oss".to_string(), + // With a dense target the header names the requested scheme (the + // per-tensor dtypes stay authoritative — experts remain mxfp4). + quant_scheme: dense.as_ref().map_or(QuantScheme::Mxfp4, |(_, _, s)| *s), + min_hw: "apple_m1".to_string(), + created: crate::chrono_now(), + base_rt_version: env!("CARGO_PKG_VERSION").to_string(), + source: SourceInfo { + // A safetensors checkpoint directory that the MLX reference loads + // natively, converted under the mirror policy (quantized tensors + // transplanted verbatim, the rest carried losslessly) — the same + // contract as the MLX quantized-checkpoint path, hence its label. + format: "mlx_safetensors".to_string(), + sha256: "".to_string(), + filename: input + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("unknown") + .to_string(), + }, + tokenizer: TokenizerBlob { + fields: crate::tokenizer_from_hf(&hf), + }, + config: ModelConfig { fields: config_map }, + metadata: Default::default(), + target_backend: TargetBackend::Metal, + quant_profile: String::new(), + alignment: AlignmentConfig::default(), + flags: HeaderFlags::QUANTIZED | HeaderFlags::HAS_MOE, + layers: (0..n_layers) + .map(|_| LayerDescriptor { + kind: LayerKind::AttentionGqa, + moe_n_experts: config.num_experts as u16, + moe_n_active: config.num_experts_per_tok as u16, + shared_attn_layer: None, + compute_hint: Some(ComputeRegion::Accelerator), + precision: LayerPrecision::default(), + }) + .collect(), + tensors: vec![], + mmproj: None, + calibration: None, + sig: None, + provenance: None, // filled below via the writer's header + }; + + // ── tensor plan ────────────────────────────────────────────────── + // (bundle name, source tensor, kind) + #[derive(Clone, Copy)] + enum Kind { + /// bf16 matrix carried verbatim as bf16. + Matrix, + /// bf16/f16 small tensor widened to f32 (norms, biases, sinks). + Widen, + } + let mut plan: Vec<(String, String, Kind, bool)> = vec![ + ( + "embed_tokens.weight".into(), + "model.embed_tokens.weight".into(), + Kind::Matrix, + true, + ), + ( + "lm_head.weight".into(), + "lm_head.weight".into(), + Kind::Matrix, + false, + ), + ( + "final_norm.weight".into(), + "model.norm.weight".into(), + Kind::Widen, + true, + ), + ]; + // (bundle name, blocks source, scales source, expected [E, out, in]) + let mut mx_plan: Vec<(String, String, String, [u64; 3])> = Vec::new(); + for l in 0..n_layers { + let p = format!("model.layers.{l}."); + let b = format!("layers.{l}."); + plan.push(( + format!("{b}input_norm.weight"), + format!("{p}input_layernorm.weight"), + Kind::Widen, + true, + )); + plan.push(( + format!("{b}post_attn_norm.weight"), + format!("{p}post_attention_layernorm.weight"), + Kind::Widen, + true, + )); + for proj in ["q", "k", "v", "o"] { + plan.push(( + format!("{b}self_attn.{proj}_proj.weight"), + format!("{p}self_attn.{proj}_proj.weight"), + Kind::Matrix, + false, + )); + plan.push(( + format!("{b}self_attn.{proj}_proj.bias"), + format!("{p}self_attn.{proj}_proj.bias"), + Kind::Widen, + true, + )); + } + plan.push(( + format!("{b}self_attn.sinks"), + format!("{p}self_attn.sinks"), + Kind::Widen, + true, + )); + plan.push(( + format!("{b}mlp.router.weight"), + format!("{p}mlp.router.weight"), + Kind::Matrix, + true, + )); + plan.push(( + format!("{b}mlp.router.bias"), + format!("{p}mlp.router.bias"), + Kind::Widen, + true, + )); + plan.push(( + format!("{b}mlp.experts.gate_up_proj.bias"), + format!("{p}mlp.experts.gate_up_proj_bias"), + Kind::Widen, + true, + )); + plan.push(( + format!("{b}mlp.experts.down_proj.bias"), + format!("{p}mlp.experts.down_proj_bias"), + Kind::Widen, + true, + )); + mx_plan.push(( + format!("{b}mlp.experts.gate_up_proj.weight"), + format!("{p}mlp.experts.gate_up_proj_blocks"), + format!("{p}mlp.experts.gate_up_proj_scales"), + [n_experts, 2 * ffn, hidden], + )); + mx_plan.push(( + format!("{b}mlp.experts.down_proj.weight"), + format!("{p}mlp.experts.down_proj_blocks"), + format!("{p}mlp.experts.down_proj_scales"), + [n_experts, hidden, ffn], + )); + } + + // Coverage: every checkpoint tensor must be claimed by the plan. + let mut claimed = std::collections::BTreeSet::new(); + for (_, src, _, _) in &plan { + claimed.insert(src.clone()); + } + for (_, blocks, scales, _) in &mx_plan { + claimed.insert(blocks.clone()); + claimed.insert(scales.clone()); + } + let unclaimed: Vec = hf + .tensor_names() + .filter(|n| !claimed.contains(*n)) + .map(|s| s.to_string()) + .collect(); + if !unclaimed.is_empty() { + bail!( + "gpt_oss: {} checkpoint tensor(s) not understood by the converter (first: {:?}) — refusing to \ + produce a bundle that silently drops weights", + unclaimed.len(), + unclaimed.first() + ); + } + + let mut writer = BaseWriter::create(output, header).context("create writer")?; + let mut prov_tensors = serde_json::Map::new(); + + let pb = indicatif::ProgressBar::new((plan.len() + mx_plan.len()) as u64); + pb.set_style( + indicatif::ProgressStyle::with_template(" transplanting [{bar:28}] {pos}/{len} {msg}") + .expect("valid progress template") + .progress_chars("=>-"), + ); + + // ── carried tensors ────────────────────────────────────────────── + for (name, src, kind, hot) in &plan { + pb.set_message(name.clone()); + let info = hf + .tensor_info(src) + .with_context(|| format!("gpt_oss: checkpoint tensor {src} missing"))? + .clone(); + let bytes = hf + .tensor_bytes(src) + .expect("tensor_bytes after tensor_info"); + // The router stays full precision under every target — it is + // tiny, and quantizing it once set a whole run's accuracy floor. + // The embedding table and lm_head stay bf16 too: gpt-oss embedding + // rows carry per-group outliers past +-17, and measured on the + // anchor even base-q8 groups cost +8% PPL there while the + // attention projections quantize cleanly. + let quantize = dense + .as_ref() + .filter(|_| { + matches!(kind, Kind::Matrix) + && !name.ends_with("mlp.router.weight") + && name != "embed_tokens.weight" + && name != "lm_head.weight" + }) + .and_then(|(bits, dtype, _)| { + let in_f = *info.shape.last().unwrap_or(&0); + match dense_group_size(*bits, in_f) { + Some(gs) => Some((*bits, *dtype, gs)), + None => { + eprintln!( + " note: {name} in_features={in_f} fits no group size — carried bf16" + ); + None + } + } + }); + if let Some((bits, qdtype, gs)) = quantize { + let w = match info.dtype { + StDtype::Bf16 | StDtype::F16 => half_bytes_to_f32(bytes, info.dtype), + StDtype::F32 => bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect(), + other => bail!("gpt_oss: unexpected dtype {:?} for {src}", other), + }; + let cfg = base_quant::RtnConfig { + bits, + group_size: gs, + symmetric: false, + scale_dtype: ScaleDtype::Bf16, + }; + let packed = base_quant::rtn::pack(&w, cfg); + let mut data = Vec::with_capacity( + packed.packed_weights.len() + packed.scales.len() + packed.biases.len(), + ); + data.extend_from_slice(&packed.packed_weights); + let scale_off = data.len() as u64; + data.extend_from_slice(&packed.scales); + let bias_off = data.len() as u64; + data.extend_from_slice(&packed.biases); + let mut e = entry(name, qdtype, info.shape.clone(), data.len() as u64, *hot); + e.scale_offset = Some(scale_off); + e.scale_length = Some(packed.scales.len() as u64); + if !packed.biases.is_empty() { + e.bias_offset = Some(bias_off); + e.bias_length = Some(packed.biases.len() as u64); + } + e.group_size = Some(gs); + e.scale_dtype = Some(ScaleDtype::Bf16); + writer.add_tensor(TensorPayload { entry: e, data }); + prov_tensors.insert( + name.clone(), + json!({ "src": [src], "quantized_to": format!("base_q{bits}"), + "transform": "rtn", "group_size": gs }), + ); + pb.inc(1); + continue; + } + let (data, dtype): (Vec, TensorDtype) = match (kind, info.dtype) { + (Kind::Matrix, StDtype::Bf16) => (bytes.to_vec(), TensorDtype::Bf16), + (Kind::Matrix, StDtype::F16) => (bytes.to_vec(), TensorDtype::F16), + (Kind::Matrix, StDtype::F32) => (bytes.to_vec(), TensorDtype::F32), + (Kind::Widen, StDtype::Bf16) => (bf16_to_f32_bytes(bytes), TensorDtype::F32), + (Kind::Widen, StDtype::F16) => (f16_to_f32_bytes(bytes), TensorDtype::F32), + (Kind::Widen, StDtype::F32) => (bytes.to_vec(), TensorDtype::F32), + (_, other) => bail!("gpt_oss: unexpected dtype {:?} for {src}", other), + }; + let e = entry(name, dtype, info.shape.clone(), data.len() as u64, *hot); + writer.add_tensor(TensorPayload { entry: e, data }); + prov_tensors.insert(name.clone(), json!({ "src": [src] })); + pb.inc(1); + } + + // ── MXFP4 expert stacks (verbatim transplant) ──────────────────── + for (name, blocks_src, scales_src, dims) in &mx_plan { + pb.set_message(name.clone()); + let binfo = hf + .tensor_info(blocks_src) + .with_context(|| format!("gpt_oss: checkpoint tensor {blocks_src} missing"))? + .clone(); + let sinfo = hf + .tensor_info(scales_src) + .with_context(|| format!("gpt_oss: checkpoint tensor {scales_src} missing"))? + .clone(); + if binfo.dtype != StDtype::U8 || sinfo.dtype != StDtype::U8 { + bail!("gpt_oss: {blocks_src}/{scales_src} must be U8 (MXFP4 blocks + E8M0 scales)"); + } + let [e, n, k] = *dims; + let groups = k / 32; + let want_blocks = vec![e, n, groups, 16]; + let want_scales = vec![e, n, groups]; + if binfo.shape != want_blocks { + bail!( + "gpt_oss: {blocks_src} shape {:?}, expected {:?}", + binfo.shape, + want_blocks + ); + } + if sinfo.shape != want_scales { + bail!( + "gpt_oss: {scales_src} shape {:?}, expected {:?}", + sinfo.shape, + want_scales + ); + } + let blocks = hf.tensor_bytes(blocks_src).expect("blocks bytes"); + let scales = hf.tensor_bytes(scales_src).expect("scales bytes"); + let mut data = Vec::with_capacity(blocks.len() + scales.len()); + data.extend_from_slice(blocks); + let scale_off = data.len() as u64; + data.extend_from_slice(scales); + let mut te = entry( + name, + TensorDtype::Mxfp4, + vec![e, n, k], + data.len() as u64, + false, + ); + te.scale_offset = Some(scale_off); + te.scale_length = Some(scales.len() as u64); + te.group_size = Some(32); + te.scale_dtype = Some(ScaleDtype::E8m0); + te.symmetric = true; + writer.add_tensor(TensorPayload { entry: te, data }); + // The tensor is the checkpoint's expert stack copied verbatim: FP4 + // codes from `_blocks`, E8M0 block exponents from `_scales`. Both + // sources are listed; the stack note says how many experts it holds. + prov_tensors.insert( + name.clone(), + json!({ + "src": [blocks_src, scales_src], + "transplanted_mxfp4": true, + "stack": { "pattern": blocks_src, "count": e }, + "scales": scales_src, + }), + ); + pb.inc(1); + } + pb.finish_and_clear(); + + writer.set_provenance(json!({ + "tensors": prov_tensors, + "dropped": [], + "policy": match &dense { + Some((bits, _, _)) => format!( + "mirror + dense target: checkpoint-quantized tensors transplanted verbatim (mxfp4 codes + e8m0 scales); \ + attention projections RTN-quantized to base_q{bits} (per-tensor `quantized_to` entries); \ + embed/lm_head/router/norms/biases carried full precision"), + None => "mirror: checkpoint-quantized tensors transplanted verbatim (mxfp4 codes + e8m0 scales), \ + unquantized tensors carried losslessly (bf16 matrices as bf16, 1-D/bias tensors widened to f32)".to_string(), + }, + })); + eprintln!( + " wrote {} carried + {} mxfp4 expert tensors", + plan.len(), + mx_plan.len() + ); + writer.finish().context("finish")?; + Ok(()) +} diff --git a/base-convert/crates/base-convert/src/hub.rs b/base-convert/crates/base-convert/src/hub.rs index c14f829..0037e9d 100644 --- a/base-convert/crates/base-convert/src/hub.rs +++ b/base-convert/crates/base-convert/src/hub.rs @@ -10,6 +10,7 @@ use crate::{AwqMode, ConvertArgs, ListArgs, PullArgs, TargetScheme}; use anyhow::{bail, Context, Result}; use base_hub::cache::{self, HubSidecar}; use base_hub::fetch::{self, Fetcher, HfFetcher}; +use base_hub::parts::{self, Artifact, Grouped}; use base_hub::registry::{MergedRegistry, ModelEntry, ModelRef, Registry, SourceKind}; use std::path::{Path, PathBuf}; use std::process::Command; @@ -61,7 +62,10 @@ fn want_source_file(name: &str) -> bool { /// `default-q4` → `q4`, `base_q8` → `q8`, `Llama-3.2-1B-Instruct-Q4` → `q4`, /// `bf16` → `bf16`. fn quant_tag(s: &str) -> String { - s.rsplit(['-', '_']).next().unwrap_or(s).to_ascii_lowercase() + s.rsplit(['-', '_']) + .next() + .unwrap_or(s) + .to_ascii_lowercase() } /// True when `tag` names a quant scheme we know how to label on disk. @@ -89,25 +93,54 @@ fn base_file_stem(f: &str) -> String { b.strip_suffix(".base").unwrap_or(b).to_string() } -/// List the `.base` artifacts a repo hosts (empty when it ships none). -fn list_base_files(fetcher: &dyn Fetcher, repo: &str, revision: &str) -> Result> { +/// List the `.base` artifacts a repo hosts (empty when it ships none). A +/// bundle the Hub's 50 GB file cap forced into `.base.part-NNN` pieces shows +/// up once, under its logical `.base` name. +fn list_base_files(fetcher: &dyn Fetcher, repo: &str, revision: &str) -> Result { let files = fetcher .list_files(repo, revision) .with_context(|| format!("listing files in {repo}@{revision}"))?; - Ok(files.into_iter().filter(|f| f.ends_with(".base")).collect()) + let grouped = parts::group(files); + for (name, why) in &grouped.malformed { + eprintln!(" skipping {name}: {why}"); + } + Ok(grouped) } -/// Pick the `.base` whose quant tag matches `want`. Falls back to the sole -/// artifact when the repo has exactly one; errors when several are present and +/// Pick the `.base` whose quant tag matches `want`. A malformed part set +/// carrying that quant is an error saying why, never a fall-through to a +/// different quant. Falls back to the sole artifact when the repo has +/// exactly one and nothing malformed; errors when several are present and /// none match. -fn select_base_file<'a>(files: &'a [String], want: &str) -> Result<&'a String> { - if let Some(f) = files.iter().find(|f| quant_tag(&base_file_stem(f)) == want) { +fn select_base_file<'a>(files: &'a Grouped, want: &str) -> Result<&'a Artifact> { + if let Some(f) = files + .artifacts + .iter() + .find(|f| quant_tag(&base_file_stem(&f.name)) == want) + { return Ok(f); } - if let [only] = files { + if let Some((name, why)) = files + .malformed + .iter() + .find(|(n, _)| quant_tag(&base_file_stem(n)) == want) + { + bail!("{name}: the {want} publication is incomplete ({why})"); + } + if let ([only], []) = (files.artifacts.as_slice(), files.malformed.as_slice()) { return Ok(only); } - let avail: Vec = files.iter().map(|f| quant_tag(&base_file_stem(f))).collect(); + let avail: Vec = files + .artifacts + .iter() + .map(|f| quant_tag(&base_file_stem(&f.name))) + .chain( + files + .malformed + .iter() + .map(|(n, _)| format!("{} (incomplete)", quant_tag(&base_file_stem(n)))), + ) + .collect(); bail!( "no pre-converted .base for quant {want:?} in this repo; it offers: {}", avail.join(", ") @@ -169,11 +202,9 @@ fn installed_single_path(reg: &MergedRegistry, id: &str) -> Result = installed.iter().filter(|r| r.id == id).collect(); match hits.as_slice() { [] => Ok(None), - [one] => Ok(Some( - one.path - .clone() - .with_context(|| format!("installed model `{id}` has no artifact path"))?, - )), + [one] => Ok(Some(one.path.clone().with_context(|| { + format!("installed model `{id}` has no artifact path") + })?)), many => { // Multiple variants installed — e.g. a universal `default-q4` cached // before the backend-qualified catalog entries existed, plus a native @@ -185,25 +216,33 @@ fn installed_single_path(reg: &MergedRegistry, id: &str) -> Result = - many.iter().copied().filter(|r| r.variant.starts_with(&native_prefix)).collect(); + let native: Vec<&ModelEntry> = many + .iter() + .copied() + .filter(|r| r.variant.starts_with(&native_prefix)) + .collect(); let pick = if native.len() == 1 { Some(native[0]) } else if native.is_empty() { - let uni: Vec<&ModelEntry> = - many.iter().copied().filter(|r| r.variant.starts_with("default-")).collect(); + let uni: Vec<&ModelEntry> = many + .iter() + .copied() + .filter(|r| r.variant.starts_with("default-")) + .collect(); (uni.len() == 1).then(|| uni[0]) } else { None }; if let Some(one) = pick { - return Ok(Some( - one.path - .clone() - .with_context(|| format!("installed model `{id}` has no artifact path"))?, - )); + return Ok(Some(one.path.clone().with_context(|| { + format!("installed model `{id}` has no artifact path") + })?)); } - let variants = many.iter().map(|r| r.variant.as_str()).collect::>().join(", "); + let variants = many + .iter() + .map(|r| r.variant.as_str()) + .collect::>() + .join(", "); bail!("model `{id}` has multiple installed variants ({variants}) — specify one as `{id}:`") } } @@ -232,7 +271,7 @@ fn installed_best_variant(reg: &MergedRegistry, id: &str, want: &str) -> Option< let runnable = |v: &str| -> bool { match v.split_once('-') { Some((slot, _)) if KNOWN_BACKENDS.contains(&slot) => slot == backend, // native only - _ => true, // default-*, bare bits, etc. + _ => true, // default-*, bare bits, etc. } }; let installed = reg.local.list().ok()?; @@ -241,7 +280,8 @@ fn installed_best_variant(reg: &MergedRegistry, id: &str, want: &str) -> Option< .filter(|r| { r.id == id && runnable(&r.variant) - && base_hub::registry::quant_bits(&r.variant).unwrap_or(r.variant.as_str()) == want_bits + && base_hub::registry::quant_bits(&r.variant).unwrap_or(r.variant.as_str()) + == want_bits }) .collect(); if matches.is_empty() { @@ -255,7 +295,11 @@ fn installed_best_variant(reg: &MergedRegistry, id: &str, want: &str) -> Option< /// Fetch a not-yet-installed model on demand, then return its artifact path. /// Prefers the pre-converted basecompute mirror; otherwise converts the source /// repo on pull. Progress (download + quantization) is shown by `cmd_pull`. -fn auto_pull_and_resolve(reg: &MergedRegistry, id: &str, want_variant: Option<&str>) -> Result { +fn auto_pull_and_resolve( + reg: &MergedRegistry, + id: &str, + want_variant: Option<&str>, +) -> Result { let pull_id = preconverted_id(reg, id); let target = want_variant .map(|v| target_from_quant(&quant_tag(v))) @@ -299,7 +343,9 @@ fn auto_pull_and_resolve(reg: &MergedRegistry, id: &str, want_variant: Option<&s /// on demand — preferring the pre-converted basecompute mirror, else converting /// the source repo — so `basert chat`/`serve ` Just Works. fn resolve_hub_model(token: &str, default_variant: Option<&str>) -> Result { - let (id, inline) = token.split_once(':').map_or((token, None), |(i, v)| (i, Some(v))); + let (id, inline) = token + .split_once(':') + .map_or((token, None), |(i, v)| (i, Some(v))); // A trailing/empty `:` is a typo, not "default variant" — fail loudly so it // doesn't silently resolve to q4. @@ -337,7 +383,9 @@ fn resolve_model_args(rest: &[String], default_variant: Option<&str>) -> Result< rest.iter() .map(|arg| { if looks_like_hub_id(arg) { - Ok(resolve_hub_model(arg, default_variant)?.to_string_lossy().into_owned()) + Ok(resolve_hub_model(arg, default_variant)? + .to_string_lossy() + .into_owned()) } else { Ok(arg.clone()) } @@ -380,6 +428,36 @@ fn extract_variant_flag(rest: &[String]) -> Result<(Option, Vec) Ok((variant, out)) } +const COMPUTEARENA_HARNESS_ENV: &str = "COMPUTEARENA_BASERT_HARNESS"; +const LEGACY_COMPUTEARENA_HARNESS_ENV: &str = "BASERT_COMPUTEARENA_HARNESS"; +const BENCHMARK_HARNESS_BINARY: &str = "basert-benchmark-harness"; + +fn bundled_computearena_harness( + exe_dir: Option<&Path>, + environment_override: bool, +) -> Option { + if environment_override { + return None; + } + let exe_dir = exe_dir?; + + // Published BaseRT packages keep the launcher and harness together. + let sibling = exe_dir.join(BENCHMARK_HARNESS_BINARY); + if sibling.is_file() { + return Some(sibling); + } + + // A source build places the Rust launcher under + // tools/base-convert/target/{debug,release}, while CMake publishes the + // native harness under build/. Recognize that layout so the documented + // developer build works exactly like an installed release. + exe_dir.ancestors().find_map(|directory| { + let workspace_manifest = directory.join("tools/base-convert/Cargo.toml"); + let harness = directory.join("build").join(BENCHMARK_HARNESS_BINARY); + (workspace_manifest.is_file() && harness.is_file()).then_some(harness) + }) +} + /// Forward `basert [args…]` to the matching runtime binary. Searches for /// `basert-` (release layout) then `baseRT_` (local dev build), /// looking next to this executable first and then on `PATH`. On success the @@ -388,18 +466,35 @@ pub fn dispatch_external(argv: Vec) -> Result<()> { use std::os::unix::process::CommandExt; let (cmd, rest) = argv.split_first().context("no command given")?; - // `--variant ` is a launcher-level model selector; strip it before - // forwarding (the runtime binary doesn't know it) and apply it during - // hub-id resolution. - let (variant_flag, rest) = extract_variant_flag(rest)?; - let rest = resolve_model_args(&rest, variant_flag.as_deref())?; - let candidates = [format!("basert-{cmd}"), format!("baseRT_{cmd}")]; + let is_computearena = cmd == "computearena"; + let (candidates, rest) = if is_computearena { + // ComputeArena is independently distributed and owns runtime + // selection. Preserve its arguments and select the BaseRT adapter. + let mut forwarded = Vec::with_capacity(rest.len() + 1); + forwarded.push("basert".to_string()); + forwarded.extend_from_slice(rest); + (vec!["computearena".to_string()], forwarded) + } else { + // `--variant ` is a launcher-level model selector; strip it before + // forwarding (the runtime binary doesn't know it) and apply it during + // hub-id resolution. + let (variant_flag, rest) = extract_variant_flag(rest)?; + let rest = resolve_model_args(&rest, variant_flag.as_deref())?; + (vec![format!("basert-{cmd}"), format!("baseRT_{cmd}")], rest) + }; // Prefer a binary sitting next to `basert` (how the release ships); fall // back to a bare name, which `Command` resolves against `PATH`. let exe_dir = std::env::current_exe() .ok() .and_then(|p| p.parent().map(Path::to_path_buf)); + let harness_environment_override = std::env::var_os(COMPUTEARENA_HARNESS_ENV).is_some() + || std::env::var_os(LEGACY_COMPUTEARENA_HARNESS_ENV).is_some(); + let bundled_harness = if is_computearena { + bundled_computearena_harness(exe_dir.as_deref(), harness_environment_override) + } else { + None + }; let mut targets: Vec = Vec::new(); if let Some(dir) = &exe_dir { for name in &candidates { @@ -413,11 +508,23 @@ pub fn dispatch_external(argv: Vec) -> Result<()> { for target in &targets { // exec() returns only on failure; ENOENT means try the next candidate. - let err = Command::new(target).args(&rest).exec(); + let mut command = Command::new(target); + command.args(&rest); + if let Some(harness) = &bundled_harness { + command.env(COMPUTEARENA_HARNESS_ENV, harness); + } + let err = command.exec(); if err.kind() != std::io::ErrorKind::NotFound { return Err(err).with_context(|| format!("launching {}", target.display())); } } + if is_computearena { + bail!( + "ComputeArena is not installed.\n\ + Install it from https://computearena.ai/quickstart, then run \ + `basert computearena` again." + ) + } if RUNTIME_COMMANDS.contains(&cmd.as_str()) { bail!( "`basert {cmd}` needs the BaseRT runtime, which wasn't found in this \ @@ -431,8 +538,14 @@ pub fn dispatch_external(argv: Vec) -> Result<()> { /// Commands served by the BaseRT runtime rather than this binary. Used only /// to shape the not-found error above; dispatch itself is name-driven, so /// commands absent from this list (or from `basert --help`) still dispatch. -const RUNTIME_COMMANDS: [&str; 6] = - ["serve", "chat", "complete", "bench", "transcribe", "profile"]; +const RUNTIME_COMMANDS: [&str; 6] = [ + "serve", + "chat", + "complete", + "bench", + "transcribe", + "profile", +]; pub fn cmd_pull(args: PullArgs) -> Result<()> { let reg = MergedRegistry::load()?; @@ -457,7 +570,13 @@ pub fn cmd_pull(args: PullArgs) -> Result<()> { // Serve it directly when it's what the user wants; otherwise grab the // requested quant straight from the same repo rather than silently // handing back the cataloged one. - ModelRef::Catalog { id, hf_repo, revision, variant, .. } => { + ModelRef::Catalog { + id, + hf_repo, + revision, + variant, + .. + } => { // Match on the QUANT BITS, not the raw variant string: the resolver // may hand back a backend-native variant (e.g. `cuda-q4mix`) that // satisfies a `q4` request but whose `quant_tag` ("q4mix") isn't the @@ -480,7 +599,10 @@ pub fn cmd_pull(args: PullArgs) -> Result<()> { ModelRef::HuggingFace { id, repo, revision } => { let fetcher = HfFetcher::new(cache::hf_staging_dir(&root))?; let base_files = list_base_files(&fetcher, repo, revision)?; - if base_files.is_empty() { + // A repo holding only a half-uploaded split bundle is still a + // `.base` repo: it gets the incomplete-publication error, not a + // conversion attempt. + if base_files.artifacts.is_empty() && base_files.malformed.is_empty() { pull_and_convert(&root, &args, id, repo, revision) } else { pull_base_direct(&root, &args, id, repo, revision, &fetcher, &base_files) @@ -492,9 +614,19 @@ pub fn cmd_pull(args: PullArgs) -> Result<()> { fn print_plan(r: &ModelRef, want: &str) { match r { ModelRef::Local { id, variant, path } => { - println!("plan: {id} [{variant}] already installed at {}", path.display()) + println!( + "plan: {id} [{variant}] already installed at {}", + path.display() + ) } - ModelRef::Catalog { id, hf_repo, file, revision, variant, .. } => { + ModelRef::Catalog { + id, + hf_repo, + file, + revision, + variant, + .. + } => { let want_bits = base_hub::registry::quant_bits(want).unwrap_or(want); if base_hub::registry::quant_bits(variant).unwrap_or(variant) == want_bits { println!( @@ -523,6 +655,7 @@ fn pull_catalog(root: &Path, r: &ModelRef) -> Result<()> { revision, variant, sha256, + parts_sha256, .. } = r else { @@ -532,14 +665,23 @@ fn pull_catalog(root: &Path, r: &ModelRef) -> Result<()> { eprintln!(" catalog: {hf_repo}/{file}@{revision} (pre-converted)"); let fetcher = HfFetcher::new(cache::hf_staging_dir(root))?; - let src = fetcher.get_file(hf_repo, revision, file)?; + // The catalog names the logical `.base`; the repo may hold it as parts. + let artifact = parts::find(&fetcher, hf_repo, revision, file)?; let vdir = cache::variant_dir(root, id, variant)?; std::fs::create_dir_all(&vdir)?; let out = cache::base_artifact_path(&vdir); + announce_parts(&artifact, &out); // Moves the staged download into place (same filesystem), so the pulled // artifact exists exactly once on disk. - fetch::install_file(&fetcher, hf_repo, &src, &out)?; + parts::install( + &fetcher, + hf_repo, + revision, + &artifact, + &out, + parts_sha256.as_deref(), + )?; let got_sha = crate::compute_sha256_streaming(&out)?; if let Some(expected) = sha256 { @@ -582,18 +724,23 @@ fn pull_base_direct( repo: &str, revision: &str, fetcher: &dyn Fetcher, - base_files: &[String], + base_files: &Grouped, ) -> Result<()> { eprintln!("basert pull v{}", env!("CARGO_PKG_VERSION")); eprintln!(" source: {repo}@{revision} (HuggingFace, pre-converted .base)"); let want = quant_token(args); - let file = select_base_file(base_files, &want)?; + let artifact = select_base_file(base_files, &want)?; + let file = &artifact.name; // Label the on-disk variant by the artifact's own quant when it carries // one (so a `-Q8.base` never lands in a `default-q4` dir); otherwise fall // back to what was requested. let file_tag = quant_tag(&base_file_stem(file)); - let variant_tag = if is_quant_tag(&file_tag) { file_tag } else { want }; + let variant_tag = if is_quant_tag(&file_tag) { + file_tag + } else { + want + }; let variant = format!("default-{variant_tag}"); eprintln!(" variant: {variant}"); eprintln!(" file: {file}"); @@ -601,19 +748,49 @@ fn pull_base_direct( let vdir = cache::variant_dir(root, id, &variant)?; std::fs::create_dir_all(&vdir)?; let out = cache::base_artifact_path(&vdir); + announce_parts(artifact, &out); - let src = fetcher.get_file(repo, revision, file)?; // Moves the staged download into place (same filesystem), so the pulled - // artifact exists exactly once on disk. - fetch::install_file(fetcher, repo, &src, &out)?; + // artifact exists exactly once on disk. A part set is reassembled beside + // `out` and renamed into place once complete. + parts::install(fetcher, repo, revision, artifact, &out, None)?; let sha = crate::compute_sha256_streaming(&out).ok(); - write_sidecar_for(&vdir, id, "huggingface", repo, None, revision, &variant, None, sha)?; + write_sidecar_for( + &vdir, + id, + "huggingface", + repo, + None, + revision, + &variant, + None, + sha, + )?; fetch::cleanup_staging(fetcher, repo); eprintln!("installed {id} [{variant}] → {}", out.display()); Ok(()) } +/// One line about a split bundle, so the user knows why the install step +/// runs on after the download bar finishes and what the disk needs. +fn announce_parts(artifact: &Artifact, out: &Path) { + if !artifact.is_split() { + return; + } + // A re-pull keeps the installed bundle until the new one is complete, + // so the honest figure for that path is one bundle more. + let replacing = if out.exists() { + " plus the installed bundle it replaces, kept until the new one is complete" + } else { + "" + }; + eprintln!( + " parts: {} (Hub-split; reassembled on install, needs bundle + one part of free disk{replacing})", + artifact.parts.len() + ); +} + /// Convert-on-pull: download an HF repo's source safetensors and run the /// existing conversion pipeline into the cache. fn pull_and_convert( @@ -648,6 +825,7 @@ fn pull_and_convert( profile: profile_path, awq_profile: None, allow_quant_from_quant: false, + no_mlx_passthrough: false, // Convert-on-pull always produces a canonical-quant bundle; // k-quant passthrough stays an explicit `convert` opt-in. kquant_passthrough: false, @@ -656,6 +834,10 @@ fn pull_and_convert( // apply to GGUF sources, which ship the tower separately. mmproj: None, mmproj_config: None, + mlx_passthrough: false, + validate: false, + direct_write: false, + imatrix: false, }; crate::cmd_convert(conv).with_context(|| format!("converting {repo}"))?; @@ -733,18 +915,19 @@ fn choose_profile( /// are downloaded. `text_config.model_type` is consulted as a fallback for /// multimodal configs that nest the language-model arch there. fn check_supported_arch(config_path: &Path, repo: &str, revision: &str) -> Result<()> { - let bytes = std::fs::read(config_path) - .with_context(|| format!("reading {}", config_path.display()))?; + let bytes = + std::fs::read(config_path).with_context(|| format!("reading {}", config_path.display()))?; let cfg: serde_json::Value = serde_json::from_slice(&bytes) .with_context(|| format!("parsing {} as JSON", config_path.display()))?; let model_type = cfg .get("model_type") .and_then(|v| v.as_str()) - .or_else(|| cfg.pointer("/text_config/model_type").and_then(|v| v.as_str())) - .ok_or_else(|| { - anyhow::anyhow!("{repo}@{revision}: config.json has no model_type field") - })?; + .or_else(|| { + cfg.pointer("/text_config/model_type") + .and_then(|v| v.as_str()) + }) + .ok_or_else(|| anyhow::anyhow!("{repo}@{revision}: config.json has no model_type field"))?; if base_arch::hf_mapper_for_model_type(model_type).is_some() { return Ok(()); @@ -961,11 +1144,103 @@ pub fn cmd_catalog_scan(org: String, out: Option, dry_run: bool) -> Res Ok(()) } +/// Write the manifest a split bundle needs beside its parts. +/// +/// `bundle` is the logical `.base` path the parts are named after +/// (`parts/GLM-5.2-Q4.base`, which need not exist); its `.part-NNN` siblings +/// are hashed in order, each and as one whole, and the result lands at +/// `.manifest.json` ready to upload with the parts. +pub fn cmd_catalog_manifest(bundle: PathBuf) -> Result<()> { + let dir = bundle + .parent() + .filter(|p| !p.as_os_str().is_empty()) + .map(Path::to_path_buf) + .unwrap_or_else(|| PathBuf::from(".")); + let stem = bundle + .file_name() + .and_then(|n| n.to_str()) + .context("bundle path has no file name")? + .to_string(); + // Every entry is read or the command fails: an entry lost to an I/O + // error could be the terminal part, and a manifest written without it + // would describe a truncated bundle as complete. + let mut parts: Vec<(u32, PathBuf)> = Vec::new(); + for entry in std::fs::read_dir(&dir).with_context(|| format!("listing {}", dir.display()))? { + let entry = entry.with_context(|| format!("reading an entry of {}", dir.display()))?; + let Some(name) = entry.file_name().to_str().map(str::to_string) else { + continue; + }; + if let Some((logical, idx)) = base_hub::parts::split_part_name(&name) { + if logical == stem { + parts.push((idx, entry.path())); + } + } + } + parts.sort(); + if parts.is_empty() { + bail!("no {stem}.part-NNN files in {}", dir.display()); + } + for (want, (have, path)) in (0u32..).zip(&parts) { + if *have != want { + bail!( + "part set is missing part {want:03} (next is {})", + path.display() + ); + } + } + let paths: Vec = parts.into_iter().map(|(_, p)| p).collect(); + eprintln!("hashing {} parts of {stem} …", paths.len()); + let manifest = base_hub::parts::build_manifest(&paths)?; + let out = dir.join(base_hub::parts::manifest_name(&stem)); + std::fs::write(&out, serde_json::to_string_pretty(&manifest)? + "\n") + .with_context(|| format!("writing {}", out.display()))?; + eprintln!( + " size: {} bytes\n sha256: {}\nwrote {}", + manifest.size, + manifest.sha256, + out.display() + ); + Ok(()) +} + #[cfg(test)] mod tests { use super::*; use base_hub::fetch::MockFetcher; + #[test] + fn bundled_computearena_harness_requires_a_sibling_and_respects_overrides() { + let tmp = tempfile::tempdir().unwrap(); + assert_eq!(bundled_computearena_harness(Some(tmp.path()), false), None); + + let harness = tmp.path().join(BENCHMARK_HARNESS_BINARY); + std::fs::write(&harness, b"fixture").unwrap(); + assert_eq!( + bundled_computearena_harness(Some(tmp.path()), false), + Some(harness) + ); + assert_eq!(bundled_computearena_harness(Some(tmp.path()), true), None); + assert_eq!(bundled_computearena_harness(None, false), None); + } + + #[test] + fn bundled_computearena_harness_supports_the_source_build_layout() { + let tmp = tempfile::tempdir().unwrap(); + let repo = tmp.path().join("baseRT"); + let exe_dir = repo.join("tools/base-convert/target/release"); + std::fs::create_dir_all(&exe_dir).unwrap(); + std::fs::write(repo.join("tools/base-convert/Cargo.toml"), b"[workspace]\n").unwrap(); + + let harness = repo.join("build").join(BENCHMARK_HARNESS_BINARY); + std::fs::create_dir_all(harness.parent().unwrap()).unwrap(); + std::fs::write(&harness, b"fixture").unwrap(); + + assert_eq!( + bundled_computearena_harness(Some(&exe_dir), false), + Some(harness) + ); + } + #[test] fn quant_tag_extracts_last_segment() { assert_eq!(quant_tag("default-q4"), "q4"); @@ -985,17 +1260,48 @@ mod tests { #[test] fn select_base_file_matches_quant_then_falls_back() { let files = vec![ - "Llama-3.2-1B-Instruct-Q4.base".to_string(), - "Llama-3.2-1B-Instruct-Q8.base".to_string(), + Artifact::whole("Llama-3.2-1B-Instruct-Q4.base"), + Artifact::whole("Llama-3.2-1B-Instruct-Q8.base"), ]; - assert_eq!(select_base_file(&files, "q4").unwrap(), &files[0]); - assert_eq!(select_base_file(&files, "q8").unwrap(), &files[1]); + let files = Grouped { + artifacts: files, + malformed: vec![], + }; + assert_eq!(select_base_file(&files, "q4").unwrap(), &files.artifacts[0]); + assert_eq!(select_base_file(&files, "q8").unwrap(), &files.artifacts[1]); // No match among several → error that lists what's available. let err = select_base_file(&files, "q2").unwrap_err().to_string(); assert!(err.contains("q4") && err.contains("q8"), "{err}"); // Sole artifact → used regardless of requested quant. - let one = vec!["model.base".to_string()]; - assert_eq!(select_base_file(&one, "q4").unwrap(), &one[0]); + let one = Grouped { + artifacts: vec![Artifact::whole("model.base")], + malformed: vec![], + }; + assert_eq!(select_base_file(&one, "q4").unwrap(), &one.artifacts[0]); + } + + #[test] + fn select_base_file_reports_an_incomplete_requested_quant() { + // A gapped Q4 upload beside a complete Q8: asking for Q4 is told + // the Q4 publication is incomplete, not quietly handed Q8. + let files = Grouped { + artifacts: vec![Artifact::whole("m-Q8.base")], + malformed: vec![( + "m-Q4.base".to_string(), + "part set is missing part 001 (found 2 parts)".to_string(), + )], + }; + let err = select_base_file(&files, "q4").unwrap_err().to_string(); + assert!(err.contains("incomplete"), "{err}"); + assert!(err.contains("missing part 001"), "{err}"); + // Q8 is still there for whoever asks for it. + assert_eq!(select_base_file(&files, "q8").unwrap().name, "m-Q8.base"); + // And a third quant sees both, the broken one marked. + let err = select_base_file(&files, "q2").unwrap_err().to_string(); + assert!( + err.contains("q4 (incomplete)") && err.contains("q8"), + "{err}" + ); } #[test] @@ -1007,11 +1313,84 @@ mod tests { std::fs::write(repo_dir.join(f), b"x").unwrap(); } let fetcher = MockFetcher::new(tmp.path()); - let mut got = list_base_files(&fetcher, "basecompute/m", "main").unwrap(); + let mut got: Vec = list_base_files(&fetcher, "basecompute/m", "main") + .unwrap() + .artifacts + .into_iter() + .map(|a| a.name) + .collect(); got.sort(); assert_eq!(got, vec!["m-Q4.base".to_string(), "m-Q8.base".to_string()]); } + #[test] + fn list_base_files_groups_a_part_set_under_its_logical_name() { + let tmp = tempfile::tempdir().unwrap(); + let repo_dir = tmp.path().join("basecompute").join("m"); + std::fs::create_dir_all(&repo_dir).unwrap(); + for f in ["m-Q4.base.part-000", "m-Q4.base.part-001", "README.md"] { + std::fs::write(repo_dir.join(f), b"x").unwrap(); + } + let fetcher = MockFetcher::new(tmp.path()); + let got = list_base_files(&fetcher, "basecompute/m", "main").unwrap(); + assert_eq!(got.artifacts.len(), 1); + assert_eq!(got.artifacts[0].name, "m-Q4.base"); + assert_eq!(got.artifacts[0].parts.len(), 2); + // The logical name carries the quant tag the selector matches on. + assert_eq!(select_base_file(&got, "q4").unwrap().name, "m-Q4.base"); + } + + #[test] + fn pull_base_direct_reassembles_a_split_bundle() { + let tmp = tempfile::tempdir().unwrap(); + let repo_dir = tmp.path().join("basecompute").join("m"); + std::fs::create_dir_all(&repo_dir).unwrap(); + // A real tiny bundle, cut in two: the install checks the stitched + // file against its own header, so arbitrary bytes will not do. + let whole = base_hub::parts::synthetic_bundle(1000); + let (head, tail) = whole.split_at(whole.len() / 2); + std::fs::write(repo_dir.join("m-Q4.base.part-000"), head).unwrap(); + std::fs::write(repo_dir.join("m-Q4.base.part-001"), tail).unwrap(); + // The manifest the publisher ships beside the parts, made the way + // they would make it. + cmd_catalog_manifest(repo_dir.join("m-Q4.base")).unwrap(); + assert!(repo_dir.join("m-Q4.base.manifest.json").exists()); + let fetcher = MockFetcher::new(tmp.path()); + let base_files = list_base_files(&fetcher, "basecompute/m", "main").unwrap(); + + let root = tmp.path().join("cache"); + let args = PullArgs { + id: "basecompute/m".into(), + profile: None, + target: TargetScheme::BaseQ4, + revision: "main".into(), + force: false, + dry_run: false, + }; + pull_base_direct( + &root, + &args, + "basecompute/m", + "basecompute/m", + "main", + &fetcher, + &base_files, + ) + .unwrap(); + + let vdir = root.join("basecompute/m/default-q4"); + assert_eq!(std::fs::read(vdir.join("model.base")).unwrap(), whole); + assert!(vdir.join("hub.json").exists()); + // No reassembly leftovers beside the installed artifact. + let stray: Vec<_> = std::fs::read_dir(&vdir) + .unwrap() + .flatten() + .map(|e| e.file_name().to_string_lossy().into_owned()) + .filter(|n| n.contains("partial")) + .collect(); + assert!(stray.is_empty(), "leftovers: {stray:?}"); + } + #[test] fn pull_base_direct_installs_requested_quant() { let tmp = tempfile::tempdir().unwrap(); @@ -1032,8 +1411,16 @@ mod tests { force: false, dry_run: false, }; - pull_base_direct(&root, &args, "basecompute/m", "basecompute/m", "main", &fetcher, &base_files) - .unwrap(); + pull_base_direct( + &root, + &args, + "basecompute/m", + "basecompute/m", + "main", + &fetcher, + &base_files, + ) + .unwrap(); // The Q8 artifact landed under the default-q8 variant dir. let out = root.join("basecompute/m/default-q8/model.base"); @@ -1053,7 +1440,8 @@ mod tests { impl StagedFetcher { fn repo_dir(&self, repo: &str) -> PathBuf { - self.staging.join(format!("models--{}", repo.replace('/', "--"))) + self.staging + .join(format!("models--{}", repo.replace('/', "--"))) } fn stage(&self, repo: &str, revision: &str, filename: &str, bytes: &[u8]) { @@ -1070,7 +1458,11 @@ mod tests { impl Fetcher for StagedFetcher { fn get_file(&self, repo: &str, revision: &str, filename: &str) -> anyhow::Result { - let p = self.repo_dir(repo).join("snapshots").join(revision).join(filename); + let p = self + .repo_dir(repo) + .join("snapshots") + .join(revision) + .join(filename); anyhow::ensure!(p.exists(), "not staged: {}", p.display()); Ok(p) } @@ -1107,8 +1499,16 @@ mod tests { force: false, dry_run: false, }; - pull_base_direct(&root, &args, "basecompute/m", "basecompute/m", "main", &fetcher, &base_files) - .unwrap(); + pull_base_direct( + &root, + &args, + "basecompute/m", + "basecompute/m", + "main", + &fetcher, + &base_files, + ) + .unwrap(); let out = root.join("basecompute/m/default-q4/model.base"); assert_eq!(std::fs::read(&out).unwrap(), b"q4-bytes"); diff --git a/base-convert/crates/base-convert/src/main.rs b/base-convert/crates/base-convert/src/main.rs index 03dbd30..41a128c 100644 --- a/base-convert/crates/base-convert/src/main.rs +++ b/base-convert/crates/base-convert/src/main.rs @@ -2,6 +2,7 @@ use anyhow::{bail, Context, Result}; use clap::{Parser, Subcommand, ValueEnum}; use std::path::PathBuf; +mod gpt_oss; mod hub; /// Hand-written top-level help: the runtime commands are dispatched via @@ -17,6 +18,7 @@ Run models: chat Chat with a model interactively complete Generate a one-shot completion bench Measure throughput + computearena Run and manage ComputeArena benchmarks Manage models: pull Download a model from the BaseRT catalog or Hugging Face @@ -67,6 +69,8 @@ enum Cmd { List(ListArgs), /// Regenerate the model catalog by scanning a published HF organization. CatalogScan(CatalogScanArgs), + /// Write the manifest a Hub-split `.base` needs beside its parts. + CatalogManifest(CatalogManifestArgs), /// Runtime commands — `serve`, `chat`, `complete`, `bench`, … — handled /// by the BaseRT runtime (dispatched in `hub::dispatch_external`). #[command(external_subcommand)] @@ -132,6 +136,18 @@ struct ConvertArgs { #[arg(long)] allow_quant_from_quant: bool, + /// Requantize MLX affine-q4 sources through f32 instead of + /// transplanting their packed bytes into `base_q4`. + /// + /// The two schemes are identical (INT4 asymmetric, group 64, f16 + /// scale + bias), so `--target base-q4` from an MLX 4-bit source + /// normally copies the codes verbatim and reproduces the source + /// weights bit-for-bit. Requantizing instead re-derives each group's + /// scale from already-quantized values, which lands on a different + /// grid — use this only to reproduce a bundle built before the + /// transplant path existed. + #[arg(long)] + no_mlx_passthrough: bool, /// GGUF sources only: copy Q4_K / Q5_K / Q6_K super-blocks into the /// bundle VERBATIM instead of dequantizing and re-packing them. /// @@ -180,6 +196,44 @@ struct ConvertArgs { /// Overrides the built-in per-projector-type defaults. #[arg(long, value_name = "PATH", requires = "mmproj")] mmproj_config: Option, + + /// Reuse an MLX checkpoint's already-quantized payloads verbatim + /// (packed nibbles + bf16 scales/biases) instead of the lossy + /// dequant→requant round trip. The written weights are + /// bit-identical to the MLX source. Requires an MLX 4-bit + /// group-size-64 source, `--target base-q4`, and no `--profile`. + /// Tensors the runtime needs at f16 (MLA k_b/v_b, the MoE router) + /// still dequantize exactly and store as f16. + #[arg(long)] + mlx_passthrough: bool, + + /// After a `--mlx-passthrough` conversion, re-open the written + /// `.base` and byte-compare every tensor against the MLX source + /// (packed/scales/biases verbatim for passthrough tensors; the + /// recomputed f16 bytes for dequantized ones). Exact equality or + /// error — this is the "oracle" gate for downstream DSA + /// validation. + #[arg(long)] + validate: bool, + + /// Stream the blob straight into the output file behind a reserved + /// header region instead of via a `.blobtmp` sibling — peak disk + /// usage becomes the bundle size instead of 2×. The header is + /// space-padded to the 64 MiB reserve (negligible on the huge + /// bundles this exists for; don't use it on small models). + #[arg(long)] + direct_write: bool, + + /// Importance-weighted RTN ("imatrix"): use the `--awq-profile` + /// sidecar's per-input-channel activation absmax as weights for the + /// per-group affine fit, instead of AWQ's weight rotation. Unlike + /// AWQ, the packed tensors still approximate the ORIGINAL weights — + /// no runtime activation scaling exists or is needed; only where + /// the quantization error lands changes (away from salient + /// channels). Requires --awq-profile; mutually exclusive with the + /// AWQ rotation path (this flag takes precedence). + #[arg(long)] + imatrix: bool, } #[derive(Parser, Debug)] @@ -226,7 +280,7 @@ struct InspectArgs { verify_checksums: bool, } -#[derive(Copy, Clone, Debug, ValueEnum)] +#[derive(Copy, Clone, Debug, PartialEq, Eq, ValueEnum)] enum TargetScheme { BaseQ2, BaseQ3, @@ -236,6 +290,9 @@ enum TargetScheme { BaseQ8, Bf16, Mxfp4, + /// NVIDIA fp4: e2m1 codes + e4m3 per-block-16 scales. Accepts the + /// `base-nvfp4` spelling used by conversion harnesses. + #[value(alias = "base-nvfp4")] Nvfp4, } @@ -281,6 +338,13 @@ struct CatalogScanArgs { dry_run: bool, } +#[derive(Parser, Debug)] +struct CatalogManifestArgs { + /// The logical `.base` path the parts are named after, e.g. + /// `parts/GLM-5.2-Q4.base` (its `.part-NNN` siblings are described). + bundle: PathBuf, +} + #[derive(Parser, Debug)] struct ListArgs { /// Also list catalog models that aren't installed yet. @@ -342,7 +406,12 @@ fn main() -> Result<()> { // subcommands (convert/pull/list/...) stay quiet, and the runtime tools // print their own banner (tools/basert_banner.h). let raw_arg1 = std::env::args().nth(1); - if raw_arg1.is_none() || matches!(raw_arg1.as_deref(), Some("-h") | Some("--help") | Some("help")) { + if raw_arg1.is_none() + || matches!( + raw_arg1.as_deref(), + Some("-h") | Some("--help") | Some("help") + ) + { print_banner(); } let args = Args::parse(); @@ -355,6 +424,7 @@ fn main() -> Result<()> { Cmd::Pull(a) => hub::cmd_pull(a), Cmd::List(a) => hub::cmd_list(a), Cmd::CatalogScan(a) => hub::cmd_catalog_scan(a.org, a.out, a.dry_run), + Cmd::CatalogManifest(a) => hub::cmd_catalog_manifest(a.bundle), Cmd::External(argv) => hub::dispatch_external(argv), } } @@ -362,8 +432,7 @@ fn main() -> Result<()> { fn cmd_keygen(args: KeygenArgs) -> Result<()> { use ed25519_dalek::SigningKey; use rand_core::OsRng; - std::fs::create_dir_all(&args.out) - .with_context(|| format!("creating {:?}", args.out))?; + std::fs::create_dir_all(&args.out).with_context(|| format!("creating {:?}", args.out))?; let mut rng = OsRng; let sk = SigningKey::generate(&mut rng); let vk = sk.verifying_key(); @@ -400,6 +469,21 @@ fn cmd_convert(args: ConvertArgs) -> Result<()> { return convert_synthetic_with_ctx(&output, &ctx); } + // gpt-oss (OpenAI, MXFP4 MoE): a safetensors checkpoint whose expert + // stacks are pre-quantized in a layout the generic paths cannot read. + // Routed before format detection (its config.json has + // `quantization_config`, not the MLX `quantization` block, so detect_format + // would send it down the dequantize-and-requantize HF path). + if args.input.is_dir() { + if let Ok(bytes) = std::fs::read(args.input.join("config.json")) { + if let Ok(cfg) = serde_json::from_slice::(&bytes) { + if gpt_oss::is_gpt_oss(&cfg) { + return gpt_oss::convert_gpt_oss(&args.input, &output, &ctx); + } + } + } + } + // Detect source format: GGUF / HF-safetensors / MLX-safetensors. use base_readers::SourceFormat; let fmt = base_readers::detect_format(&args.input) @@ -418,6 +502,52 @@ fn cmd_convert(args: ConvertArgs) -> Result<()> { } } +/// Build the typed per-layer descriptors for the header. Homogeneous +/// transformers (empty `layer_types`) keep the historical all-GQA +/// layout; hybrid configs get their real schedule so the header states +/// which layers are SSM vs attention vs MoE-FFN and how many experts +/// they carry. +fn layer_descriptors_from_config( + config: &base_arch::ArchConfig, +) -> Vec { + use base_format::{ComputeRegion, LayerDescriptor, LayerKind, LayerPrecision}; + (0..config.num_hidden_layers as usize) + .map(|i| { + let kind = match config.layer_types.get(i).map(String::as_str) { + // Nemotron-H vocabulary. + Some("mamba") => LayerKind::Ssm, + Some("moe") => LayerKind::MoeFfn, + Some("mlp") => LayerKind::DenseMlp, + // Qwen3.5 hybrid vocabulary: Gated-DeltaNet layers are + // recurrent-state layers, closest to Ssm. + Some("linear_attention") => LayerKind::Ssm, + // "attention" / "full_attention" / unknown / homogeneous. + _ => LayerKind::AttentionGqa, + }; + let is_moe_layer = kind == LayerKind::MoeFfn; + LayerDescriptor { + kind, + moe_n_experts: if is_moe_layer { + config.num_experts as u16 + } else { + 0 + }, + moe_n_active: if is_moe_layer { + config.num_experts_per_tok as u16 + } else { + 0 + }, + shared_attn_layer: None, + compute_hint: Some(ComputeRegion::Accelerator), + precision: LayerPrecision { + force_fp32_ssm: kind == LayerKind::Ssm, + ..LayerPrecision::default() + }, + } + }) + .collect() +} + /// Real-model conversion: read a GGUF, dequant per-tensor to f32, /// remap tensor names to canonical .base convention, re-quantize to the /// target scheme, write the .base file. @@ -432,15 +562,14 @@ fn convert_gguf( ) -> Result<()> { use base_arch::source_mapper_for_gguf; use base_format::{ - AlignmentConfig, BaseReader, BaseWriter, ComputeRegion, Header, HeaderFlags, LayerKind, - LayerDescriptor, LayerPrecision, ModelConfig, QuantScheme, SourceInfo, TargetBackend, TensorDtype, - TensorFlags, TensorPayload, TokenizerBlob, + AlignmentConfig, BaseReader, BaseWriter, ComputeRegion, Header, HeaderFlags, ModelConfig, + QuantScheme, SourceInfo, TargetBackend, TensorDtype, TensorFlags, TensorPayload, + TokenizerBlob, }; use base_readers::gguf::{dequant_to_f32, ggml_type_name, GgmlType, GgufFile}; let target = ctx.target; - let gguf = GgufFile::open(input) - .with_context(|| format!("opening GGUF {:?}", input))?; + let gguf = GgufFile::open(input).with_context(|| format!("opening GGUF {:?}", input))?; let arch = gguf .arch() .ok_or_else(|| anyhow::anyhow!("GGUF missing general.architecture"))?; @@ -552,21 +681,23 @@ fn convert_gguf( metadata: Default::default(), target_backend: TargetBackend::Metal, quant_profile: ctx.profile_name().unwrap_or("").to_string(), - alignment: AlignmentConfig::default(), + // Accel tensors align to the 16 KiB Apple page (default is 64 B): + // the runtime's BaseWeightStore can then always split its chunked + // no-copy mmap at a tensor start. Mixed-dtype bundles with 64 B + // alignment can run a whole max_buffer_size window without a + // page-aligned start (seen on the GLM 5.2 q4/q5/q6 production + // bundle), forcing overlap-mapped splits or per-tensor copies. + // Padding cost: < tensor_count × 16 KiB — noise on any real model. + alignment: AlignmentConfig { + accel_align_log2: 14, + ..Default::default() + }, flags: HeaderFlags::QUANTIZED, - layers: (0..config.num_hidden_layers) - .map(|_| LayerDescriptor { - kind: LayerKind::AttentionGqa, - moe_n_experts: 0, - moe_n_active: 0, - shared_attn_layer: None, - compute_hint: Some(ComputeRegion::Accelerator), - precision: LayerPrecision::default(), - }) - .collect(), + layers: layer_descriptors_from_config(&config), tensors: vec![], mmproj: None, calibration: None, + provenance: None, sig: None, }; @@ -624,9 +755,8 @@ fn convert_gguf( // permutation only ever moves whole rows). let unpermuted_bytes = match mapper.rope_unpermute_heads(&canonical, &config) { Some(n_heads) => { - let out = unpermute_rope_rows(info, raw, n_heads).with_context(|| { - format!("rope row un-permute for {:?}", info.name) - })?; + let out = unpermute_rope_rows(info, raw, n_heads) + .with_context(|| format!("rope row un-permute for {:?}", info.name))?; unpermuted += 1; Some(out) } @@ -668,7 +798,14 @@ fn convert_gguf( || info.name.ends_with(".ssm_a"); let is_ssm_sensitive = is_ssm_a || canonical.ends_with(".ssm.dt_bias") - || canonical.ends_with(".ssm.d"); + || canonical.ends_with(".ssm.d") + // Grouped RMS-norm gains ([groups, d] — 2-D, so the 1-D + // norm check misses them) and the short depthwise conv: + // tiny tensors on the recurrent path; group-quantizing + // them wrecks the state update. Keep f32 on CPU. + || canonical.ends_with(".ssm.norm.weight") + || canonical.ends_with(".ssm.conv1d.weight") + || canonical.ends_with(".ssm.conv1d.bias"); // SSM A-matrix (and adjacent SSM scalars) MUST stay f32 in CPU // region — quantizing them produces NaN after ~100 recurrent // steps. Regular 1-D norm weights and the embed/lm_head pair @@ -678,6 +815,20 @@ fn convert_gguf( // memory bloat). let is_norm_like = info.shape.len() == 1; let is_embedding = canonical == "embed_tokens.weight" || canonical == "lm_head.weight"; + // Emit at f16 (not the target quant) for two precision-sensitive + // GLM 5.2 cases: + // * MLA k_b/v_b up-projections — consumed by dedicated absorb + // kernels that read the weights as RAW f16 (not the quant GEMM + // path) and encode the exact GGUF ne0-fastest 3-D layout; + // quantizing would corrupt those raw-half reads outright. + // * The MoE router (ffn_gate_inp → mlp.router) — llama.cpp keeps + // it F32; base_qN routing of this tiny [dim, n_experts] tensor + // can flip borderline sigmoid+bias top-k selections and diverge + // the whole expert mixture from the reference. + let is_mla_absorb = + canonical.ends_with(".k_b_proj.weight") || canonical.ends_with(".v_b_proj.weight"); + let is_moe_router = canonical.ends_with(".router.weight"); + let force_f16 = is_mla_absorb || is_moe_router; let (entry, data) = if is_ssm_sensitive { let mut flags = TensorFlags::empty(); @@ -704,12 +855,12 @@ fn convert_gguf( symmetric: false, flags, checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; let data: Vec = f32s.iter().flat_map(|f| f.to_le_bytes()).collect(); entry.length = data.len() as u64; (entry, data) - } else if is_norm_like { + } else if is_norm_like || force_f16 { // 1-D norms (and biases caught by the same shape check) are // always emitted at f16. A profile's catch-all `**.weight` // rule typically targets a quant bit-width; quantizing a @@ -742,8 +893,8 @@ fn convert_gguf( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; (entry, bytes) } else if is_embedding { // Embed/lm_head: pack at the target quant scheme. Earlier @@ -752,11 +903,11 @@ fn convert_gguf( // (Llama-3.2-1B at MLX-direct ships embed at 4-bit ≈ 131 MB — // cache-friendly). Quantizing embed matches the source format // and is what the embedding_lookup_q4 kernel expects. + let in_features = info.shape.last().copied().map(|d| d as usize); let (packed, dtype) = if ctx.profile.is_some() { - let in_features = info.shape.last().copied().map(|d| d as usize); ctx.pack_tensor(&canonical, &f32s, in_features)? } else { - pack_for_target(&f32s, target)? + pack_for_target_rows(&f32s, target, in_features, &canonical)? }; let mut data = Vec::with_capacity( packed.packed_weights.len() + packed.scales.len() + packed.biases.len(), @@ -806,17 +957,17 @@ fn convert_gguf( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; entry.length = data.len() as u64; (entry, data) } else { // Quantize to target scheme, Accelerator region. + let in_features = info.shape.last().copied().map(|d| d as usize); let (packed, dtype) = if ctx.profile.is_some() { - let in_features = info.shape.last().copied().map(|d| d as usize); ctx.pack_tensor(&canonical, &f32s, in_features)? } else { - pack_for_target(&f32s, target)? + pack_for_target_rows(&f32s, target, in_features, &canonical)? }; let mut data = Vec::with_capacity( packed.packed_weights.len() + packed.scales.len() + packed.biases.len(), @@ -867,8 +1018,8 @@ fn convert_gguf( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; let _ = entry.length; // length will be overwritten by writer entry.length = data.len() as u64; (entry, data) @@ -916,8 +1067,8 @@ fn convert_gguf( let hf_cfg = match mmproj_config { Some(p) => { - let bytes = std::fs::read(p) - .with_context(|| format!("reading --mmproj-config {:?}", p))?; + let bytes = + std::fs::read(p).with_context(|| format!("reading --mmproj-config {:?}", p))?; let v: serde_json::Value = serde_json::from_slice(&bytes) .with_context(|| format!("parsing --mmproj-config {:?}", p))?; // Tower geometry still comes from the GGUF; this only @@ -937,8 +1088,8 @@ fn convert_gguf( for info in mm.tensors.iter() { // Unknown names are fatal, not skipped: a silently dropped tower // weight produces plausible-looking output, not an error. - let canonical = base_arch::muse_glimmer::map_mmproj_gguf_name(&info.name) - .ok_or_else(|| { + let canonical = + base_arch::muse_glimmer::map_mmproj_gguf_name(&info.name).ok_or_else(|| { anyhow::anyhow!( "unmapped mmproj tensor {:?} — refusing to drop a tower weight", info.name @@ -1070,11 +1221,9 @@ fn convert_gguf( offset: 0, length: 0, scale_offset: (!packed.scales.is_empty()).then_some(scale_off), - scale_length: (!packed.scales.is_empty()) - .then_some(packed.scales.len() as u64), + scale_length: (!packed.scales.is_empty()).then_some(packed.scales.len() as u64), bias_offset: (!packed.biases.is_empty()).then_some(bias_off), - bias_length: (!packed.biases.is_empty()) - .then_some(packed.biases.len() as u64), + bias_length: (!packed.biases.is_empty()).then_some(packed.biases.len() as u64), awq_scale_offset: None, awq_scale_length: None, group_size: (packed.group_size > 0).then_some(packed.group_size), @@ -1190,7 +1339,11 @@ fn unpermute_rope_rows( } if n_rows % n_heads != 0 { - bail!("row count {} not divisible by head count {}", n_rows, n_heads); + bail!( + "row count {} not divisible by head count {}", + n_rows, + n_heads + ); } let hd = n_rows / n_heads; if hd % 2 != 0 { @@ -1275,11 +1428,7 @@ fn kquant_passthrough_entry( } /// Convert from an HF safetensors directory. -fn convert_hf( - input: &std::path::Path, - output: &std::path::Path, - ctx: &QuantContext, -) -> Result<()> { +fn convert_hf(input: &std::path::Path, output: &std::path::Path, ctx: &QuantContext) -> Result<()> { use base_arch::hf_mapper_for_model_type; use base_readers::hf::HfDir; let hf = HfDir::open(input)?; @@ -1350,7 +1499,10 @@ fn convert_hf( } } if ids.len() > config.eos_token_ids.len() + 1 { - eprintln!(" eos: stop ids {:?} (merged generation_config.json)", ids); + eprintln!( + " eos: stop ids {:?} (merged generation_config.json)", + ids + ); } config.eos_token_id = ids[0]; config.eos_token_ids = ids[1..].to_vec(); @@ -1361,12 +1513,28 @@ fn convert_hf( let provider = HfTensorProvider { hf: &hf }; let mmproj_cfg = mmproj_config_from_hf(&hf); let config_for_permute = config.clone(); + // A modelopt NVFP4 checkpoint (hf_quant_config.json, quant_algo NVFP4) + // is a *quantized* source: its own quantization decisions are the + // bundle's (mirror policy). Quantized tensors are transplanted + // verbatim; tensors it keeps unquantized are carried losslessly. + let nvfp4_source = std::fs::read_to_string(input.join("hf_quant_config.json")) + .ok() + .and_then(|t| serde_json::from_str::(&t).ok()) + .map(|v| v["quantization"]["quant_algo"].as_str() == Some("NVFP4")) + .unwrap_or(false); + if nvfp4_source { + eprintln!(" source: NVFP4-quantized checkpoint (mirror policy: transplant + carry)"); + } convert_generic( input, output, ctx, mapper.canonical_arch(), - "hf_safetensors", + if nvfp4_source { + "nvfp4_safetensors" + } else { + "hf_safetensors" + }, config, &provider, hf.tensor_names().map(|s| s.to_string()).collect(), @@ -1374,6 +1542,12 @@ fn convert_hf( mmproj_cfg, &|n| mapper.norm_shift(n), &|n| mapper.rope_permute_heads(n, &config_for_permute), + &|n| mapper.value_transform(n), + mapper.shape_fastest_first(), + // An unquantized HF source has nothing to mirror — every tensor + // is the target scheme's business. A quantized NVFP4 source + // mirrors. + nvfp4_source, &|n| mapper.row_rms_normalize(n, &config_for_permute), ) } @@ -1441,11 +1615,11 @@ fn convert_whisper( // Special-token metadata — hard requirement. A whisper bundle // without exact token ids mis-transcribes silently (the historic // vocab-size-table bug), so no tokenizer.json = no conversion. - let tokenizer_json = hf - .tokenizer_json - .as_ref() - .context("whisper: model dir has no tokenizer.json (required for whisper.* token metadata)")?; - let mut metadata = whisper::token_metadata_from_tokenizer(tokenizer_json, config.vocab_size as i64)?; + let tokenizer_json = hf.tokenizer_json.as_ref().context( + "whisper: model dir has no tokenizer.json (required for whisper.* token metadata)", + )?; + let mut metadata = + whisper::token_metadata_from_tokenizer(tokenizer_json, config.vocab_size as i64)?; // large-v3-turbo distillation dropped the translation task; record it so // the runtime rejects task=translate instead of silently emitting // source-language text (see whisper::supports_translate for the config @@ -1492,11 +1666,14 @@ fn convert_whisper( } } } - eprintln!(" mapped: {} tensors kept, {} dropped", mapped.len(), dropped); + eprintln!( + " mapped: {} tensors kept, {} dropped", + mapped.len(), + dropped + ); // Sanity guard: every contract-required tensor must be present. - let have: std::collections::BTreeSet<&str> = - mapped.iter().map(|(_, c)| c.as_str()).collect(); + let have: std::collections::BTreeSet<&str> = mapped.iter().map(|(_, c)| c.as_str()).collect(); let missing: Vec = whisper::required_tensor_names(&config) .into_iter() .filter(|t| !have.contains(t.as_str())) @@ -1601,6 +1778,7 @@ fn convert_whisper( mmproj: None, calibration: None, sig: None, + provenance: None, }; let mut writer = BaseWriter::create(output, header).context("create writer")?; @@ -1638,102 +1816,100 @@ fn convert_whisper( (ComputeRegion::Accelerator, ResidencyHint::Warm) }; - let (entry, data) = if quantizing - && shape.len() == 2 - && whisper::is_quantizable_linear(canonical) - { - // Profile-routed block linear: pack at the canonical scheme, - // [W | scales | biases] in one blob — exactly the layout the - // LLM quant tensors use. A profile may still route these to - // f16 (whisper-f16.json); pack_tensor then returns raw f16 - // bytes with empty scale/bias streams and the entry - // degenerates to the plain case. - let in_features = shape.last().copied().map(|d| d as usize); - let (packed, dtype) = ctx - .pack_tensor(canonical, &f32s, in_features) - .with_context(|| format!("packing {canonical}"))?; - let mut data = Vec::with_capacity( - packed.packed_weights.len() + packed.scales.len() + packed.biases.len(), - ); - data.extend_from_slice(&packed.packed_weights); - let scale_off = data.len() as u64; - data.extend_from_slice(&packed.scales); - let bias_off = data.len() as u64; - data.extend_from_slice(&packed.biases); - let entry = base_format::TensorEntry { - name: canonical.clone(), - dtype, - shape, - offset: 0, - length: data.len() as u64, - scale_offset: if !packed.scales.is_empty() { - Some(scale_off) - } else { - None - }, - scale_length: if !packed.scales.is_empty() { - Some(packed.scales.len() as u64) - } else { - None - }, - bias_offset: if !packed.biases.is_empty() { - Some(bias_off) - } else { - None - }, - bias_length: if !packed.biases.is_empty() { - Some(packed.biases.len() as u64) - } else { - None - }, - awq_scale_offset: None, - awq_scale_length: None, - group_size: if packed.group_size > 0 { - Some(packed.group_size) - } else { - None - }, - layout: None, - residency: Some(residency), - compute_region: region, - scale_dtype: packed.scale_dtype, - symmetric: false, - flags: TensorFlags::empty(), - checksum_xxh64: None, - source_ggml_type: None, - }; - (entry, data) - } else { - // Everything else — and every tensor on the default path — - // is raw f16. - let bytes: Vec = f32s - .iter() - .flat_map(|&f| half::f16::from_f32(f).to_le_bytes()) - .collect(); - let entry = base_format::TensorEntry { - name: canonical.clone(), - dtype: TensorDtype::F16, - shape, - offset: 0, - length: bytes.len() as u64, - scale_offset: None, - scale_length: None, - bias_offset: None, - bias_length: None, - awq_scale_offset: None, - awq_scale_length: None, - group_size: None, - layout: None, - residency: Some(residency), - compute_region: region, - scale_dtype: None, - symmetric: false, - flags: TensorFlags::empty(), - checksum_xxh64: None, - source_ggml_type: None, + let (entry, data) = + if quantizing && shape.len() == 2 && whisper::is_quantizable_linear(canonical) { + // Profile-routed block linear: pack at the canonical scheme, + // [W | scales | biases] in one blob — exactly the layout the + // LLM quant tensors use. A profile may still route these to + // f16 (whisper-f16.json); pack_tensor then returns raw f16 + // bytes with empty scale/bias streams and the entry + // degenerates to the plain case. + let in_features = shape.last().copied().map(|d| d as usize); + let (packed, dtype) = ctx + .pack_tensor(canonical, &f32s, in_features) + .with_context(|| format!("packing {canonical}"))?; + let mut data = Vec::with_capacity( + packed.packed_weights.len() + packed.scales.len() + packed.biases.len(), + ); + data.extend_from_slice(&packed.packed_weights); + let scale_off = data.len() as u64; + data.extend_from_slice(&packed.scales); + let bias_off = data.len() as u64; + data.extend_from_slice(&packed.biases); + let entry = base_format::TensorEntry { + name: canonical.clone(), + dtype, + shape, + offset: 0, + length: data.len() as u64, + scale_offset: if !packed.scales.is_empty() { + Some(scale_off) + } else { + None + }, + scale_length: if !packed.scales.is_empty() { + Some(packed.scales.len() as u64) + } else { + None + }, + bias_offset: if !packed.biases.is_empty() { + Some(bias_off) + } else { + None + }, + bias_length: if !packed.biases.is_empty() { + Some(packed.biases.len() as u64) + } else { + None + }, + awq_scale_offset: None, + awq_scale_length: None, + group_size: if packed.group_size > 0 { + Some(packed.group_size) + } else { + None + }, + layout: None, + residency: Some(residency), + compute_region: region, + scale_dtype: packed.scale_dtype, + symmetric: false, + flags: TensorFlags::empty(), + checksum_xxh64: None, + source_ggml_type: None, + }; + (entry, data) + } else { + // Everything else — and every tensor on the default path — + // is raw f16. + let bytes: Vec = f32s + .iter() + .flat_map(|&f| half::f16::from_f32(f).to_le_bytes()) + .collect(); + let entry = base_format::TensorEntry { + name: canonical.clone(), + dtype: TensorDtype::F16, + shape, + offset: 0, + length: bytes.len() as u64, + scale_offset: None, + scale_length: None, + bias_offset: None, + bias_length: None, + awq_scale_offset: None, + awq_scale_length: None, + group_size: None, + layout: None, + residency: Some(residency), + compute_region: region, + scale_dtype: None, + symmetric: false, + flags: TensorFlags::empty(), + checksum_xxh64: None, + source_ggml_type: None, + }; + (entry, bytes) }; - (entry, bytes) - }; writer.add_tensor(TensorPayload { entry, data }); pb.inc(1); } @@ -1786,7 +1962,8 @@ fn convert_mlx( CANONICAL_QUANT_SPEC.md, canonical-quant requires fp16/bf16/fp32 source). \ Re-fetch the fp16/bf16 HF checkpoint, or pass --allow-quant-from-quant to \ accept the compounded quant error.", - mlx.quant.bits, mlx.quant.group_size + mlx.quant.bits, + mlx.quant.group_size ); } let mapper = hf_mapper_for_model_type(model_type) @@ -1801,6 +1978,17 @@ fn convert_mlx( config.intermediate_size, config.vocab_size ); + if ctx.mlx_passthrough { + if mlx.quant.bits != 4 || mlx.quant.group_size != 64 { + bail!( + "--mlx-passthrough requires a 4-bit group-size-64 MLX source (this one is \ + {}-bit gs={})", + mlx.quant.bits, + mlx.quant.group_size + ); + } + eprintln!(" passthrough: reusing MLX q4 payloads verbatim (bit-identical weights)"); + } let names: Vec = mlx .hf .tensor_names() @@ -1818,16 +2006,170 @@ fn convert_mlx( "mlx_safetensors", config, &provider, - names, + names.clone(), &tokenizer_from_hf(&mlx.hf), mmproj_cfg, &|n| mapper.norm_shift(n), &|n| mapper.rope_permute_heads(n, &config_for_permute), + &|n| mapper.value_transform(n), + mapper.shape_fastest_first(), + // Quantized MLX source: mirror its per-tensor quantization + // decisions (transplant what it quantized, carry what it kept + // unquantized). + true, &|n| mapper.row_rms_normalize(n, &config_for_permute), - ) + )?; + if ctx.validate { + validate_mlx_bundle(output, &mlx, &names, mapper.canonical_arch(), &|n| { + mapper.norm_shift(n) + })?; + } + Ok(()) +} + +/// `--validate` gate for `--mlx-passthrough` conversions: re-open the +/// written `.base` and byte-compare every main-bundle tensor against +/// the MLX source. Acceptance is exact equality, not "close" — this is +/// what makes the bundle a trustworthy oracle for index-level DSA +/// validation later. +/// +/// Coverage by stored dtype: +/// * `BaseQ4`/`BaseQ8` with a quantized MLX source — packed nibbles, +/// scales and biases must match the source verbatim. +/// * `BaseQ4`/`BaseQ8` from an unquantized source — recomputed via +/// the same deterministic pack path and compared. +/// * `F16` — recomputed dequant→f16 (norm shift applied on 1-D) and +/// compared. +/// * anything else (F32 SSM tensors etc.) is counted as skipped. +/// +/// Note: assumes source names map 1:1 to written tensors (true for MLX +/// sources — the stacking/splitting providers are no-ops there). +fn validate_mlx_bundle( + output: &std::path::Path, + mlx: &base_readers::mlx::MlxDir, + source_names: &[String], + canonical_arch: &str, + norm_shift: &dyn Fn(&str) -> f32, +) -> Result<()> { + use base_format::{BaseReader, TensorDtype}; + let reader = BaseReader::open(output).context("re-opening written .base for --validate")?; + let (mut n_pass, mut n_f16, mut n_repack, mut n_skip) = (0usize, 0usize, 0usize, 0usize); + for n in source_names { + let Some(Canonical::Main(canonical)) = to_canonical_name(n, canonical_arch) else { + continue; + }; + let entry = reader + .header() + .tensors + .iter() + .find(|t| t.name == canonical) + .with_context(|| format!("--validate: {canonical} missing from written .base"))? + .clone(); + let data = reader.tensor_bytes(&canonical)?; + match entry.dtype { + TensorDtype::BaseQ4 | TensorDtype::BaseQ8 => { + let bits = if entry.dtype == TensorDtype::BaseQ4 { + 4 + } else { + 8 + }; + let gs = entry.group_size.unwrap_or(64); + let scale_off = entry.scale_offset.unwrap_or(data.len() as u64) as usize; + let bias_off = entry.bias_offset.unwrap_or(data.len() as u64) as usize; + let (got_packed, got_scales, got_biases) = ( + &data[..scale_off], + &data[scale_off..bias_off], + &data[bias_off..], + ); + match mlx.tensor_packed(n, bits, gs)? { + Some(p) => { + // Scales/biases are compared against what the WRITE + // path emits, not the raw source bytes: base_q4 stores + // f16, so a bf16-scaled checkpoint (mlx-lm >= 0.20) is + // narrowed on the way in. Comparing raw bf16 here + // failed every such tensor deterministically, which + // made --validate unusable on current checkpoints. The + // packed weights themselves are still byte-verbatim. + let (want_scales, _) = + base_readers::mlx::narrow_to_f16_le(p.scales, p.scale_dtype)?; + let (want_biases, _) = + base_readers::mlx::narrow_to_f16_le(p.biases, p.scale_dtype)?; + ensure_bytes_eq(&canonical, "packed weights", got_packed, p.packed)?; + ensure_bytes_eq(&canonical, "scales", got_scales, &want_scales)?; + ensure_bytes_eq(&canonical, "biases", got_biases, &want_biases)?; + n_pass += 1; + } + None => { + let f32s = mlx.tensor_to_f32(n)?; + let target = if bits == 4 { + TargetScheme::BaseQ4 + } else { + TargetScheme::BaseQ8 + }; + let (packed, _) = pack_for_target(&f32s, target)?; + ensure_bytes_eq( + &canonical, + "packed weights", + got_packed, + &packed.packed_weights, + )?; + ensure_bytes_eq(&canonical, "scales", got_scales, &packed.scales)?; + ensure_bytes_eq(&canonical, "biases", got_biases, &packed.biases)?; + n_repack += 1; + } + } + } + TensorDtype::F16 => { + let mut f32s = mlx.tensor_to_f32(n)?; + if entry.shape.len() == 1 { + let s = norm_shift(&canonical); + if s != 0.0 { + for v in f32s.iter_mut() { + *v += s; + } + } + } + let expect: Vec = f32s + .iter() + .flat_map(|&f| half::f16::from_f32(f).to_le_bytes()) + .collect(); + ensure_bytes_eq(&canonical, "f16 payload", data, &expect)?; + n_f16 += 1; + } + _ => { + n_skip += 1; + } + } + } + eprintln!( + " validate: OK — byte-identical to source ({n_pass} passthrough, {n_f16} f16, \ + {n_repack} repacked, {n_skip} skipped)" + ); + Ok(()) } -fn tokenizer_from_hf(hf: &base_readers::hf::HfDir) -> std::collections::BTreeMap { +fn ensure_bytes_eq(canonical: &str, what: &str, got: &[u8], want: &[u8]) -> Result<()> { + if got.len() != want.len() { + bail!( + "--validate: {canonical} {what}: written length {} != source {}", + got.len(), + want.len() + ); + } + if got != want { + let idx = got + .iter() + .zip(want.iter()) + .position(|(a, b)| a != b) + .unwrap_or(0); + bail!("--validate: {canonical} {what}: byte mismatch at offset {idx}"); + } + Ok(()) +} + +fn tokenizer_from_hf( + hf: &base_readers::hf::HfDir, +) -> std::collections::BTreeMap { use serde_json::json; let mut m = std::collections::BTreeMap::new(); m.insert("tokenizer_type".into(), json!("hf")); @@ -2074,7 +2416,10 @@ fn mmproj_config_from_gguf( m.insert("out_hidden_size".into(), json!(out_hidden)); m.insert("projector_hidden_size".into(), json!(projector_hidden)); - if let Some(v) = resolve_u64("image_token_id", defaults.as_ref().map(|d| d.image_token_id)) { + if let Some(v) = resolve_u64( + "image_token_id", + defaults.as_ref().map(|d| d.image_token_id), + ) { m.insert("image_token_id".into(), v); } if let Some(v) = resolve_u64( @@ -2085,9 +2430,7 @@ fn mmproj_config_from_gguf( } if let Some(v) = resolve_u64( "vision_soft_tokens_per_image", - defaults - .as_ref() - .map(|d| d.vision_soft_tokens_per_image), + defaults.as_ref().map(|d| d.vision_soft_tokens_per_image), ) { m.insert("vision_soft_tokens_per_image".into(), v); } @@ -2114,16 +2457,15 @@ fn mmproj_config_from_gguf( /// text-only models. Captures the bits the runtime needs to drive the /// vision / audio prefill paths: tower configs, multimodal token IDs, /// soft-token counts, and image-pooling parameters. -fn mmproj_config_from_hf(hf: &base_readers::hf::HfDir) -> std::collections::BTreeMap { +fn mmproj_config_from_hf( + hf: &base_readers::hf::HfDir, +) -> std::collections::BTreeMap { let mut m = std::collections::BTreeMap::new(); let cfg = &hf.config; // Multimodal sub-configs — passed through verbatim. Runtime parses // hidden_size / num_hidden_layers / patch_size / etc. from these. - for key in [ - "vision_config", - "audio_config", - ] { + for key in ["vision_config", "audio_config"] { if let Some(v) = cfg.get(key) { m.insert(key.into(), v.clone()); } @@ -2180,26 +2522,181 @@ fn mmproj_config_from_hf(hf: &base_readers::hf::HfDir) -> std::collections::BTre } /// Abstraction over GGUF vs HF vs MLX so the shared convert logic -/// doesn't care where bytes come from. All sources are dequantized to -/// f32 and re-packed via the profile-driven canonical path. +/// doesn't care where bytes come from. Sources are dequantized to f32 +/// and re-packed via the profile-driven canonical path — except where +/// `packed_base_q4` can transplant an identically-schemed tensor whole. +/// True when every value survives an f32 -> f16 -> f32 round trip +/// bit-exactly (the "carry unquantized" dtype choice must be lossless — +/// in-range bf16 always passes because bf16's 7 mantissa bits fit in +/// f16's 10; an f32 source generally does not). +fn f16_lossless(vals: &[f32]) -> bool { + vals.iter() + .all(|&v| half::f16::from_f32(v).to_f32().to_bits() == v.to_bits()) +} + +/// True when every value survives an f32 -> bf16 -> f32 round trip +/// bit-exactly (covers bf16 sources whose values left f16's range). +fn bf16_lossless(vals: &[f32]) -> bool { + vals.iter() + .all(|&v| half::bf16::from_f32(v).to_f32().to_bits() == v.to_bits()) +} + +/// Stable label for a value transform, recorded in provenance so +/// verification tooling can re-derive the expectation independently. +fn vt_label(vt: base_arch::ValueTransform) -> &'static str { + match vt { + base_arch::ValueTransform::NegExp => "neg_exp", + } +} + trait TensorProvider { fn source_shape(&self, name: &str) -> Result>; fn to_f32(&self, name: &str) -> Result>; + + /// Hand back this tensor already in `base_q4`'s layout, when the + /// source stores the identical scheme so the bytes can be + /// transplanted rather than dequantized and requantized. + /// + /// Only MLX affine-q4 sources answer this; every other provider + /// takes the default `None` and goes through the f32 path. See + /// `base_readers::mlx::MlxDir::packed_base_q4` for why the round + /// trip is not the identity, and therefore worth avoiding. + fn packed_base_q4(&self, _name: &str, _group_size: u32) -> Result> { + Ok(None) + } + + /// Whether the *source checkpoint* stores this tensor quantized. + /// Drives the mirror policy: a tensor the source keeps unquantized + /// must be carried unquantized, whatever the target scheme would + /// otherwise do to a 2-D weight (a requantized MoE router — 0.03% + /// of the parameters — once set a whole correctness run's accuracy + /// floor). Only MLX sources answer `true`; unquantized sources are + /// wholly the target scheme's business. + fn source_is_quantized(&self, _name: &str) -> bool { + false + } + + /// Hand back this tensor already in the bundle's `nvfp4` layout when + /// the source stores NVFP4 (packed e2m1 code bytes + e4m3 per-block-16 + /// scale bytes, both copied verbatim — a transplant, never a + /// requantization). HF NVFP4 safetensors sources answer this; every + /// other provider takes the default `None`. + fn packed_nvfp4(&self, _name: &str) -> Result> { + Ok(None) + } } struct HfTensorProvider<'a> { hf: &'a base_readers::hf::HfDir, } +impl<'a> HfTensorProvider<'a> { + /// The e4m3 block-scale sibling of an NVFP4-quantized `.weight`, when + /// the checkpoint stores one (modelopt convention: `.weight` U8 + /// packed codes + `.weight_scale` F8_E4M3 per-block-16 scales). + fn nvfp4_scale_sibling(&self, name: &str) -> Option { + use base_readers::safetensors::StDtype; + let stem = name.strip_suffix(".weight")?; + let info = self.hf.tensor_info(name)?; + if info.dtype != StDtype::U8 { + return None; + } + let scale = format!("{stem}.weight_scale"); + self.hf.tensor_info(&scale)?; + Some(scale) + } +} impl<'a> TensorProvider for HfTensorProvider<'a> { fn source_shape(&self, name: &str) -> Result> { - self.hf + let info = self + .hf .tensor_info(name) - .map(|t| t.shape.clone()) - .ok_or_else(|| anyhow::anyhow!("tensor {name} missing")) + .ok_or_else(|| anyhow::anyhow!("tensor {name} missing"))?; + let mut shape = info.shape.clone(); + // NVFP4-quantized weights store two e2m1 codes per byte; report + // the logical (unpacked) shape like the MLX provider does. + if self.nvfp4_scale_sibling(name).is_some() { + if let Some(last) = shape.last_mut() { + *last *= 2; + } + } + Ok(shape) } fn to_f32(&self, name: &str) -> Result> { self.hf.tensor_to_f32(name) } + fn source_is_quantized(&self, name: &str) -> bool { + self.nvfp4_scale_sibling(name).is_some() + } + fn packed_nvfp4(&self, name: &str) -> Result> { + let Some(scale_name) = self.nvfp4_scale_sibling(name) else { + return Ok(None); + }; + let info = self + .hf + .tensor_info(name) + .expect("checked by sibling lookup"); + let codes = self + .hf + .tensor_bytes(name) + .ok_or_else(|| anyhow::anyhow!("tensor bytes for {name} missing"))?; + let scales = self + .hf + .tensor_bytes(&scale_name) + .ok_or_else(|| anyhow::anyhow!("tensor bytes for {scale_name} missing"))?; + let total_values: u64 = info.shape.iter().product::() * 2; + if codes.len() as u64 * 2 != total_values { + bail!( + "{name}: {} packed bytes for {} logical values (expected 2 codes/byte)", + codes.len(), + total_values + ); + } + if scales.len() as u64 * 16 != total_values { + bail!( + "{name}: {} e4m3 scale bytes for {} values (expected one per block of 16)", + scales.len(), + total_values + ); + } + // The global f32 scale (`dequant = code × block_scale × scale_2`) + // and the calibrated activation scale (`input_scale`, needed for + // vLLM-recipe static activation quantization) ride as an f32 pair + // in the tensor's bias region so runtime kernels find them in the + // same slab as the codes and block scales. Both are *also* carried + // as their own f32 sidecar tensors for verification. + let stem = name + .strip_suffix(".weight") + .expect("checked by sibling lookup"); + let scale2 = self + .hf + .tensor_bytes(&format!("{stem}.weight_scale_2")) + .ok_or_else(|| anyhow::anyhow!("{stem}.weight_scale_2 missing"))?; + if scale2.len() != 4 { + bail!( + "{stem}.weight_scale_2 is not a single f32 ({} bytes)", + scale2.len() + ); + } + let in_scale = self + .hf + .tensor_bytes(&format!("{stem}.input_scale")) + .ok_or_else(|| anyhow::anyhow!("{stem}.input_scale missing"))?; + if in_scale.len() != 4 { + bail!( + "{stem}.input_scale is not a single f32 ({} bytes)", + in_scale.len() + ); + } + let mut biases = scale2.to_vec(); + biases.extend_from_slice(in_scale); + Ok(Some(base_quant::Packed { + packed_weights: codes.to_vec(), + scales: scales.to_vec(), + biases, + group_size: base_quant::nvfp4::GROUP_SIZE as u32, + scale_dtype: Some(base_format::ScaleDtype::E4m3), + })) + } } struct MlxTensorProvider<'a> { @@ -2219,6 +2716,33 @@ impl<'a> TensorProvider for MlxTensorProvider<'a> { fn to_f32(&self, name: &str) -> Result> { self.mlx.tensor_to_f32(name) } + fn packed_base_q4(&self, name: &str, group_size: u32) -> Result> { + let Some(p) = self.mlx.packed_base_q4(name, group_size)? else { + return Ok(None); + }; + if p.out_of_f16_range > 0 { + // bf16 scales carry a wider exponent than f16. If any left + // f16's range the transplanted tensor would decode to + // inf/NaN — worse than requantizing, so refuse this tensor + // and let the caller fall back. + eprintln!( + " note: {name}: {} bf16 scale/bias value(s) outside f16 range — \ + requantizing this tensor instead of transplanting", + p.out_of_f16_range + ); + return Ok(None); + } + Ok(Some(base_quant::Packed { + packed_weights: p.packed_weights, + scales: p.scales, + biases: p.biases, + group_size: p.group_size, + scale_dtype: Some(base_format::ScaleDtype::F16), + })) + } + fn source_is_quantized(&self, name: &str) -> bool { + self.mlx.unpacked_shape(name).is_some() + } } /// Wraps a provider to stack HF-mainline MoE *per-expert* tensors into the @@ -2239,7 +2763,10 @@ struct StackingProvider<'a> { impl<'a> StackingProvider<'a> { fn build(inner: &'a dyn TensorProvider, source_names: &[String]) -> Self { use std::collections::{BTreeMap, HashSet}; - const MARK: &str = ".mlp.experts."; + // `.experts.` rather than `.mlp.experts.`: Nemotron-H NVFP4 + // checkpoints hold per-expert tensors under `…mixer.experts.{e}.…`. + // `shared_experts.` cannot match (no dot before "experts"). + const MARK: &str = ".experts."; let mut groups: BTreeMap> = BTreeMap::new(); let mut consumed: HashSet = HashSet::new(); for n in source_names { @@ -2251,7 +2778,19 @@ impl<'a> StackingProvider<'a> { if e_str.is_empty() || !e_str.bytes().all(|b| b.is_ascii_digit()) { continue; // already-fused (`experts.gate_proj.weight`) — leave to canon } - if !matches!(tail, ".gate_proj.weight" | ".up_proj.weight" | ".down_proj.weight") { + // `.weight_scale` (the e4m3 block scales) is deliberately NOT + // stacked: it is consumed byte-verbatim by the nvfp4 weight + // transplant. The f32 scalar sidecars stack into small arrays. + if !matches!( + tail, + ".gate_proj.weight" + | ".up_proj.weight" + | ".down_proj.weight" + | ".up_proj.weight_scale_2" + | ".down_proj.weight_scale_2" + | ".up_proj.input_scale" + | ".down_proj.input_scale" + ) { continue; } let e: usize = e_str.parse().unwrap_or(usize::MAX); @@ -2264,10 +2803,17 @@ impl<'a> StackingProvider<'a> { es.sort_by_key(|(e, _)| *e); stacks.insert(v, es.into_iter().map(|(_, n)| n).collect()); } - let mut rewritten: Vec = - source_names.iter().filter(|n| !consumed.contains(*n)).cloned().collect(); + let mut rewritten: Vec = source_names + .iter() + .filter(|n| !consumed.contains(*n)) + .cloned() + .collect(); rewritten.extend(stacks.keys().cloned()); - StackingProvider { inner, stacks, rewritten } + StackingProvider { + inner, + stacks, + rewritten, + } } fn rewritten_names(&self) -> Vec { self.rewritten.clone() @@ -2298,6 +2844,53 @@ impl TensorProvider for StackingProvider<'_> { None => self.inner.to_f32(name), } } + fn packed_base_q4(&self, name: &str, group_size: u32) -> Result> { + // A stacked tensor is assembled from several source tensors with + // independent scales — there are no contiguous bytes to hand + // over. Only pass through for names we don't stack. + match self.stacks.get(name) { + Some(_) => Ok(None), + None => self.inner.packed_base_q4(name, group_size), + } + } + fn packed_nvfp4(&self, name: &str) -> Result> { + // Unlike base_q4 above, an nvfp4 stack CAN be transplanted: the + // bundle's stacked layout is per-expert sections concatenated in + // expert order — codes `[E][out][in/2]`, then block scales + // `[E][out][in/16]`, then the per-expert f32 global scales `[E]` + // in the bias region. + let Some(parts) = self.stacks.get(name) else { + return self.inner.packed_nvfp4(name); + }; + let mut codes: Vec = Vec::new(); + let mut scales: Vec = Vec::new(); + let mut biases: Vec = Vec::new(); + for p in parts { + let Some(packed) = self.inner.packed_nvfp4(p)? else { + return Ok(None); // mixed stack — let the f32 path decide + }; + codes.extend_from_slice(&packed.packed_weights); + scales.extend_from_slice(&packed.scales); + biases.extend_from_slice(&packed.biases); + } + Ok(Some(base_quant::Packed { + packed_weights: codes, + scales, + biases, + group_size: base_quant::nvfp4::GROUP_SIZE as u32, + scale_dtype: Some(base_format::ScaleDtype::E4m3), + })) + } + fn source_is_quantized(&self, name: &str) -> bool { + match self.stacks.get(name) { + // Report the constituents' storage so the mirror policy sees + // through the virtual name (all experts share one scheme). + Some(parts) => parts + .first() + .is_some_and(|p| self.inner.source_is_quantized(p)), + None => self.inner.source_is_quantized(name), + } + } } /// Inverse of [`StackingProvider`]: some HF checkpoints ship *fused* attention @@ -2355,20 +2948,34 @@ impl<'a> SplittingProvider<'a> { } splits.insert( format!("{before}.self_attn.q_proj.weight"), - SplitSpec { src: n.clone(), row_off: 0, row_cnt: q }, + SplitSpec { + src: n.clone(), + row_off: 0, + row_cnt: q, + }, ); splits.insert( format!("{before}.self_attn.k_proj.weight"), - SplitSpec { src: n.clone(), row_off: q, row_cnt: k }, + SplitSpec { + src: n.clone(), + row_off: q, + row_cnt: k, + }, ); splits.insert( format!("{before}.self_attn.v_proj.weight"), - SplitSpec { src: n.clone(), row_off: q + k, row_cnt: v }, + SplitSpec { + src: n.clone(), + row_off: q + k, + row_cnt: v, + }, ); consumed.insert(n.clone()); } else if let Some(before) = n.strip_suffix(".mlp.gate_up_proj.weight") { if ffn == 0 { - bail!("fused {n}: intermediate_size must be set in config to split gate_up_proj"); + bail!( + "fused {n}: intermediate_size must be set in config to split gate_up_proj" + ); } let shape = inner.source_shape(n)?; let rows = shape.first().copied().unwrap_or(0); @@ -2381,20 +2988,35 @@ impl<'a> SplittingProvider<'a> { } splits.insert( format!("{before}.mlp.gate_proj.weight"), - SplitSpec { src: n.clone(), row_off: 0, row_cnt: ffn }, + SplitSpec { + src: n.clone(), + row_off: 0, + row_cnt: ffn, + }, ); splits.insert( format!("{before}.mlp.up_proj.weight"), - SplitSpec { src: n.clone(), row_off: ffn, row_cnt: ffn }, + SplitSpec { + src: n.clone(), + row_off: ffn, + row_cnt: ffn, + }, ); consumed.insert(n.clone()); } } - let mut rewritten: Vec = - source_names.iter().filter(|n| !consumed.contains(*n)).cloned().collect(); + let mut rewritten: Vec = source_names + .iter() + .filter(|n| !consumed.contains(*n)) + .cloned() + .collect(); rewritten.extend(splits.keys().cloned()); - Ok(SplittingProvider { inner, splits, rewritten }) + Ok(SplittingProvider { + inner, + splits, + rewritten, + }) } fn rewritten_names(&self) -> Vec { self.rewritten.clone() @@ -2429,6 +3051,159 @@ impl TensorProvider for SplittingProvider<'_> { None => self.inner.to_f32(name), } } + fn packed_base_q4(&self, name: &str, group_size: u32) -> Result> { + match self.splits.get(name) { + Some(_) => Ok(None), + None => self.inner.packed_base_q4(name, group_size), + } + } + fn packed_nvfp4(&self, name: &str) -> Result> { + match self.splits.get(name) { + Some(_) => Ok(None), + None => self.inner.packed_nvfp4(name), + } + } + fn source_is_quantized(&self, name: &str) -> bool { + match self.splits.get(name) { + Some(spec) => self.inner.source_is_quantized(&spec.src), + None => self.inner.source_is_quantized(name), + } + } +} + +/// GLM 5.2 (glm_dsa) HF checkpoints ship the MLA up-projection FUSED as +/// `self_attn.kv_b_proj.weight` `[n_heads·(qk_nope+v_head), kv_lora]`, but +/// the runtime's absorb kernels read the SPLIT per-head forms (k_b with the +/// per-head block transposed). Expose MLX-sanitize-equivalent virtual +/// tensors (mlx-lm `deepseek_v32.sanitize`: +/// `embed_q = kv_b[:, :nope, :].swapaxes(-1,-2)`, +/// `unembed_out = kv_b[:, nope:, :]`) and drop the fused source; +/// `glm_hf_canonical` renames them to `k_b_proj`/`v_b_proj` and the +/// force-f16 rule stores them raw f16 — the same layout an MLX-sourced +/// bundle carries (GLM5.2_DSA.md §3a). Pure pass-through for every other +/// arch and for sources that already ship them split. +struct GlmKvbSplitProvider<'a> { + inner: &'a dyn TensorProvider, + /// virtual name -> (fused source name, is_k_b) + splits: std::collections::BTreeMap, + rewritten: Vec, + n_heads: usize, + qk_nope: usize, + v_head: usize, + kv_lora: usize, +} +impl<'a> GlmKvbSplitProvider<'a> { + fn build( + inner: &'a dyn TensorProvider, + source_names: &[String], + config: &base_arch::ArchConfig, + canonical_arch: &str, + ) -> Self { + let mut splits = std::collections::BTreeMap::new(); + let mut rewritten = Vec::with_capacity(source_names.len()); + let applies = canonical_arch == "glm_dsa" + && config.kv_lora_rank > 0 + && config.qk_nope_head_dim > 0 + && config.v_head_dim > 0; + for n in source_names { + if applies && n.ends_with(".self_attn.kv_b_proj.weight") { + let kb = n.replace(".kv_b_proj.", ".embed_q."); + let vb = n.replace(".kv_b_proj.", ".unembed_out."); + splits.insert(kb.clone(), (n.clone(), true)); + splits.insert(vb.clone(), (n.clone(), false)); + rewritten.push(kb); + rewritten.push(vb); + } else { + rewritten.push(n.clone()); + } + } + GlmKvbSplitProvider { + inner, + splits, + rewritten, + n_heads: config.num_attention_heads as usize, + qk_nope: config.qk_nope_head_dim as usize, + v_head: config.v_head_dim as usize, + kv_lora: config.kv_lora_rank as usize, + } + } + fn rewritten_names(&self) -> Vec { + self.rewritten.clone() + } +} +impl TensorProvider for GlmKvbSplitProvider<'_> { + fn source_shape(&self, name: &str) -> Result> { + match self.splits.get(name) { + Some((_, true)) => Ok(vec![ + self.n_heads as u64, + self.kv_lora as u64, + self.qk_nope as u64, + ]), + Some((_, false)) => Ok(vec![ + self.n_heads as u64, + self.v_head as u64, + self.kv_lora as u64, + ]), + None => self.inner.source_shape(name), + } + } + fn to_f32(&self, name: &str) -> Result> { + let Some((src, is_kb)) = self.splits.get(name) else { + return self.inner.to_f32(name); + }; + let (h, nope, v, lora) = (self.n_heads, self.qk_nope, self.v_head, self.kv_lora); + let head_dim = nope + v; + let data = self.inner.to_f32(src)?; + if data.len() != h * head_dim * lora { + bail!( + "kv_b_proj {} has {} elements, expected n_heads({h})*(qk_nope({nope})+v_head({v}))*kv_lora({lora})", + src, + data.len() + ); + } + // Source is row-major [h*head_dim, lora]; row r of head hh is + // data[(hh*head_dim + r)*lora ..][..lora]. + if *is_kb { + // k_b: [h, lora, nope] — per-head transpose of the first `nope` rows. + let mut out = vec![0f32; h * lora * nope]; + for hh in 0..h { + let src_base = hh * head_dim * lora; + let dst_base = hh * lora * nope; + for r in 0..nope { + for l in 0..lora { + out[dst_base + l * nope + r] = data[src_base + r * lora + l]; + } + } + } + Ok(out) + } else { + // v_b: [h, v, lora] — the last `v` rows per head, layout kept. + let mut out = Vec::with_capacity(h * v * lora); + for hh in 0..h { + let start = (hh * head_dim + nope) * lora; + out.extend_from_slice(&data[start..start + v * lora]); + } + Ok(out) + } + } + fn packed_base_q4(&self, name: &str, group_size: u32) -> Result> { + match self.splits.get(name) { + Some(_) => Ok(None), + None => self.inner.packed_base_q4(name, group_size), + } + } + fn packed_nvfp4(&self, name: &str) -> Result> { + match self.splits.get(name) { + Some(_) => Ok(None), + None => self.inner.packed_nvfp4(name), + } + } + fn source_is_quantized(&self, name: &str) -> bool { + match self.splits.get(name) { + Some((src, _)) => self.inner.source_is_quantized(src), + None => self.inner.source_is_quantized(name), + } + } } #[allow(clippy::too_many_arguments)] @@ -2445,14 +3220,33 @@ fn convert_generic( mmproj_config: std::collections::BTreeMap, norm_shift: &dyn Fn(&str) -> f32, rope_permute: &dyn Fn(&str) -> Option, + value_transform: &dyn Fn(&str) -> Option, + shape_fastest_first: bool, + mirror_unquantized: bool, row_rms_normalize: &dyn Fn(&str) -> Option, ) -> Result<()> { use base_format::{ - AlignmentConfig, BaseReader, BaseWriter, ComputeRegion, Header, HeaderFlags, LayerKind, - LayerDescriptor, LayerPrecision, ModelConfig, QuantScheme, SourceInfo, TargetBackend, TensorDtype, - TensorFlags, TensorPayload, TokenizerBlob, + AlignmentConfig, BaseReader, BaseWriter, ComputeRegion, Header, HeaderFlags, ModelConfig, + QuantScheme, SourceInfo, TargetBackend, TensorDtype, TensorFlags, TensorPayload, + TokenizerBlob, }; let target = ctx.target; + // Transplanting only makes sense when the bundle's scheme *is* the + // source's scheme. A `--profile` run picks a scheme per tensor, so + // the packed source bytes may not be what that tensor should end up + // as — leave those to the f32 path. + let allow_q4_passthrough = + ctx.q4_passthrough && ctx.profile.is_none() && target == TargetScheme::BaseQ4; + // NVFP4 sources (modelopt HF checkpoints) transplant the same way: + // packed e2m1 codes + e4m3 block scales copied verbatim. + let allow_nvfp4_passthrough = ctx.profile.is_none() && target == TargetScheme::Nvfp4; + // Mirror policy (quantized sources only): the checkpoint's own + // quantization decisions are the bundle's. Tensors it quantized are + // transplanted (above); tensors it kept unquantized are carried + // unquantized below, never requantized. Only meaningful alongside the + // transplant — a profile / non-matching target already opted out of + // holding the source's exact weights. + let mirror = mirror_unquantized && (allow_q4_passthrough || allow_nvfp4_passthrough); let quant_scheme = match target { TargetScheme::BaseQ2 => QuantScheme::BaseQ2, TargetScheme::BaseQ3 => QuantScheme::BaseQ3, @@ -2493,6 +3287,18 @@ fn convert_generic( let source_names = splitting.rewritten_names(); let provider: &dyn TensorProvider = &splitting; + // GLM 5.2: derive the split MLA up-projections (k_b/v_b in absorb + // layout) from the fused kv_b_proj the HF bf16 checkpoint ships. + let kvb = GlmKvbSplitProvider::build(provider, &source_names, &config, canonical_arch); + if !kvb.splits.is_empty() { + eprintln!( + " mla: split {} fused kv_b_proj into k_b/v_b (absorb layout, f16)", + kvb.splits.len() / 2 + ); + } + let source_names = kvb.rewritten_names(); + let provider: &dyn TensorProvider = &kvb; + // Map source names → canonical. Use the same llama-style map that // GGUF uses for blk.N.* tensors, plus a HF-style pass-through for // `model.layers.N.*` already-canonical names. Multimodal towers @@ -2500,19 +3306,19 @@ fn convert_generic( // whether to materialize them. let mut mapped: Vec<(String, String)> = Vec::new(); let mut mmproj_mapped: Vec<(String, String)> = Vec::new(); - let mut dropped = 0usize; + let mut dropped_names: Vec = Vec::new(); for n in &source_names { match to_canonical_name(n, canonical_arch) { Some(Canonical::Main(c)) => mapped.push((n.clone(), c)), Some(Canonical::Mmproj(c)) => mmproj_mapped.push((n.clone(), c)), - None => dropped += 1, + None => dropped_names.push(n.clone()), } } eprintln!( " mapped: {} tensors kept, {} mmproj, {} dropped", mapped.len(), mmproj_mapped.len(), - dropped + dropped_names.len() ); let mut tok_fields = tokenizer_fields.clone(); @@ -2533,11 +3339,21 @@ fn convert_generic( let mut found_any = false; for (src_name, canonical) in &mapped { // canonical = "layers.{N}.mlp.down_proj.weight" - let Some(rest) = canonical.strip_prefix("layers.") else { continue }; - let Some((idx_str, tail)) = rest.split_once('.') else { continue }; - if tail != "mlp.down_proj.weight" { continue } - let Ok(layer) = idx_str.parse::() else { continue }; - if layer >= per_layer_ffn.len() { continue } + let Some(rest) = canonical.strip_prefix("layers.") else { + continue; + }; + let Some((idx_str, tail)) = rest.split_once('.') else { + continue; + }; + if tail != "mlp.down_proj.weight" { + continue; + } + let Ok(layer) = idx_str.parse::() else { + continue; + }; + if layer >= per_layer_ffn.len() { + continue; + } // down_proj shape (HF unpacked): [hidden_size, ffn_size] if let Ok(shape) = provider.source_shape(src_name) { if shape.len() == 2 { @@ -2551,7 +3367,9 @@ fn convert_generic( // so the runtime has a value for every layer. let fallback = config.intermediate_size; for v in per_layer_ffn.iter_mut() { - if *v == 0 { *v = fallback; } + if *v == 0 { + *v = fallback; + } } // Only emit per_layer_ffn when FFN width actually varies // layer-to-layer (Gemma 4 E2B: 6144 for own-KV, 12288 for @@ -2588,22 +3406,24 @@ fn convert_generic( }, metadata: Default::default(), target_backend: TargetBackend::Metal, - quant_profile: ctx.profile_name().unwrap_or("").to_string(), - alignment: AlignmentConfig::default(), - flags: HeaderFlags::QUANTIZED, - layers: (0..config.num_hidden_layers) - .map(|_| LayerDescriptor { - kind: LayerKind::AttentionGqa, - moe_n_experts: 0, - moe_n_active: 0, - shared_attn_layer: None, - compute_hint: Some(ComputeRegion::Accelerator), - precision: LayerPrecision::default(), - }) - .collect(), + quant_profile: ctx.profile_name().unwrap_or("").to_string(), + // Accel tensors align to the 16 KiB Apple page (default is 64 B): + // the runtime's BaseWeightStore can then always split its chunked + // no-copy mmap at a tensor start. Mixed-dtype bundles with 64 B + // alignment can run a whole max_buffer_size window without a + // page-aligned start (seen on the GLM 5.2 q4/q5/q6 production + // bundle), forcing overlap-mapped splits or per-tensor copies. + // Padding cost: < tensor_count × 16 KiB — noise on any real model. + alignment: AlignmentConfig { + accel_align_log2: 14, + ..Default::default() + }, + flags: HeaderFlags::QUANTIZED, + layers: layer_descriptors_from_config(&config), tensors: vec![], mmproj: None, calibration: None, + provenance: None, sig: None, }; @@ -2613,8 +3433,12 @@ fn convert_generic( let has_moe = source_names .iter() .any(|n| n.contains("experts") || n.contains("_exps") || n.contains("_shexp")); - let has_ssm = source_names.iter().any(|n| n.contains(".ssm.") || n.contains("ssm_")); - let has_attn = source_names.iter().any(|n| n.contains("self_attn") || n.contains("attn_q")); + let has_ssm = source_names.iter().any(|n| n.contains(".ssm.") || n.contains("ssm_")) + // HF Nemotron-H names carry no ".ssm." — detect on the canonical side. + || mapped.iter().any(|(_, c)| c.contains(".ssm.")); + let has_attn = source_names + .iter() + .any(|n| n.contains("self_attn") || n.contains("attn_q")); if has_moe { header.flags |= HeaderFlags::HAS_MOE; } @@ -2625,7 +3449,15 @@ fn convert_generic( } } - let mut writer = BaseWriter::create(output, header).context("create writer")?; + // 64 MiB header reserve in direct-write mode: ~2k tensor entries at + // ~400 B JSON each is well under 8 MiB even for GLM 5.2's 79-layer MoE; + // the padding cost is invisible against a 400+ GB bundle. + let mut writer = if ctx.direct_write { + BaseWriter::create_direct(output, header, 64 * 1024 * 1024) + .context("create writer (direct)")? + } else { + BaseWriter::create(output, header).context("create writer")? + }; let pb = indicatif::ProgressBar::new(mapped.len() as u64); pb.set_style( @@ -2633,13 +3465,111 @@ fn convert_generic( .expect("valid progress template") .progress_chars("=>-"), ); + let mut n_passthrough = 0usize; + let mut n_carried = 0usize; + let mut prov_tensors: std::collections::BTreeMap = + Default::default(); for (src_name, canonical) in &mapped { pb.set_message(canonical.clone()); - let shape = provider.source_shape(src_name)?; + let src_shape = provider.source_shape(src_name)?; + // The reduction dim is the source's last (HF C-order [out, in]); + // read it before any reordering of the reported shape. + let in_features = src_shape.last().copied().map(|d| d as usize); + let mut shape = if shape_fastest_first { + src_shape.iter().rev().copied().collect::>() + } else { + src_shape + }; + // Scalar sidecars (NVFP4 `weight_scale_2` / `input_scale`) are + // rank-0 in safetensors; the bundle stores them as [1]. + if shape.is_empty() { + shape = vec![1]; + } + // Nemotron HF conv1d weights arrive as [channels, 1, taps]; the + // canonical bundle shape is [1, taps, channels] (identical flat + // data — taps contiguous per channel — so this is metadata only, + // matching what MLX-sourced bundles record). + if canonical.ends_with(".ssm.conv1d.weight") && shape.len() == 3 && shape[1] == 1 { + shape = vec![1, shape[0], shape[2]]; + } + let shape = shape; + let ndims = shape.len(); - let mut f32s = provider - .to_f32(src_name) - .with_context(|| format!("reading {src_name}"))?; + let is_ssm_a = canonical == "ssm.a_log" + || canonical.ends_with(".ssm.a_log") + || src_name.ends_with(".ssm_a") + || is_gdn_a_log(canonical); + let is_ssm_sensitive = is_ssm_a + || canonical.ends_with(".ssm.dt_bias") + || canonical.ends_with(".ssm.d") + // Mamba-2 conv1d and the grouped gated norm. The scan runs + // its state in f32 and these feed it directly; the GGUF path + // keeps them f32 for the same reason, and a bundle that + // quantized them would not be interchangeable with one that + // did not. + || canonical.ends_with(".ssm.conv1d.weight") + || canonical.ends_with(".ssm.conv1d.bias") + || canonical.ends_with(".ssm.norm.weight") + || is_gdn_f32(canonical); + // NVFP4 sidecars: the per-tensor global scale and the calibrated + // activation scale. Always f32 — they parameterize the dequant + // itself, and the verify gate compares them value-exactly. + let is_sidecar = + canonical.ends_with(".weight_scale_2") || canonical.ends_with(".input_scale"); + // See GGUF path above for rationale on the f16/GPU route. Qwen3.5 GDN + // conv1d + per-head gate projections also route here (f16 on GPU). + let is_norm_like = !is_sidecar && (shape.len() == 1 || is_gdn_f16(canonical)); + let is_embedding = canonical == "embed_tokens.weight" || canonical == "lm_head.weight"; + // GLM MLA k_b/v_b projections and router weights are raw-f16 + // consumers and must never enter the quantized transplant path. + let force_f16 = canonical_arch == "glm_dsa" + && (canonical.ends_with(".k_b_proj.weight") + || canonical.ends_with(".v_b_proj.weight") + || canonical.ends_with(".mlp.router.weight")); + // Mirror policy: this tensor is unquantized in the (quantized) + // source checkpoint, so it must stay unquantized in the bundle. + let src_unquantized = mirror && !provider.source_is_quantized(src_name); + + // Zero-loss path: when the source already stores exactly the + // target scheme, transplant the packed bytes instead of + // dequantizing and requantizing. Skipped for tensors that don't + // reach the quantized branches at all, and for any tensor whose + // rows get permuted on the way through (the permutation is + // defined on values, not on packed groups). + let passthrough = if (allow_q4_passthrough || allow_nvfp4_passthrough) + && !is_ssm_sensitive + && !is_norm_like + && !is_sidecar + && !force_f16 + && rope_permute(canonical).is_none() + { + if allow_nvfp4_passthrough { + provider + .packed_nvfp4(src_name) + .with_context(|| format!("transplanting {src_name}"))? + } else { + provider + .packed_base_q4(src_name, base_quant::base_q4::GROUP_SIZE as u32) + .with_context(|| format!("transplanting {src_name}"))? + } + } else { + None + }; + let transplant_dtype = if allow_nvfp4_passthrough { + TensorDtype::Nvfp4 + } else { + TensorDtype::BaseQ4 + }; + + // Reading to f32 is the expensive step (and for a 30B model, the + // memory-hungry one) — skip it entirely when transplanting. + let mut f32s = if passthrough.is_some() { + Vec::new() + } else { + provider + .to_f32(src_name) + .with_context(|| format!("reading {src_name}"))? + }; // Per-arch hook: bake the +1 unit-offset into Gemma 3's // zero-centered RMSNorm gamma so the runtime can use the plain @@ -2656,6 +3586,14 @@ fn convert_generic( } } + // Per-arch hook: undo a stored reparameterization (Mamba-2 keeps + // the state-transition matrix as `A_log`; the scan wants + // `A = -exp(A_log)`). Applied before packing so the bundle holds + // the value the kernel consumes, matching the GGUF path. + if let Some(vt) = value_transform(canonical) { + vt.apply(&mut f32s); + } + // Per-arch hook: normalize HF "split-half" rotary q/k row layout to // the interleaved layout the runtime rope kernels (and GGUF sources) // use. Mirrors `convert_hf_to_gguf.py::LlamaModel.permute`. @@ -2689,22 +3627,9 @@ fn convert_generic( } let f32s_for = || -> &[f32] { &f32s }; + let transplanted = passthrough.is_some(); - let is_ssm_a = canonical == "ssm.a_log" - || canonical.ends_with(".ssm.a_log") - || src_name.ends_with(".ssm_a") - || is_gdn_a_log(canonical); - let is_ssm_sensitive = is_ssm_a - || canonical.ends_with(".ssm.dt_bias") - || canonical.ends_with(".ssm.d") - || is_gdn_f32(canonical); - // See GGUF path above for rationale on the f16/GPU route. Qwen3.5 GDN - // conv1d + per-head gate projections also route here (f16 on GPU). - let is_norm_like = shape.len() == 1 || is_gdn_f16(canonical); - let is_embedding = - canonical == "embed_tokens.weight" || canonical == "lm_head.weight"; - - let (entry, data) = if is_ssm_sensitive { + let (entry, data) = if is_ssm_sensitive || is_sidecar { let mut flags = TensorFlags::empty(); if is_ssm_a { flags |= TensorFlags::SSM_A_MATRIX; @@ -2729,25 +3654,41 @@ fn convert_generic( symmetric: false, flags, checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; let data: Vec = f32s_for().iter().flat_map(|f| f.to_le_bytes()).collect(); entry.length = data.len() as u64; (entry, data) - } else if is_norm_like { + } else if is_norm_like || force_f16 { // 1-D norms (and biases caught by the same shape check) are // always emitted at f16. A profile's catch-all `**.weight` // rule typically targets a quant bit-width; quantizing a // per-channel norm-gain wrecks the model. Override the // profile here so a profile that omits explicit norm - // patterns still produces a working bundle. - let (bytes, dtype) = ( - f32s_for() - .iter() - .flat_map(|&f| half::f16::from_f32(f).to_le_bytes()) - .collect::>(), - TensorDtype::F16, - ); + // patterns still produces a working bundle. GLM k_b/v_b and + // router tensors also route here via force_f16. + // + // Under the mirror policy, f16 must also represent the + // source's values losslessly — an f32 source (Nemotron's + // `e_score_correction_bias`) whose values don't round-trip + // is widened to f32 instead of silently rounded. + let (bytes, dtype) = if !force_f16 && src_unquantized && !f16_lossless(f32s_for()) { + ( + f32s_for() + .iter() + .flat_map(|f| f.to_le_bytes()) + .collect::>(), + TensorDtype::F32, + ) + } else { + ( + f32s_for() + .iter() + .flat_map(|&f| half::f16::from_f32(f).to_le_bytes()) + .collect::>(), + TensorDtype::F16, + ) + }; let entry = base_format::TensorEntry { name: canonical.clone(), dtype, @@ -2768,16 +3709,85 @@ fn convert_generic( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; + (entry, bytes) + } else if src_unquantized { + // Mirror policy, the 2-D case: the checkpoint keeps this + // tensor unquantized (Nemotron's bf16 MoE router is the + // canonical example), so requantizing it would make the + // bundle differ from the reference weights — carry it at the + // narrowest dtype that represents every value losslessly. + // The runtime serves f16/bf16/f32 2-D weights natively + // (f16 zero-copy; bf16/f32 converted to f16 at load). + n_carried += 1; + let (bytes, dtype) = if f16_lossless(f32s_for()) { + ( + f32s_for() + .iter() + .flat_map(|&f| half::f16::from_f32(f).to_le_bytes()) + .collect::>(), + TensorDtype::F16, + ) + } else if bf16_lossless(f32s_for()) { + ( + f32s_for() + .iter() + .flat_map(|&f| half::bf16::from_f32(f).to_le_bytes()) + .collect::>(), + TensorDtype::Bf16, + ) + } else { + ( + f32s_for() + .iter() + .flat_map(|f| f.to_le_bytes()) + .collect::>(), + TensorDtype::F32, + ) + }; + let entry = base_format::TensorEntry { + name: canonical.clone(), + dtype, + shape, + offset: 0, + length: bytes.len() as u64, + scale_offset: None, + scale_length: None, + bias_offset: None, + bias_length: None, + awq_scale_offset: None, + awq_scale_length: None, + group_size: None, + layout: None, + residency: Some(if is_embedding { + base_format::ResidencyHint::Hot + } else { + base_format::ResidencyHint::Warm + }), + compute_region: if is_embedding { + ComputeRegion::Gpu + } else { + ComputeRegion::Accelerator + }, + scale_dtype: None, + symmetric: false, + flags: TensorFlags::empty(), + checksum_xxh64: None, + source_ggml_type: None, + }; (entry, bytes) } else if is_embedding { // See GGUF path above for rationale on embed quantization. - let in_features = shape.last().copied().map(|d| d as usize); - let (packed, dtype) = if ctx.profile.is_some() { - ctx.pack_tensor(canonical, f32s_for(), in_features)? - } else { - pack_for_target(f32s_for(), target)? + let (packed, dtype) = match passthrough { + Some(p) => { + n_passthrough += 1; + (p, transplant_dtype) + } + None if ctx.profile.is_some() => { + ctx.pack_tensor(canonical, f32s_for(), in_features)? + } + None => pack_for_target_rows(f32s_for(), target, in_features, canonical)?, }; let mut data = Vec::with_capacity( packed.packed_weights.len() + packed.scales.len() + packed.biases.len(), @@ -2827,16 +3837,20 @@ fn convert_generic( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; entry.length = data.len() as u64; (entry, data) } else { - let in_features = shape.last().copied().map(|d| d as usize); - let (packed, dtype) = if ctx.profile.is_some() { - ctx.pack_tensor(canonical, f32s_for(), in_features)? - } else { - pack_for_target(f32s_for(), target)? + let (packed, dtype) = match passthrough { + Some(p) => { + n_passthrough += 1; + (p, transplant_dtype) + } + None if ctx.profile.is_some() => { + ctx.pack_tensor(canonical, f32s_for(), in_features)? + } + None => pack_for_target_rows(f32s_for(), target, in_features, canonical)?, }; let mut data = Vec::with_capacity( packed.packed_weights.len() + packed.scales.len() + packed.biases.len(), @@ -2887,17 +3901,64 @@ fn convert_generic( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; entry.length = data.len() as u64; (entry, data) }; writer.add_tensor(TensorPayload { entry, data }); + + // Record where this tensor came from and what was done to it, so + // verification tooling can check the bundle against the source + // checkpoint without a hand-maintained name table. + let mut rec = base_format::TensorProvenance { + transplanted, + carried: src_unquantized, + transform: value_transform(canonical).map(|vt| vt_label(vt).to_string()), + permuted: ndims >= 2 && rope_permute(canonical).is_some(), + ..Default::default() + }; + if ndims == 1 { + let s = norm_shift(canonical); + if s != 0.0 { + rec.norm_shift = Some(s); + } + } + if let Some(spec) = splitting.splits.get(src_name) { + rec.src = vec![spec.src.clone()]; + rec.rows = Some([spec.row_off, spec.row_cnt]); + } else if let Some(parts) = stacking.stacks.get(src_name) { + rec.stack = Some(base_format::StackRef { + pattern: src_name.replace(".experts.", ".experts.{e}."), + count: parts.len() as u32, + }); + } else { + rec.src = vec![src_name.clone()]; + } + prov_tensors.insert(canonical.clone(), rec); pb.inc(1); } pb.finish_and_clear(); - eprintln!(" quantized {} tensors", mapped.len()); + if n_passthrough > 0 || n_carried > 0 { + eprintln!( + " quantized {} tensors ({} transplanted from the source's identical \ + base_q4 scheme — no requantization; {} carried unquantized per the \ + source's own storage)", + mapped.len(), + n_passthrough, + n_carried + ); + } else { + eprintln!(" quantized {} tensors", mapped.len()); + } + writer.set_provenance(base_format::Provenance { + schema: 1, + mirror, + dropped: dropped_names, + mmproj: mmproj_mapped.iter().map(|(s, _)| s.clone()).collect(), + tensors: prov_tensors, + }); // Multimodal towers — preserve their HF names verbatim and route // them into the mmproj sub-bundle. Tower weights stay on the same @@ -2958,8 +4019,8 @@ fn convert_generic( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; (entry, bytes) } else { // Pad if needed so pack's group-size invariant holds for @@ -2976,14 +4037,12 @@ fn convert_generic( let pack_n: Result<_> = if ctx.profile.is_some() { ctx.pack_tensor(canonical, f32s_for(), in_features) } else { - pack_for_target(f32s_for(), target) + pack_for_target_rows(f32s_for(), target, in_features, canonical) }; match pack_n { Ok((packed, dtype)) => { let mut data = Vec::with_capacity( - packed.packed_weights.len() - + packed.scales.len() - + packed.biases.len(), + packed.packed_weights.len() + packed.scales.len() + packed.biases.len(), ); data.extend_from_slice(&packed.packed_weights); let scale_off = data.len() as u64; @@ -3030,8 +4089,8 @@ fn convert_generic( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; entry.length = data.len() as u64; (entry, data) } @@ -3064,8 +4123,8 @@ fn convert_generic( symmetric: false, flags: TensorFlags::empty(), checksum_xxh64: None, - source_ggml_type: None, -}; + source_ggml_type: None, + }; (entry, bytes) } } @@ -3238,6 +4297,15 @@ fn to_canonical_name(name: &str, arch: &str) -> Option { return nomic_bert_hf_rename(name).map(Canonical::Main); } + // Nemotron-H names everything `backbone.layers.N.mixer.*` regardless + // of whether the mixer is a Mamba-2 scan, an attention block or an + // MoE FFN, so the generic `model.layers.N.*` table below cannot tell + // them apart. Its own rename splits them back out. + if arch.starts_with("nemotron_h") && (name.starts_with("backbone.") || name == "lm_head.weight") + { + return base_arch::nemotron::nemotron_hf_rename(name).map(Canonical::Main); + } + // Strip HF naming prefix. Accept multiple multimodal wrapper // orderings: // - `model.language_model.*` (Gemma 4 26B-A4B, mainline HF) @@ -3259,6 +4327,14 @@ fn to_canonical_name(name: &str, arch: &str) -> Option { return None; } + // GLM 5.2 (glm_dsa) MLX/HF checkpoints — dedicated rename table so the + // bundle carries the same canonical names the GGUF mapper emits + // (header-equivalent across sources) without routing GLM through the + // generic MoE renames below (which would emit the ffn_*_exps forms). + if arch == "glm_dsa" { + return glm_hf_canonical(stripped).map(Canonical::Main); + } + // HF native names → canonical. if stripped != name { if stripped == "rotary_emb.inv_freq" { @@ -3326,7 +4402,10 @@ fn to_canonical_name(name: &str, arch: &str) -> Option { // the canonical name without going through the legacy // `mlp.experts.X_proj.weight ↔ ffn_X_exps.weight` rule. canon = canon - .replace(".mlp.experts.gate_up_proj.weight", ".ffn_gate_up_exps.weight") + .replace( + ".mlp.experts.gate_up_proj.weight", + ".ffn_gate_up_exps.weight", + ) .replace(".mlp.experts.down_proj.weight", ".ffn_down_exps.weight") .replace(".mlp.experts.gate_proj.weight", ".ffn_gate_exps.weight") .replace(".mlp.experts.up_proj.weight", ".ffn_up_exps.weight") @@ -3428,6 +4507,106 @@ fn to_canonical_name(name: &str, arch: &str) -> Option { base_arch::llama::map_llama_style(name).map(Canonical::Main) } +/// Candidate calibration-sidecar keys for a canonical bundle tensor name, +/// most-specific first. The baseRT collector records the RUNTIME dispatch +/// names (`dispatch_gemm`'s tensor_name), which differ from bundle-canonical +/// for several GLM tensors; routed experts additionally fall back to the +/// shared expert's stats — the shared expert consumes the IDENTICAL +/// activation vector (post ffn-norm), so its per-channel absmax is exactly +/// the right importance for the routed gate/up projections (and the closest +/// available proxy for down). +fn imatrix_stats_keys(canonical: &str) -> Vec { + let mut keys = vec![canonical.to_string()]; + if let Some(rest) = canonical.strip_prefix("layers.") { + if let Some((idx, tail)) = rest.split_once('.') { + let mapped: Option> = match tail { + "self_attn.o_proj.weight" => { + Some(vec![format!("layers.{idx}.attention.output.weight")]) + } + // Gate and up consume the IDENTICAL activation vector, and + // the fused gate|up decode dispatches only record the gate + // name — gate stats are exact for up, not a proxy. + "mlp.gate_proj.weight" => Some(vec![format!("layers.{idx}.ffn.gate.weight")]), + "mlp.up_proj.weight" => Some(vec![ + format!("layers.{idx}.ffn.up.weight"), + format!("layers.{idx}.ffn.gate.weight"), + ]), + "mlp.down_proj.weight" => Some(vec![format!("layers.{idx}.ffn.down.weight")]), + "mlp.shared_expert.gate_proj.weight" => { + Some(vec![format!("layers.{idx}.ffn_gate_shexp.weight")]) + } + "mlp.shared_expert.up_proj.weight" => Some(vec![ + format!("layers.{idx}.ffn_up_shexp.weight"), + format!("layers.{idx}.ffn_gate_shexp.weight"), + ]), + "mlp.shared_expert.down_proj.weight" => { + Some(vec![format!("layers.{idx}.ffn_down_shexp.weight")]) + } + // Routed experts aren't captured (bespoke MoE dispatches). + // Their gate/up input IS the shared expert's input (post + // ffn-norm), so the shexp stats are exact; shexp down is + // the closest available proxy for expert down. + "mlp.experts.gate_proj.weight" => Some(vec![ + format!("layers.{idx}.ffn_gate_exps.weight"), + format!("layers.{idx}.ffn_gate_shexp.weight"), + ]), + "mlp.experts.up_proj.weight" => Some(vec![ + format!("layers.{idx}.ffn_up_exps.weight"), + format!("layers.{idx}.ffn_up_shexp.weight"), + format!("layers.{idx}.ffn_gate_shexp.weight"), + ]), + "mlp.experts.down_proj.weight" => Some(vec![ + format!("layers.{idx}.ffn_down_exps.weight"), + format!("layers.{idx}.ffn_down_shexp.weight"), + ]), + _ => None, + }; + if let Some(m) = mapped { + keys.extend(m); + } + } + } + keys +} + +/// GLM 5.2 MLX/HF tensor names → the canonical names the GGUF +/// `GlmDsaMapper` emits (see `base-arch/src/glm.rs`), so GGUF- and +/// MLX-sourced `.base` bundles differ only where the checkpoints +/// genuinely differ. Input is the wrapper-stripped name +/// (`layers.N.…`, `lm_head.weight`, `norm.weight`). +/// +/// MLX-specific structure (from mlx-lm's `deepseek_v32.py`, which +/// `glm_moe_dsa` subclasses): +/// * `self_attn.embed_q` / `self_attn.unembed_out` are the per-head +/// nope/v split of the original `kv_b_proj` — exactly our +/// `k_b_proj` / `v_b_proj`, in the same memory layout +/// ([head, kv_lora, nope] row-major == GGUF ne0-fastest +/// [nope, kv_lora, head]; same for v_b). +/// * `indexer.wk` / `indexer.wq_b` ↔ GGUF `indexer.attn_k` / +/// `indexer.attn_q_b`; indexer weights exist only on the 21 "full" +/// layers (`indexer_types`) — the GGUF export carries junk copies +/// on all 79, so the MLX bundle legitimately has fewer tensors. +/// * `mlp.switch_mlp.*` is the stacked-expert block; `mlp.gate` is +/// the router; `mlp.shared_experts` (plural) is the single +/// ungated shared expert. +fn glm_hf_canonical(stripped: &str) -> Option { + if stripped == "norm.weight" { + return Some("final_norm.weight".to_string()); + } + Some( + stripped + .replace(".mlp.switch_mlp.", ".mlp.experts.") + .replace(".mlp.shared_experts.", ".mlp.shared_expert.") + .replace(".mlp.gate.weight", ".mlp.router.weight") + .replace(".self_attn.embed_q.", ".self_attn.k_b_proj.") + .replace(".self_attn.unembed_out.", ".self_attn.v_b_proj.") + .replace(".indexer.wk.", ".indexer.k_proj.") + .replace(".indexer.wq_b.", ".indexer.q_b_proj.") + .replace(".input_layernorm.", ".input_norm.") + .replace(".post_attention_layernorm.", ".post_attn_norm."), + ) +} + fn compute_sha256_streaming(path: &std::path::Path) -> Result { use sha2::{Digest, Sha256}; let f = std::fs::File::open(path)?; @@ -3447,17 +4626,14 @@ fn compute_sha256_streaming(path: &std::path::Path) -> Result { /// driven by profile.resolve; otherwise falls back to convert_synthetic /// for v1.0 behavior. Useful for end-to-end smoke testing the canonical /// pipeline without a real model checkpoint. -fn convert_synthetic_with_ctx( - output: &std::path::Path, - ctx: &QuantContext, -) -> Result<()> { +fn convert_synthetic_with_ctx(output: &std::path::Path, ctx: &QuantContext) -> Result<()> { if ctx.profile.is_none() { return convert_synthetic(output, ctx.target); } use base_format::{ - AlignmentConfig, BaseReader, BaseWriter, ComputeRegion, Header, HeaderFlags, LayerKind, - LayerDescriptor, LayerPrecision, ModelConfig, QuantScheme, SourceInfo, TargetBackend, - TensorFlags, TensorPayload, TokenizerBlob, + AlignmentConfig, BaseReader, BaseWriter, ComputeRegion, Header, HeaderFlags, + LayerDescriptor, LayerKind, LayerPrecision, ModelConfig, QuantScheme, SourceInfo, + TargetBackend, TensorFlags, TensorPayload, TokenizerBlob, }; use std::collections::BTreeMap; @@ -3495,10 +4671,7 @@ fn convert_synthetic_with_ctx( }, metadata: Default::default(), target_backend: TargetBackend::Metal, - quant_profile: ctx - .profile_name() - .unwrap_or("") - .to_string(), + quant_profile: ctx.profile_name().unwrap_or("").to_string(), alignment: AlignmentConfig::default(), flags: HeaderFlags::QUANTIZED | HeaderFlags::TIED_EMBEDDINGS, layers: (0..n_layers) @@ -3514,6 +4687,7 @@ fn convert_synthetic_with_ctx( tensors: vec![], mmproj: None, calibration: None, + provenance: None, sig: None, }; @@ -3588,7 +4762,9 @@ fn convert_synthetic_with_ctx( }; // Embedding (always bf16 per default profiles). - let embed: Vec = (0..vocab * hidden).map(|i| ((i as f32) % 17.0) * 0.01).collect(); + let embed: Vec = (0..vocab * hidden) + .map(|i| ((i as f32) % 17.0) * 0.01) + .collect(); emit( &mut writer, "model.embed_tokens.weight", @@ -3698,7 +4874,9 @@ fn convert_synthetic_with_ctx( ComputeRegion::Gpu, )?; - writer.finish().context("writing canonical synthetic bundle")?; + writer + .finish() + .context("writing canonical synthetic bundle")?; // Read it back; verify the canonical fields populated correctly. let reader = BaseReader::open(output).context("reopen canonical bundle")?; @@ -3717,9 +4895,9 @@ fn convert_synthetic_with_ctx( /// disk, and verify it reads back. fn convert_synthetic(output: &std::path::Path, target: TargetScheme) -> Result<()> { use base_format::{ - AlignmentConfig, BaseReader, BaseWriter, ComputeRegion, Header, HeaderFlags, LayerKind, - LayerDescriptor, LayerPrecision, ModelConfig, QuantScheme, SourceInfo, TargetBackend, TensorDtype, - TensorFlags, TensorPayload, TokenizerBlob, + AlignmentConfig, BaseReader, BaseWriter, ComputeRegion, Header, HeaderFlags, + LayerDescriptor, LayerKind, LayerPrecision, ModelConfig, QuantScheme, SourceInfo, + TargetBackend, TensorDtype, TensorFlags, TensorPayload, TokenizerBlob, }; use std::collections::BTreeMap; @@ -3781,13 +4959,16 @@ fn convert_synthetic(output: &std::path::Path, target: TargetScheme) -> Result<( tensors: vec![], mmproj: None, calibration: None, + provenance: None, sig: None, }; let mut writer = BaseWriter::create(output, header.clone()).context("create writer")?; // Embedding (GPU region, bf16). - let embed: Vec = (0..vocab * hidden).map(|i| ((i as f32) % 17.0) * 0.01).collect(); + let embed: Vec = (0..vocab * hidden) + .map(|i| ((i as f32) % 17.0) * 0.01) + .collect(); let embed_bytes: Vec = embed .iter() .flat_map(|&f| half::bf16::from_f32(f).to_le_bytes()) @@ -3937,6 +5118,54 @@ fn pack_for_target( } } +/// Row-aware sibling of [`pack_for_target`] for the profile-less path: the +/// runtime GEMV / GEMM / MoE kernels index `K / group_size` scales per row, +/// so a group may never span a row boundary. When the in-features dim is not +/// a multiple of the scheme's default group size, pick the largest smaller +/// group (down to 32) that divides it — the kernels read the group size from +/// the tensor header — and fall back to bf16 only when none does. Before +/// this, `--target base-q8` packed Nemotron-3-Nano's routed down experts +/// (K = 1856 = 14.5 x 128) with row-spanning groups: every scale was read +/// against the wrong weights and the bundle decoded garbage (PPL ~1e6). +fn pack_for_target_rows( + weights: &[f32], + target: TargetScheme, + in_features: Option, + name: &str, +) -> Result<(base_quant::Packed, base_format::TensorDtype)> { + use base_format::TensorDtype; + let (default_gs, dtype): (usize, TensorDtype) = match target { + TargetScheme::BaseQ4 => (base_quant::base_q4::GROUP_SIZE, TensorDtype::BaseQ4), + TargetScheme::BaseQ6 => (base_quant::base_q6::GROUP_SIZE, TensorDtype::BaseQ6), + TargetScheme::BaseQ8 => (base_quant::base_q8::GROUP_SIZE, TensorDtype::BaseQ8), + _ => return pack_for_target(weights, target), + }; + let k = match in_features { + Some(k) if k > 0 && k % default_gs != 0 => k, + _ => return pack_for_target(weights, target), + }; + let mut gs = default_gs; + while gs > 32 && k % gs != 0 { + gs /= 2; + } + if k % gs != 0 || weights.len() % gs != 0 { + eprintln!( + " note: {name} has in_features={k} (not a multiple of any {target:?} group size >= 32); \ + falling back to bf16 — quant grouping would misalign scales." + ); + return Ok((pack_bf16(weights), TensorDtype::Bf16)); + } + eprintln!( + " note: {name} has in_features={k} (not a multiple of gs={default_gs}); packing {target:?} at gs={gs}" + ); + let packed = match target { + TargetScheme::BaseQ4 => base_quant::base_q4::pack_with_group_size(weights, gs), + TargetScheme::BaseQ6 => base_quant::base_q6::pack_with_group_size(weights, gs), + _ => base_quant::base_q8::pack_with_group_size(weights, gs), + }; + Ok((packed, dtype)) +} + /// Wrap fp32 weights as bf16 raw bytes (no quant, no scales). fn pack_bf16(weights: &[f32]) -> base_quant::Packed { let bytes: Vec = weights @@ -3964,8 +5193,23 @@ struct QuantContext { target: TargetScheme, /// Bypass the spec's already-quantized-source rejection. allow_quant_from_quant: bool, + /// Transplant MLX affine-q4 tensors into `base_q4` verbatim instead + /// of requantizing them through f32. + q4_passthrough: bool, /// Copy GGUF Q4_K/Q5_K/Q6_K super-blocks through verbatim. kquant_passthrough: bool, + /// MLX sources only: reuse the source's packed q4 payloads + /// verbatim instead of dequant→requant (see `--mlx-passthrough`). + mlx_passthrough: bool, + /// After an `--mlx-passthrough` conversion, byte-compare the + /// written `.base` against the MLX source. + validate: bool, + /// Direct-write the blob behind a reserved header region (no + /// `.blobtmp`, no 2× disk peak) — see `--direct-write`. + direct_write: bool, + /// Importance-weighted RTN from the awq_profile sidecar — see + /// `--imatrix`. + imatrix: bool, } impl QuantContext { @@ -3989,13 +5233,34 @@ impl QuantContext { "--awq-profile requires --profile (AWQ only applies to canonical bit-widths from a profile)" ); } + if args.mlx_passthrough { + if profile.is_some() { + bail!("--mlx-passthrough and --profile are mutually exclusive (passthrough reuses the source's quant verbatim)"); + } + if !matches!(args.target, TargetScheme::BaseQ4) { + bail!("--mlx-passthrough requires --target base-q4 (the MLX 4-bit layout)"); + } + } + if args.validate && !args.mlx_passthrough { + bail!( + "--validate requires --mlx-passthrough (it byte-compares against the MLX source)" + ); + } + if args.imatrix && args.awq_profile.is_none() { + bail!("--imatrix requires --awq-profile (the calibration absmax supplies the channel weights)"); + } Ok(Self { profile, awq_profile, awq_config: base_awq::AwqConfig::default(), target: args.target, allow_quant_from_quant: args.allow_quant_from_quant, + q4_passthrough: !args.no_mlx_passthrough, kquant_passthrough: args.kquant_passthrough, + mlx_passthrough: args.mlx_passthrough, + validate: args.validate, + direct_write: args.direct_write, + imatrix: args.imatrix, }) } @@ -4021,7 +5286,7 @@ impl QuantContext { use base_format::TensorDtype; // Without a profile: legacy uniform-target behavior. let Some(profile) = &self.profile else { - return pack_for_target(weights, self.target); + return pack_for_target_rows(weights, self.target, in_features, name); }; let resolved = profile .resolve_or_err(name) @@ -4058,18 +5323,16 @@ impl QuantContext { TensorDtype::F32, )) } - TensorDtype::Mxfp4 => { - Ok((base_quant::mxfp4::pack(weights), TensorDtype::Mxfp4)) - } - TensorDtype::Nvfp4 => { - Ok((base_quant::nvfp4::pack(weights), TensorDtype::Nvfp4)) - } + TensorDtype::Mxfp4 => Ok((base_quant::mxfp4::pack(weights), TensorDtype::Mxfp4)), + TensorDtype::Nvfp4 => Ok((base_quant::nvfp4::pack(weights), TensorDtype::Nvfp4)), dtype @ (TensorDtype::BaseQ2 | TensorDtype::BaseQ3 | TensorDtype::BaseQ4 | TensorDtype::BaseQ5 | TensorDtype::BaseQ6 - | TensorDtype::BaseQ8) => self.pack_canonical(name, weights, in_features, dtype, resolved), + | TensorDtype::BaseQ8) => { + self.pack_canonical(name, weights, in_features, dtype, resolved) + } } } @@ -4122,24 +5385,64 @@ impl QuantContext { } } + // Importance-weighted RTN ("imatrix", --imatrix): fit each group's + // (scale, bias) under per-input-channel activation weights from the + // sidecar. Runtime-free (weights still approximate the originals) — + // takes precedence over the AWQ rotation path, which would require + // inference-side activation scaling the runtime doesn't implement. + if self.imatrix && !cfg.symmetric { + if let (Some(awq), Some(in_feat)) = (&self.awq_profile, in_features) { + if in_feat > 0 && in_feat % cfg.group_size as usize == 0 { + for key in imatrix_stats_keys(name) { + if let Some(absmax) = awq.absmax(&key) { + if absmax.len() == in_feat { + // Second moment proxy: importance = absmax². + let w: Vec = absmax.iter().map(|&a| a * a).collect(); + let packed = + base_quant::rtn::pack_weighted(weights, cfg, &w, in_feat); + return Ok((packed, dtype)); + } + eprintln!( + " imatrix: skipping {name} — stats {key:?} len {} != in_features {}", + absmax.len(), + in_feat + ); + break; + } + } + } + } + // No stats → plain RTN below (uniform importance). + } + // AWQ pre-process: only when sidecar carries an absmax for - // this tensor and we know the in_features dim. - let weights_for_pack: Vec = match (&self.awq_profile, in_features) { + // this tensor and we know the in_features dim. Never under + // --imatrix — rotation requires runtime activation scaling. + let weights_for_pack: Vec = match ( + if self.imatrix { + &None + } else { + &self.awq_profile + }, + in_features, + ) { (Some(awq), Some(in_feat)) => { if let Some(absmax) = awq.absmax(name) { if absmax.len() == in_feat { - let plan = self - .awq_config - .search(weights, in_feat, absmax, bits, cfg.group_size, cfg.symmetric); + let plan = self.awq_config.search( + weights, + in_feat, + absmax, + bits, + cfg.group_size, + cfg.symmetric, + ); // Rotate weights. The runtime undoes the rotation by // pre-multiplying activations with `plan.scales`; // those scales are stored alongside the rotated // tensor in the `.base` header so the runtime can // recover the original output. - eprintln!( - " awq: {name} α={:.2} mse={:.4e}", - plan.alpha, plan.mse - ); + eprintln!(" awq: {name} α={:.2} mse={:.4e}", plan.alpha, plan.mse); base_awq::awq_apply(weights, in_feat, &plan.scales) } else { eprintln!( @@ -4163,8 +5466,8 @@ impl QuantContext { } fn cmd_sign(args: SignArgs) -> Result<()> { - let key_bytes = std::fs::read(&args.key) - .with_context(|| format!("reading key {:?}", args.key))?; + let key_bytes = + std::fs::read(&args.key).with_context(|| format!("reading key {:?}", args.key))?; let key = base_sign::signing_key_from_bytes(&key_bytes)?; base_sign::sign_base_file(&args.input, &args.output, &key, &args.key_id)?; eprintln!("signed -> {}", args.output.display()); @@ -4173,8 +5476,8 @@ fn cmd_sign(args: SignArgs) -> Result<()> { fn cmd_verify(args: VerifyArgs) -> Result<()> { use ed25519_dalek::VerifyingKey; - let bytes = std::fs::read(&args.pubkey) - .with_context(|| format!("reading pubkey {:?}", args.pubkey))?; + let bytes = + std::fs::read(&args.pubkey).with_context(|| format!("reading pubkey {:?}", args.pubkey))?; if bytes.len() != 32 { bail!("ed25519 public key must be 32 bytes, got {}", bytes.len()); } @@ -4232,7 +5535,12 @@ fn cmd_inspect(args: InspectArgs) -> Result<()> { let slots = reader.slots()?; println!("n_slots: {}", slots.len()); for s in &slots { - println!(" slot kind={:?} raw=0x{:04x} len={}", s.kind(), s.kind_raw, s.payload.len()); + println!( + " slot kind={:?} raw=0x{:04x} len={}", + s.kind(), + s.kind_raw, + s.payload.len() + ); } if args.verify_checksums { @@ -4324,7 +5632,11 @@ mod canonical_name_tests { "layers.0.linear_attn.dt_bias", ), ] { - assert_eq!(main_canon(src, "qwen35"), Some(want.to_string()), "src={src}"); + assert_eq!( + main_canon(src, "qwen35"), + Some(want.to_string()), + "src={src}" + ); } // Full-attention block: standard Qwen QK-norm + gate passthrough, @@ -4396,7 +5708,10 @@ mod canonical_name_tests { "layers.0.mlp.gate_proj.weight", "layers.3.self_attn.q_proj.weight", ] { - assert!(!is_gdn_f32(n) && !is_gdn_f16(n), "unexpected override for {n}"); + assert!( + !is_gdn_f32(n) && !is_gdn_f16(n), + "unexpected override for {n}" + ); } } @@ -4441,6 +5756,76 @@ mod canonical_name_tests { ); } + // GLM 5.2 MLX (mlx-community/GLM-5.2-4bit) — the dedicated glm_dsa + // rename table must reproduce the GGUF mapper's canonical names. + #[test] + fn glm_dsa_mlx_tensor_canonicalization() { + let c = |n: &str| main_canon(n, "glm_dsa"); + // MoE: stacked experts, router, bias, shared expert (plural → singular). + assert_eq!( + c("model.layers.5.mlp.switch_mlp.down_proj.weight").as_deref(), + Some("layers.5.mlp.experts.down_proj.weight") + ); + assert_eq!( + c("model.layers.5.mlp.gate.weight").as_deref(), + Some("layers.5.mlp.router.weight") + ); + assert_eq!( + c("model.layers.5.mlp.gate.e_score_correction_bias").as_deref(), + Some("layers.5.mlp.gate.e_score_correction_bias") + ); + assert_eq!( + c("model.layers.5.mlp.shared_experts.up_proj.weight").as_deref(), + Some("layers.5.mlp.shared_expert.up_proj.weight") + ); + // MLA absorb pair: MLX pre-absorbed names → k_b/v_b. + assert_eq!( + c("model.layers.5.self_attn.embed_q.weight").as_deref(), + Some("layers.5.self_attn.k_b_proj.weight") + ); + assert_eq!( + c("model.layers.5.self_attn.unembed_out.weight").as_deref(), + Some("layers.5.self_attn.v_b_proj.weight") + ); + // Indexer (only the 21 "full" layers carry these). + assert_eq!( + c("model.layers.6.self_attn.indexer.wk.weight").as_deref(), + Some("layers.6.self_attn.indexer.k_proj.weight") + ); + assert_eq!( + c("model.layers.6.self_attn.indexer.wq_b.weight").as_deref(), + Some("layers.6.self_attn.indexer.q_b_proj.weight") + ); + assert_eq!( + c("model.layers.6.self_attn.indexer.k_norm.bias").as_deref(), + Some("layers.6.self_attn.indexer.k_norm.bias") + ); + assert_eq!( + c("model.layers.6.self_attn.indexer.weights_proj.weight").as_deref(), + Some("layers.6.self_attn.indexer.weights_proj.weight") + ); + // Norms + globals. + assert_eq!( + c("model.layers.5.input_layernorm.weight").as_deref(), + Some("layers.5.input_norm.weight") + ); + assert_eq!( + c("model.layers.5.post_attention_layernorm.weight").as_deref(), + Some("layers.5.post_attn_norm.weight") + ); + assert_eq!(c("model.norm.weight").as_deref(), Some("final_norm.weight")); + assert_eq!(c("lm_head.weight").as_deref(), Some("lm_head.weight")); + assert_eq!( + c("model.embed_tokens.weight").as_deref(), + Some("embed_tokens.weight") + ); + // Already-canonical MLA names pass through untouched. + assert_eq!( + c("model.layers.5.self_attn.kv_a_proj_with_mqa.weight").as_deref(), + Some("layers.5.self_attn.kv_a_proj_with_mqa.weight") + ); + } + // HF mainline Gemma 4 26B-A4B uses `model.language_model.layers.N.experts.X_proj` // bare (no `.weight` suffix; gate+up are 3D fused into a single tensor). #[test] @@ -4525,10 +5910,7 @@ mod canonical_name_tests { #[test] fn shared_expert_not_rewritten_by_experts_rules() { assert_eq!( - main_canon( - "model.layers.3.mlp.shared_expert.gate_proj.weight", - "qwen3", - ), + main_canon("model.layers.3.mlp.shared_expert.gate_proj.weight", "qwen3",), Some("layers.3.mlp.shared_expert.gate_proj.weight".to_string()), ); } @@ -4572,7 +5954,10 @@ mod canonical_name_tests { fn nomic_bert_hf_canonical_names() { let cases: &[(&str, &str)] = &[ ("embeddings.word_embeddings.weight", "embed_tokens.weight"), - ("embeddings.token_type_embeddings.weight", "token_types.weight"), + ( + "embeddings.token_type_embeddings.weight", + "token_types.weight", + ), ("emb_ln.weight", "token_embd_norm.weight"), ("emb_ln.bias", "token_embd_norm.bias"), ( @@ -4624,7 +6009,10 @@ mod canonical_name_tests { main_canon("embeddings.position_embeddings.weight", "nomic-bert"), None, ); - assert_eq!(main_canon("0.auto_model.pooler.dense.weight", "nomic-bert"), None); + assert_eq!( + main_canon("0.auto_model.pooler.dense.weight", "nomic-bert"), + None + ); } } @@ -4685,7 +6073,9 @@ mod stacking_tests { // rewritten names: per-expert dropped, attn kept, virtuals added let rw = sp.rewritten_names(); - assert!(rw.iter().any(|n| n == "model.layers.0.self_attn.q_proj.weight")); + assert!(rw + .iter() + .any(|n| n == "model.layers.0.self_attn.q_proj.weight")); assert!(rw.iter().any(|n| n == gate)); assert!(!rw.iter().any(|n| n.contains(".mlp.experts.0."))); } @@ -4773,7 +6163,10 @@ mod splitting_tests { assert_eq!(sp.source_shape(u).unwrap(), vec![3, 2]); // qkv layout = [q rows 0..4 | k rows 4..6 | v rows 6..8] - assert_eq!(sp.to_f32(q).unwrap(), vec![0., 1., 10., 11., 20., 21., 30., 31.]); + assert_eq!( + sp.to_f32(q).unwrap(), + vec![0., 1., 10., 11., 20., 21., 30., 31.] + ); assert_eq!(sp.to_f32(k).unwrap(), vec![40., 41., 50., 51.]); assert_eq!(sp.to_f32(v).unwrap(), vec![60., 61., 70., 71.]); // gate_up layout = [gate rows 0..3 | up rows 3..6], gate first @@ -4782,7 +6175,9 @@ mod splitting_tests { // rewritten: fused dropped, o_proj passed through, virtuals present let rw = sp.rewritten_names(); - assert!(rw.iter().any(|n| n == "model.layers.0.self_attn.o_proj.weight")); + assert!(rw + .iter() + .any(|n| n == "model.layers.0.self_attn.o_proj.weight")); assert!(rw.iter().any(|n| n == q)); assert!(!rw.iter().any(|n| n.ends_with("qkv_proj.weight"))); assert!(!rw.iter().any(|n| n.ends_with("gate_up_proj.weight"))); @@ -4803,7 +6198,7 @@ mod splitting_tests { #[test] fn mismatched_config_is_rejected() { let mock = SplitMock; // qkv_proj is [8, 2] - // nq=3 implies q+k+v = 3*2 + 2*(1*2) = 10 ≠ 8 rows → loud error, no garbage. + // nq=3 implies q+k+v = 3*2 + 2*(1*2) = 10 ≠ 8 rows → loud error, no garbage. let bad = ArchConfig { head_dim: 2, num_attention_heads: 3, @@ -4831,7 +6226,8 @@ fn rope_permute_rows(f32s: &[f32], rows: usize, cols: usize, n_heads: u32) -> Ve for k in 0..2usize { let dst = h * hd + 2 * j + k; let src = h * hd + k * half + j; - out[dst * cols..(dst + 1) * cols].copy_from_slice(&f32s[src * cols..(src + 1) * cols]); + out[dst * cols..(dst + 1) * cols] + .copy_from_slice(&f32s[src * cols..(src + 1) * cols]); } } } @@ -4914,7 +6310,10 @@ mod gguf_passthrough_tests { let original: Vec = (0..rows * cols).map(|i| i as f32).collect(); // What llama.cpp would have written into the GGUF. let permuted = rope_permute_rows(&original, rows, cols, n_heads); - assert_ne!(permuted, original, "forward permute must actually move rows"); + assert_ne!( + permuted, original, + "forward permute must actually move rows" + ); let out = unpermute_rope_rows( &info(cols as u64, rows as u64, GgmlType::F32), @@ -4957,12 +6356,12 @@ mod gguf_passthrough_tests { let (n_heads, hd) = (2u32, 4usize); let rows = n_heads as usize * hd; let row_bytes = 2 * 144; // 512 elements / 256 per block x 144 B - // Give every row a distinct byte pattern so a mis-shuffle shows up. + // Give every row a distinct byte pattern so a mis-shuffle shows up. let src: Vec = (0..rows) .flat_map(|r| std::iter::repeat_n((r as u8) + 1, row_bytes)) .collect(); - let out = unpermute_rope_rows(&info(512, rows as u64, GgmlType::Q4K), &src, n_heads) - .unwrap(); + let out = + unpermute_rope_rows(&info(512, rows as u64, GgmlType::Q4K), &src, n_heads).unwrap(); assert_eq!(out.len(), src.len()); // Expected destination rows: dst[h*HD + k*HD/2 + j] = src[h*HD + 2j + k]. // Head 0 src rows 0,1,2,3 → dst 0,2,1,3. diff --git a/base-convert/crates/base-convert/tests/hub_e2e.rs b/base-convert/crates/base-convert/tests/hub_e2e.rs index b2e3a36..a119d2b 100644 --- a/base-convert/crates/base-convert/tests/hub_e2e.rs +++ b/base-convert/crates/base-convert/tests/hub_e2e.rs @@ -51,3 +51,154 @@ fn list_discovers_synthetic_model() { assert!(table.contains("basecompute/demo"), "table: {table}"); assert!(table.contains("installed"), "table: {table}"); } + +#[cfg(unix)] +#[test] +fn computearena_dispatches_to_standalone_cli_with_basert_adapter() { + use std::os::unix::fs::PermissionsExt; + + let tmp = tempfile::tempdir().unwrap(); + let executable = tmp.path().join("computearena"); + std::fs::write(&executable, "#!/bin/sh\nprintf '%s\\n' \"$@\"\n").unwrap(); + let mut permissions = std::fs::metadata(&executable).unwrap().permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&executable, permissions).unwrap(); + + let output = Command::new(bin()) + .args(["computearena", "--harness", "/tmp/harness", "run"]) + .env("PATH", tmp.path()) + .output() + .expect("run basert computearena"); + assert!( + output.status.success(), + "stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert_eq!( + String::from_utf8(output.stdout).unwrap(), + "basert\n--harness\n/tmp/harness\nrun\n" + ); +} + +#[cfg(unix)] +#[test] +fn computearena_dispatch_exposes_a_bundled_harness_without_overriding_the_user() { + use std::os::unix::fs::PermissionsExt; + + let tmp = tempfile::tempdir().unwrap(); + let launcher = tmp.path().join("basert"); + std::fs::copy(bin(), &launcher).unwrap(); + + let computearena = tmp.path().join("computearena"); + std::fs::write( + &computearena, + "#!/bin/sh\nprintf '%s\\n' \"$@\"\nprintf 'harness=%s\\n' \"$COMPUTEARENA_BASERT_HARNESS\"\n", + ) + .unwrap(); + let mut permissions = std::fs::metadata(&computearena).unwrap().permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&computearena, permissions).unwrap(); + + let harness = tmp.path().join("basert-benchmark-harness"); + std::fs::write(&harness, b"fixture").unwrap(); + + let output = Command::new(&launcher) + .args(["computearena", "list"]) + .env("PATH", tmp.path()) + .env_remove("COMPUTEARENA_BASERT_HARNESS") + .env_remove("BASERT_COMPUTEARENA_HARNESS") + .output() + .expect("run bundled basert computearena"); + assert!( + output.status.success(), + "stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert_eq!( + String::from_utf8(output.stdout).unwrap(), + format!("basert\nlist\nharness={}\n", harness.display()) + ); + + let output = Command::new(&launcher) + .args(["computearena", "list"]) + .env("PATH", tmp.path()) + .env("COMPUTEARENA_BASERT_HARNESS", "/user/selected/harness") + .output() + .expect("run basert computearena with an explicit harness environment"); + assert!(output.status.success()); + assert_eq!( + String::from_utf8(output.stdout).unwrap(), + "basert\nlist\nharness=/user/selected/harness\n" + ); +} + +#[cfg(unix)] +#[test] +fn computearena_dispatch_exposes_the_source_build_harness() { + use std::os::unix::fs::PermissionsExt; + + let tmp = tempfile::tempdir().unwrap(); + let repo = tmp.path().join("baseRT"); + let launcher_dir = repo.join("tools/base-convert/target/release"); + std::fs::create_dir_all(&launcher_dir).unwrap(); + std::fs::write(repo.join("tools/base-convert/Cargo.toml"), b"[workspace]\n").unwrap(); + + let launcher = launcher_dir.join("basert"); + std::fs::copy(bin(), &launcher).unwrap(); + + let cli_dir = tmp.path().join("cli"); + std::fs::create_dir_all(&cli_dir).unwrap(); + let computearena = cli_dir.join("computearena"); + std::fs::write( + &computearena, + "#!/bin/sh\nprintf '%s\\n' \"$@\"\nprintf 'harness=%s\\n' \"$COMPUTEARENA_BASERT_HARNESS\"\n", + ) + .unwrap(); + let mut permissions = std::fs::metadata(&computearena).unwrap().permissions(); + permissions.set_mode(0o755); + std::fs::set_permissions(&computearena, permissions).unwrap(); + + let harness = repo.join("build/basert-benchmark-harness"); + std::fs::create_dir_all(harness.parent().unwrap()).unwrap(); + std::fs::write(&harness, b"fixture").unwrap(); + + let output = Command::new(&launcher) + .args(["computearena", "list"]) + .env("PATH", &cli_dir) + .env_remove("COMPUTEARENA_BASERT_HARNESS") + .env_remove("BASERT_COMPUTEARENA_HARNESS") + .output() + .expect("run source-built basert computearena"); + assert!( + output.status.success(), + "stderr: {}", + String::from_utf8_lossy(&output.stderr) + ); + assert_eq!( + String::from_utf8(output.stdout).unwrap(), + format!("basert\nlist\nharness={}\n", harness.display()) + ); +} + +#[test] +fn missing_computearena_points_to_the_public_quickstart() { + let tmp = tempfile::tempdir().unwrap(); + let output = Command::new(bin()) + .arg("computearena") + .env("PATH", tmp.path()) + .output() + .expect("run basert without computearena installed"); + + assert!(!output.status.success()); + let stderr = String::from_utf8(output.stderr).unwrap(); + assert!( + stderr.contains("ComputeArena is not installed."), + "{stderr}" + ); + assert!( + stderr.contains("https://computearena.ai/quickstart"), + "{stderr}" + ); + assert!(stderr.contains("`basert computearena` again"), "{stderr}"); + assert!(!stderr.contains("--harness"), "{stderr}"); +} diff --git a/base-convert/crates/base-convert/tests/inventory.rs b/base-convert/crates/base-convert/tests/inventory.rs index 7c35b29..95a523e 100644 --- a/base-convert/crates/base-convert/tests/inventory.rs +++ b/base-convert/crates/base-convert/tests/inventory.rs @@ -18,9 +18,10 @@ use std::process::Command; /// Known-not-yet-supported GGUF variants. Each entry is a substring /// match against the filename. Maintaining this list documents the /// coverage frontier. -const EXPECTED_SKIP: &[(&str, &str)] = &[ - ("mmproj", "standalone mmproj bundles handled as sub-bundles, not primary"), -]; +const EXPECTED_SKIP: &[(&str, &str)] = &[( + "mmproj", + "standalone mmproj bundles handled as sub-bundles, not primary", +)]; fn models_dir() -> Option { // BASE_MODELS_DIR is the explicit override; otherwise look for a diff --git a/base-convert/crates/base-convert/tests/whisper_quant.rs b/base-convert/crates/base-convert/tests/whisper_quant.rs index e3c1916..3c9e736 100644 --- a/base-convert/crates/base-convert/tests/whisper_quant.rs +++ b/base-convert/crates/base-convert/tests/whisper_quant.rs @@ -122,13 +122,25 @@ fn write_synthetic_whisper_dir(dir: &Path) { let mut t: Vec<(String, Vec)> = vec![ ("model.encoder.conv1.weight".into(), vec![D_MODEL, MELS, 3]), ("model.encoder.conv1.bias".into(), vec![D_MODEL]), - ("model.encoder.conv2.weight".into(), vec![D_MODEL, D_MODEL, 3]), + ( + "model.encoder.conv2.weight".into(), + vec![D_MODEL, D_MODEL, 3], + ), ("model.encoder.conv2.bias".into(), vec![D_MODEL]), - ("model.encoder.embed_positions.weight".into(), vec![SRC_POS, D_MODEL]), + ( + "model.encoder.embed_positions.weight".into(), + vec![SRC_POS, D_MODEL], + ), ("model.encoder.layer_norm.weight".into(), vec![D_MODEL]), ("model.encoder.layer_norm.bias".into(), vec![D_MODEL]), - ("model.decoder.embed_tokens.weight".into(), vec![VOCAB, D_MODEL]), - ("model.decoder.embed_positions.weight".into(), vec![TGT_POS, D_MODEL]), + ( + "model.decoder.embed_tokens.weight".into(), + vec![VOCAB, D_MODEL], + ), + ( + "model.decoder.embed_positions.weight".into(), + vec![TGT_POS, D_MODEL], + ), ("model.decoder.layer_norm.weight".into(), vec![D_MODEL]), ("model.decoder.layer_norm.bias".into(), vec![D_MODEL]), ]; @@ -214,8 +226,16 @@ fn assert_quant_bundle( assert_eq!(h.arch, "whisper"); assert_eq!(h.quant_scheme, scheme); assert_eq!(h.quant_profile, profile_name); - assert!(h.flags.contains(HeaderFlags::QUANTIZED), "flags: {:?}", h.flags); - assert!(h.flags.contains(HeaderFlags::TIED_EMBEDDINGS), "flags: {:?}", h.flags); + assert!( + h.flags.contains(HeaderFlags::QUANTIZED), + "flags: {:?}", + h.flags + ); + assert!( + h.flags.contains(HeaderFlags::TIED_EMBEDDINGS), + "flags: {:?}", + h.flags + ); let quant_names = expected_quant_linears(); let mut seen_quant = 0usize; @@ -306,7 +326,11 @@ fn whisper_default_stays_all_f16() { assert_eq!(h.arch, "whisper"); assert_eq!(h.quant_scheme, QuantScheme::F16); assert_eq!(h.quant_profile, ""); - assert!(!h.flags.contains(HeaderFlags::QUANTIZED), "flags: {:?}", h.flags); + assert!( + !h.flags.contains(HeaderFlags::QUANTIZED), + "flags: {:?}", + h.flags + ); assert!(h.flags.contains(HeaderFlags::TIED_EMBEDDINGS)); for t in h.tensors.iter() { let numel: u64 = t.shape.iter().product(); diff --git a/base-convert/crates/base-format/src/header.rs b/base-convert/crates/base-format/src/header.rs index 484a121..befca4b 100644 --- a/base-convert/crates/base-format/src/header.rs +++ b/base-convert/crates/base-format/src/header.rs @@ -273,7 +273,6 @@ pub enum ComputeRegion { Cpu, } - /// Per-region alignment in log2 bytes. Stored in the header so the runtime /// knows the packed layout without hardcoding assumptions, and so the /// converter can target different hardware page sizes. @@ -505,6 +504,71 @@ pub struct Signature { pub signature: String, } +/// Conversion provenance: where every bundle tensor came from and what +/// was done to it, written by the converter at conversion time. The +/// runtime ignores this block entirely; it exists so verification +/// tooling can check a bundle against its source checkpoint without a +/// hand-maintained per-arch name table, and can prove coverage in both +/// directions (every bundle tensor accounted for, every source tensor +/// consumed). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct Provenance { + pub schema: u32, + /// True when the converter ran the mirror policy on a quantized + /// source: tensors the source stores quantized are transplanted, + /// tensors it keeps unquantized are carried unquantized. + #[serde(default, skip_serializing_if = "is_false")] + pub mirror: bool, + /// Source tensors intentionally not represented in the bundle. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub dropped: Vec, + /// Source tensors routed into the mmproj sub-bundle. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub mmproj: Vec, + /// Bundle tensor name -> provenance record. + pub tensors: BTreeMap, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct TensorProvenance { + /// Source tensor name(s) this bundle tensor was built from. Empty + /// only when `stack` describes the sources instead. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub src: Vec, + /// Packed bytes copied verbatim from the source's identical scheme. + #[serde(default, skip_serializing_if = "is_false")] + pub transplanted: bool, + /// Stored unquantized because the source stores it unquantized + /// (mirror policy). + #[serde(default, skip_serializing_if = "is_false")] + pub carried: bool, + /// Reparameterization baked into the stored values (e.g. "neg_exp" + /// for Mamba-2 `A = -exp(A_log)`). Verifiers re-derive the expected + /// values from this label independently. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub transform: Option, + /// Additive shift baked into 1-D norm gains (Gemma-style +1). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub norm_shift: Option, + /// Rows were reordered (HF split-half rotary -> interleaved). + #[serde(default, skip_serializing_if = "is_false")] + pub permuted: bool, + /// `[row_offset, row_count]` slice of the source tensor (fused + /// qkv_proj / gate_up_proj splits). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub rows: Option<[u64; 2]>, + /// Stacked from `count` per-expert source tensors; `pattern` + /// contains `{e}` where the expert index goes. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub stack: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StackRef { + pub pattern: String, + pub count: u32, +} + /// Top-level header. Serialized as JSON with sorted keys for /// reproducible signing. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -565,6 +629,24 @@ pub struct Header { pub calibration: Option, #[serde(skip_serializing_if = "Option::is_none", default)] pub sig: Option, + + /// Conversion provenance: how every bundle tensor relates to the + /// source checkpoint's tensors. Written by the converter so a + /// fidelity gate can verify the bundle against the checkpoint without + /// a hand-maintained name table. Shape: + /// + /// ```json + /// {"tensors": {"": {"src": ["", ...], + /// "transplanted": true, // packed codes copied verbatim + /// "transform": "neg_exp", // declared reparameterization + /// "stack": {"pattern": "...{e}...", "count": N}, + /// "permuted": true, "rows": [off, cnt]}}, + /// "dropped": [""]} + /// ``` + /// + /// Absent on bundles from converter paths that do not record it yet. + #[serde(skip_serializing_if = "Option::is_none", default)] + pub provenance: Option, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -584,7 +666,7 @@ pub struct MmprojBundle { /// Per-layer kind. Drives runtime dispatch: which forward path to run /// (attention vs SSM), which KV / SSM buffers to allocate, and whether /// the layer feeds into an MoE FFN. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] #[serde(rename_all = "snake_case")] pub enum LayerKind { /// Standard dense multi-head attention (no GQA). @@ -604,6 +686,37 @@ pub enum LayerKind { DenseMoe, /// Standard transformer block (attention + dense MLP). DenseMlp, + /// FFN-only MoE block: no attention/SSM mixer in the layer. + /// Nemotron-H MoE interleaves pure-Mamba, pure-attention and + /// pure-MoE-FFN blocks — each block has exactly one of the three. + MoeFfn, + /// Forward-compat fallback: a layer kind this build does not know. + /// Older tools must not hard-fail on a newer bundle's header just to + /// print an inventory (the C++ runtime never validated this enum, + /// and an old `basert inspect` once refused a Nemotron bundle over + /// exactly this). Never written by a converter. + Unknown, +} + +// Deserialize by hand so unknown strings land on `Unknown` instead of a +// serde hard error (`#[serde(other)]` is not available on externally +// tagged unit enums). +impl<'de> serde::Deserialize<'de> for LayerKind { + fn deserialize>(d: D) -> Result { + let s = String::deserialize(d)?; + Ok(match s.as_str() { + "attention_dense" => Self::AttentionDense, + "attention_gqa" => Self::AttentionGqa, + "attention_sliding" => Self::AttentionSliding, + "ssm" => Self::Ssm, + "ssm_moe" => Self::SsmMoe, + "attention_moe" => Self::AttentionMoe, + "dense_moe" => Self::DenseMoe, + "dense_mlp" => Self::DenseMlp, + "moe_ffn" => Self::MoeFfn, + _ => Self::Unknown, + }) + } } /// Per-layer precision overrides. Absent = inherit bundle default. diff --git a/base-convert/crates/base-format/src/lib.rs b/base-convert/crates/base-format/src/lib.rs index 4087e21..00e1db0 100644 --- a/base-convert/crates/base-format/src/lib.rs +++ b/base-convert/crates/base-format/src/lib.rs @@ -11,8 +11,9 @@ mod writer; pub use error::{Error, Result}; pub use header::{ AlignmentConfig, CalibrationInfo, ComputeRegion, Header, HeaderFlags, LayerDescriptor, - LayerKind, LayerPrecision, Layout, ModelConfig, QuantScheme, ResidencyHint, ScaleDtype, - Signature, SourceInfo, TargetBackend, TensorDtype, TensorEntry, TensorFlags, TokenizerBlob, + LayerKind, LayerPrecision, Layout, ModelConfig, Provenance, QuantScheme, ResidencyHint, + ScaleDtype, Signature, SourceInfo, StackRef, TargetBackend, TensorDtype, TensorEntry, + TensorFlags, TensorProvenance, TokenizerBlob, }; pub use reader::BaseReader; pub use slots::{read_slots, write_slots, Slot, SlotFlags, SlotKind}; diff --git a/base-convert/crates/base-format/src/reader.rs b/base-convert/crates/base-format/src/reader.rs index 6d42546..d8bd623 100644 --- a/base-convert/crates/base-format/src/reader.rs +++ b/base-convert/crates/base-format/src/reader.rs @@ -103,10 +103,11 @@ impl BaseReader { let header_len = u64::from_le_bytes(prefix[8..16].try_into().unwrap()); let mut header_bytes = vec![0u8; header_len as usize]; - file.read_exact(&mut header_bytes).map_err(|e| match e.kind() { - std::io::ErrorKind::UnexpectedEof => Error::HeaderOverflow(header_len), - _ => Error::Io(e), - })?; + file.read_exact(&mut header_bytes) + .map_err(|e| match e.kind() { + std::io::ErrorKind::UnexpectedEof => Error::HeaderOverflow(header_len), + _ => Error::Io(e), + })?; Ok(Header::from_json_bytes(&header_bytes)?) } @@ -216,10 +217,7 @@ impl BaseReader { .iter() .flat_map(|m| m.tensors.iter()) .map(|t| self.blob_offset + t.offset + t.length); - let blob_end = main_end - .chain(mmproj_end) - .max() - .unwrap_or(self.blob_offset); + let blob_end = main_end.chain(mmproj_end).max().unwrap_or(self.blob_offset); (blob_end + 7) & !7u64 } diff --git a/base-convert/crates/base-format/src/writer.rs b/base-convert/crates/base-format/src/writer.rs index 6767383..c23220b 100644 --- a/base-convert/crates/base-format/src/writer.rs +++ b/base-convert/crates/base-format/src/writer.rs @@ -2,8 +2,8 @@ use crate::header::{Header, MmprojBundle, TensorDtype, TensorEntry}; use crate::slots::{write_slots, Slot}; use crate::{Error, Result, BLOB_ALIGNMENT, FORMAT_VERSION, MAGIC, PREFIX_LEN}; use std::fs::File; -use std::io::{BufWriter, Seek, Write}; -use std::path::Path; +use std::io::{BufReader, BufWriter, Seek, SeekFrom, Write}; +use std::path::{Path, PathBuf}; /// A tensor payload staged for writing. /// @@ -14,6 +14,24 @@ pub struct TensorPayload { pub data: Vec, } +/// Streaming blob sink: the weights blob is written incrementally to a +/// sibling temp file as tensors are added, so the writer never holds the +/// whole (multi-hundred-GB) blob in RAM. Only the tiny per-tensor +/// `TensorEntry` metadata is retained. At `finish` the header is +/// serialized (all offsets/lengths/checksums are already known) and the +/// temp blob is streamed into the final file after it. The on-disk format +/// is byte-identical to the buffered path. +struct BlobStream { + file: BufWriter, + path: PathBuf, + /// Bytes written to the blob so far (blob-relative cursor). Tensor + /// `entry.offset` values are relative to the blob start, matching the + /// reader's `blob_offset + entry.offset` addressing. + cursor: u64, + lm_entries: Vec, + mmproj_entries: Vec, +} + /// Writer for the `.base` single-file format. /// /// Usage: @@ -22,9 +40,16 @@ pub struct TensorPayload { /// 3. Call `finish` to commit — partitions tensors by `compute_region`, /// assigns per-region alignments, serializes the canonical-JSON header, /// writes prefix + padding + blob. +/// +/// File-backed writers (via [`BaseWriter::create`]) STREAM the blob to a +/// temp file (constant memory). Generic writers (via [`BaseWriter::new`], +/// e.g. an in-memory `Cursor` in tests) buffer payloads in RAM. pub struct BaseWriter { inner: W, header: Header, + /// Streaming blob sink (Some for file-backed `create`, None for the + /// buffered generic path). + stream: Option, payloads: Vec, /// Tensors destined for the multimodal sub-bundle (vision/audio /// towers + projector). Stored in the same weights blob as the LM @@ -40,19 +65,100 @@ pub struct BaseWriter { /// separate config file. mmproj_config: std::collections::BTreeMap, slots: Vec, + /// First error encountered while streaming a payload to the temp blob + /// (add_tensor is infallible for API compatibility); reported at finish. + stream_error: Option, + /// Direct-write mode (`create_direct`): the blob streams straight into + /// the final file after a fixed reserved header region, so no `.blobtmp` + /// sibling (and no 2× disk peak) exists. Holds the reserved header byte + /// count; the header JSON is space-padded to exactly this length at + /// finish so the reader's `align_up(prefix + header_len)` blob-start + /// computation lands on the offset the blob was streamed at. + direct_reserve: Option, } impl BaseWriter> { pub fn create>(path: P, header: Header) -> Result { + let path = path.as_ref(); let file = File::create(path)?; + // Sibling temp file for the streamed blob (same directory ⇒ same + // filesystem, so the final copy is a fast local sequential I/O). + let tmp_path = { + let mut s = path.as_os_str().to_os_string(); + s.push(".blobtmp"); + PathBuf::from(s) + }; + // Read+write: we stream the blob in, then read it back at finish to + // copy it into the final file (File::create is write-only → EBADF on read). + let tmp = std::fs::OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open(&tmp_path)?; + Ok(Self { + inner: BufWriter::new(file), + header, + stream: Some(BlobStream { + file: BufWriter::new(tmp), + path: tmp_path, + cursor: 0, + lm_entries: Vec::new(), + mmproj_entries: Vec::new(), + }), + payloads: Vec::new(), + mmproj_payloads: Vec::new(), + mmproj_arch: None, + mmproj_config: std::collections::BTreeMap::new(), + slots: Vec::new(), + stream_error: None, + direct_reserve: None, + }) + } + + /// Direct-write variant of [`BaseWriter::create`]: reserve + /// `header_reserve` bytes for the header up front and stream the blob + /// straight into the final file at the fixed blob start — no `.blobtmp` + /// sibling, so peak disk usage is the bundle size instead of 2×. The + /// header JSON is space-padded (valid trailing whitespace) to exactly + /// the reserve at finish, which keeps the reader's + /// `align_up(prefix + header_len)` blob-start computation equal to the + /// offset the blob was streamed at; finish errors if the header doesn't + /// fit (re-run with a larger reserve). Intended for very large bundles + /// where the temp-blob copy would not fit on disk. + pub fn create_direct>( + path: P, + header: Header, + header_reserve: u64, + ) -> Result { + let path = path.as_ref(); + let file = File::create(path)?; // header handle (truncates) + // Second handle for the blob stream, positioned at the fixed blob + // start. Writing beyond EOF leaves the header region as a hole + // until finish backfills it. + let mut blob = std::fs::OpenOptions::new() + .read(true) + .write(true) + .open(path)?; + let blob_start = align_up(PREFIX_LEN + header_reserve, BLOB_ALIGNMENT); + blob.seek(SeekFrom::Start(blob_start))?; Ok(Self { inner: BufWriter::new(file), header, + stream: Some(BlobStream { + file: BufWriter::new(blob), + path: path.to_path_buf(), + cursor: 0, + lm_entries: Vec::new(), + mmproj_entries: Vec::new(), + }), payloads: Vec::new(), mmproj_payloads: Vec::new(), mmproj_arch: None, mmproj_config: std::collections::BTreeMap::new(), slots: Vec::new(), + stream_error: None, + direct_reserve: Some(header_reserve), }) } } @@ -62,16 +168,61 @@ impl BaseWriter { Self { inner, header, + stream: None, payloads: Vec::new(), mmproj_payloads: Vec::new(), mmproj_arch: None, mmproj_config: std::collections::BTreeMap::new(), slots: Vec::new(), + stream_error: None, + direct_reserve: None, } } + /// Stream one payload into the temp blob, assigning its aligned + /// blob-relative offset and computing its checksum. The payload's + /// `data` is dropped immediately after the write, keeping memory flat. + fn stream_payload( + alignment: crate::header::AlignmentConfig, + s: &mut BlobStream, + payload: TensorPayload, + is_mmproj: bool, + ) -> Result<()> { + let align = alignment.align_for(payload.entry.compute_region); + let aligned = align_up(s.cursor, align); + if aligned > s.cursor { + write_zeros(&mut s.file, (aligned - s.cursor) as usize)?; + } + let mut entry = payload.entry; + entry.offset = aligned; + entry.length = payload.data.len() as u64; + if entry.checksum_xxh64.is_none() { + entry.checksum_xxh64 = Some(xxhash_rust::xxh64::xxh64(&payload.data, 0)); + } + s.file.write_all(&payload.data)?; + s.cursor = aligned + payload.data.len() as u64; + if is_mmproj { + s.mmproj_entries.push(entry); + } else { + s.lm_entries.push(entry); + } + Ok(()) + } + pub fn add_tensor(&mut self, payload: TensorPayload) { - self.payloads.push(payload); + if let Some(s) = self.stream.as_mut() { + // Streaming path: file-backed writers never buffer the blob. + // (add_tensor's signature is infallible for API compatibility; + // a write error here is surfaced by re-checking at finish via + // the BufWriter's retained error state on flush.) + let alignment = self.header.alignment; + if let Err(e) = Self::stream_payload(alignment, s, payload, false) { + // Stash the error to report at finish; keep the API simple. + self.stream_error.get_or_insert(e); + } + } else { + self.payloads.push(payload); + } } /// Add a tensor that belongs to the multimodal sub-bundle @@ -79,7 +230,17 @@ impl BaseWriter { /// into the same weights blob, but its entry lands under /// `header.mmproj.tensors` instead of `header.tensors`. pub fn add_mmproj_tensor(&mut self, payload: TensorPayload) { - self.mmproj_payloads.push(payload); + if let Some(s) = self.stream.as_mut() { + // mmproj tensors continue the blob after all LM tensors; callers + // add every LM tensor before the first mmproj tensor, so the + // streamed order matches the buffered path. + let alignment = self.header.alignment; + if let Err(e) = Self::stream_payload(alignment, s, payload, true) { + self.stream_error.get_or_insert(e); + } + } else { + self.mmproj_payloads.push(payload); + } } /// Set the mmproj sub-bundle arch tag (e.g. "gemma4_vision_audio"). @@ -104,7 +265,48 @@ impl BaseWriter { self.slots.push(slot); } + /// Attach conversion provenance to the header. Call before `finish()`; + /// accepts both the structured converter record and JSON-producing + /// architecture-specific paths. + pub fn set_provenance(&mut self, prov: T) { + self.header.provenance = Some(serde_json::to_value(prov).expect("serializing provenance")); + } + + /// Stamp `target_backend` from CONTENT instead of trusting the caller's + /// default (every construction site hardcodes Metal, so CUDA-only bundles + /// carried a `metal` tag and failed a kernel lookup only after a full + /// download + load). The one CUDA-only content class today: base_q6 MoE + /// expert slabs (`*_exps.*` tensors) — Metal ships no q6 MoE kernels. + /// bf16-scale q8/q4 bundles stay `metal` (universal): both backends carry + /// the `_sbf16` kernel families. + /// + /// Call from EVERY finish path. It used to live inline in the buffered + /// `finish` only, which the streaming rewrite silently turned into dead + /// code for real bundles: `create`/`create_direct` both set `stream`, so + /// they return through `finish_streaming`/`finish_direct` and never + /// reached it. Requires `header.tensors` to already be populated. + fn stamp_target_backend_from_content(&mut self) { + let q6_experts = |ts: &[TensorEntry]| { + ts.iter() + .any(|t| t.name.contains("_exps.") && t.dtype == TensorDtype::BaseQ6) + }; + if q6_experts(&self.header.tensors) { + self.header.target_backend = crate::header::TargetBackend::CudaSm121; + } + } + pub fn finish(mut self) -> Result<()> { + // Surface any error stashed while streaming payloads to the temp blob. + if let Some(e) = self.stream_error.take() { + return Err(e); + } + if let Some(reserve) = self.direct_reserve { + return self.finish_direct(reserve); + } + if self.stream.is_some() { + return self.finish_streaming(); + } + let alignment = self.header.alignment; // Assign each tensor an offset honoring its compute-region's @@ -127,22 +329,7 @@ impl BaseWriter { } self.header.tensors = entries; - // Stamp `target_backend` from CONTENT instead of trusting the - // caller's default (every construction site used to hardcode - // Metal, so CUDA-only bundles carried a `metal` tag and failed a - // kernel lookup only after a full download + load). The one - // CUDA-only content class today: base_q6 MoE expert slabs - // (`*_exps.*` tensors) — Metal ships no q6 MoE kernels. bf16-scale - // q8/q4 bundles stay `metal` (universal): both backends carry the - // `_sbf16` kernel families. - if self - .header - .tensors - .iter() - .any(|t| t.name.contains("_exps.") && t.dtype == TensorDtype::BaseQ6) - { - self.header.target_backend = crate::header::TargetBackend::CudaSm121; - } + self.stamp_target_backend_from_content(); // Multimodal sub-bundle entries land in the same weights blob, // continuing past the LM tensors. Their entries go into @@ -239,6 +426,157 @@ impl BaseWriter { self.inner.flush()?; Ok(()) } + + /// Streaming commit (file-backed writers). The blob was written to a + /// temp file as tensors were added, so here we only finalize the header + /// (offsets/lengths/checksums already recorded) and copy the temp blob + /// into the final file after the header. Constant memory throughout. + fn finish_streaming(mut self) -> Result<()> { + let mut s = self.stream.take().expect("finish_streaming without stream"); + + // Finalize the header from the streamed entries. + self.header.tensors = std::mem::take(&mut s.lm_entries); + self.stamp_target_backend_from_content(); + if !s.mmproj_entries.is_empty() { + let arch = self + .mmproj_arch + .clone() + .unwrap_or_else(|| "mmproj".to_string()); + self.header.mmproj = Some(MmprojBundle { + arch, + config: std::mem::take(&mut self.mmproj_config), + tensors: std::mem::take(&mut s.mmproj_entries), + }); + } + + // Remove the temp blob on EVERY exit from here on, not just the happy + // path. A failure below (the destination filling mid-copy is the + // realistic one) used to return through `?` and strand a model-sized + // `.blobtmp` next to the partial output — hundreds of GB that the user + // has to find by hand, and that makes every retry fail for space. + struct TmpGuard(std::path::PathBuf); + impl Drop for TmpGuard { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } + } + let _tmp_guard = TmpGuard(s.path.clone()); + + // Flush + rewind the temp blob for reading. + s.file.flush()?; + let mut tmp = s.file.into_inner().map_err(|e| e.into_error())?; + tmp.seek(SeekFrom::Start(0))?; + + // Prefix + header. + let header_json = self.header.to_canonical_json().map_err(Error::Json)?; + let header_len = header_json.len() as u64; + self.inner.write_all(&MAGIC)?; + self.inner.write_all(&FORMAT_VERSION.to_le_bytes())?; + self.inner.write_all(&header_len.to_le_bytes())?; + self.inner.write_all(&header_json)?; + + // Pad to blob start (BLOB_ALIGNMENT), then copy the streamed blob. + let header_end = PREFIX_LEN + header_len; + let blob_start = align_up(header_end, BLOB_ALIGNMENT); + write_zeros(&mut self.inner, (blob_start - header_end) as usize)?; + { + let mut reader = BufReader::new(&mut tmp); + let copied = std::io::copy(&mut reader, &mut self.inner)?; + debug_assert_eq!(copied, s.cursor, "streamed blob length mismatch"); + } + + // Extension slots after the blob (8-byte aligned). + if !self.slots.is_empty() { + let pad = ((8 - (s.cursor % 8)) % 8) as usize; + if pad > 0 { + self.inner.write_all(&[0u8; 8][..pad])?; + } + write_slots(&mut self.inner, &self.slots)?; + } + + self.inner.flush()?; + drop(tmp); + // _tmp_guard removes the blob here, and on every `?` above. + Ok(()) + } + + /// Finish for direct-write mode (`create_direct`): the blob already sits + /// at its final offset, so this only appends the extension slots to the + /// blob stream and backfills the reserved header region — no copy, no + /// temp file to delete. + fn finish_direct(mut self, reserve: u64) -> Result<()> { + let mut s = self.stream.take().expect("finish_direct without stream"); + + self.header.tensors = std::mem::take(&mut s.lm_entries); + self.stamp_target_backend_from_content(); + if !s.mmproj_entries.is_empty() { + let arch = self + .mmproj_arch + .clone() + .unwrap_or_else(|| "mmproj".to_string()); + self.header.mmproj = Some(MmprojBundle { + arch, + config: std::mem::take(&mut self.mmproj_config), + tensors: std::mem::take(&mut s.mmproj_entries), + }); + } + + // Extension slots continue past the blob (8-byte aligned) on the + // blob handle — it is already positioned at the blob end. + if !self.slots.is_empty() { + let pad = ((8 - (s.cursor % 8)) % 8) as usize; + if pad > 0 { + s.file.write_all(&[0u8; 8][..pad])?; + } + write_slots(&mut s.file, &self.slots)?; + } + s.file.flush()?; + drop(s); + + // Backfill the reserved header region. The JSON is space-padded to + // exactly `reserve` bytes (trailing whitespace is valid JSON), so + // header_len = reserve and the reader's blob-start computation + // matches the offset the blob was streamed at. + let header_json = self.header.to_canonical_json().map_err(Error::Json)?; + if header_json.len() as u64 > reserve { + return Err(Error::Io(std::io::Error::other(format!( + "direct-write header ({} bytes) exceeds the reserved {} bytes — \ + re-run with a larger header reserve", + header_json.len(), + reserve + )))); + } + self.inner.write_all(&MAGIC)?; + self.inner.write_all(&FORMAT_VERSION.to_le_bytes())?; + self.inner.write_all(&reserve.to_le_bytes())?; + self.inner.write_all(&header_json)?; + write_fill( + &mut self.inner, + b' ', + (reserve - header_json.len() as u64) as usize, + )?; + let header_end = PREFIX_LEN + reserve; + let blob_start = align_up(header_end, BLOB_ALIGNMENT); + write_zeros(&mut self.inner, (blob_start - header_end) as usize)?; + self.inner.flush()?; + Ok(()) + } +} + +/// Write `n` zero bytes to `w` without allocating an n-sized buffer. +fn write_zeros(w: &mut W, n: usize) -> std::io::Result<()> { + write_fill(w, 0u8, n) +} + +/// Write `n` copies of `byte` to `w` without allocating an n-sized buffer. +fn write_fill(w: &mut W, byte: u8, mut n: usize) -> std::io::Result<()> { + let buf = [byte; 8192]; + while n > 0 { + let chunk = n.min(buf.len()); + w.write_all(&buf[..chunk])?; + n -= chunk; + } + Ok(()) } fn align_up(x: u64, align: u64) -> u64 { diff --git a/base-convert/crates/base-format/tests/roundtrip.rs b/base-convert/crates/base-format/tests/roundtrip.rs index f29de66..475c25e 100644 --- a/base-convert/crates/base-format/tests/roundtrip.rs +++ b/base-convert/crates/base-format/tests/roundtrip.rs @@ -33,6 +33,7 @@ fn make_header() -> Header { tensors: vec![], mmproj: None, calibration: None, + provenance: None, sig: None, } } @@ -241,7 +242,9 @@ fn writer_fills_xxhash64_and_reader_verifies() { let reader = BaseReader::open(tmp.path()).unwrap(); let t = &reader.header().tensors[0]; assert_eq!(t.checksum_xxh64, Some(expected)); - reader.verify_tensor("checked").expect("checksum should match"); + reader + .verify_tensor("checked") + .expect("checksum should match"); } #[test] @@ -455,7 +458,10 @@ fn rejects_unknown_version() { bytes.extend_from_slice(&0u64.to_le_bytes()); std::fs::write(tmp.path(), &bytes).unwrap(); let err = expect_err(BaseReader::open(tmp.path())); - assert!(matches!(err, base_format::Error::UnsupportedVersion(999, _))); + assert!(matches!( + err, + base_format::Error::UnsupportedVersion(999, _) + )); } #[test] diff --git a/base-convert/crates/base-hub/Cargo.toml b/base-convert/crates/base-hub/Cargo.toml index 0dd506c..25de633 100644 --- a/base-convert/crates/base-hub/Cargo.toml +++ b/base-convert/crates/base-hub/Cargo.toml @@ -21,6 +21,10 @@ ureq = "3" # The HTTP client handed to hf-hub, so downloads carry our own timeouts. # Version must track hf-hub's reqwest, or the `Client` types do not unify. reqwest = { version = "0.13", default-features = false } +# flock for the per-destination reassembly lock. +libc = "0.2" +# A revision is one path segment of the tree URL, whatever characters it has. +percent-encoding = "2" [dev-dependencies] tempfile.workspace = true diff --git a/base-convert/crates/base-hub/catalog.json b/base-convert/crates/base-hub/catalog.json index 274274d..34c7f07 100644 --- a/base-convert/crates/base-hub/catalog.json +++ b/base-convert/crates/base-hub/catalog.json @@ -1,6 +1,6 @@ { "schema": 1, - "updated": "2026-08-19", + "updated": "2026-09-08", "models": [ { "id": "basecompute/Llama-3.1-8B-Instruct", @@ -1132,6 +1132,72 @@ "quant": "default-q8", "size": 62423808, "sha256": "3dbaaad2a23d0c68f1bac54963aadeb151c4bd38cf92d70cc2d27bf626feeb07" + }, + { + "id": "basecompute/gpt-oss-20b", + "hf_repo": "basecompute/gpt-oss-20b", + "file": "gpt-oss-20b-MXFP4.base", + "revision": "main", + "source_repo": "openai/gpt-oss-20b", + "arch": "gpt_oss", + "quant": "default-bf16", + "size": 13789700096, + "sha256": "bbd2324ad5fca191cc016daa5d696ee45c0dc96def18a14a77aa4f0799463f00" + }, + { + "id": "basecompute/gpt-oss-20b", + "hf_repo": "basecompute/gpt-oss-20b", + "file": "gpt-oss-20b-Q8.base", + "revision": "main", + "source_repo": "openai/gpt-oss-20b", + "arch": "gpt_oss", + "quant": "default-q8", + "size": 13183361024, + "sha256": "f95fe4a1c1fd69a8771515b8673c353afc49ecfce97fc48565477aa7e5235772" + }, + { + "id": "basecompute/gpt-oss-20b", + "hf_repo": "basecompute/gpt-oss-20b", + "file": "gpt-oss-20b-Q4.base", + "revision": "main", + "source_repo": "openai/gpt-oss-20b", + "arch": "gpt_oss", + "quant": "default-q4", + "size": 12873900032, + "sha256": "6b7ae00571be1b9d4e5204c34e9fe22c19ee69c5dd6e51d32900cf8d41091c7f" + }, + { + "id": "basecompute/NVIDIA-Nemotron-3-Nano-30B-A3B", + "hf_repo": "basecompute/NVIDIA-Nemotron-3-Nano-30B-A3B", + "file": "NVIDIA-Nemotron-3-Nano-30B-A3B-Q4.base", + "revision": "main", + "source_repo": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", + "arch": "nemotron_h_moe", + "quant": "default-q4", + "size": 17774510080, + "sha256": "1b13d5d33fdc7312ca6ac1fd23bc41c74c2fddc3219e038bf6060a7e6dae28fb" + }, + { + "id": "basecompute/NVIDIA-Nemotron-3-Nano-30B-A3B", + "hf_repo": "basecompute/NVIDIA-Nemotron-3-Nano-30B-A3B", + "file": "NVIDIA-Nemotron-3-Nano-30B-A3B-Q8.base", + "revision": "main", + "source_repo": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", + "arch": "nemotron_h_moe", + "quant": "default-q8", + "size": 33035943936, + "sha256": "417ff3cfc93fed22d3be71f0fcfb0317bb436c404656d793f9cbb8281e2bd7a4" + }, + { + "id": "basecompute/NVIDIA-Nemotron-3-Nano-30B-A3B", + "hf_repo": "basecompute/NVIDIA-Nemotron-3-Nano-30B-A3B", + "file": "NVIDIA-Nemotron-3-Nano-30B-A3B-BF16.base", + "revision": "main", + "source_repo": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", + "arch": "nemotron_h_moe", + "quant": "bf16", + "size": 63165857792, + "sha256": "f1bbee6ad05870c33d09a4ffeb2f6480b2fe13e34d4d8aaead8b302f115535e1" } ] } diff --git a/base-convert/crates/base-hub/src/catalog.rs b/base-convert/crates/base-hub/src/catalog.rs index 1aecf50..3096998 100644 --- a/base-convert/crates/base-hub/src/catalog.rs +++ b/base-convert/crates/base-hub/src/catalog.rs @@ -68,6 +68,13 @@ pub struct CatalogEntry { /// Optional integrity check for the downloaded `.base`. #[serde(default, skip_serializing_if = "Option::is_none")] pub sha256: Option, + /// For a bundle the Hub's file cap split into `.part-NNN` pieces: + /// the sha256 of each part, in order. The Hub knows only these, never a + /// whole-file hash, so this is what proves the listing still describes + /// the bundle `sha256` was pinned against, and what each part is checked + /// against as it lands. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parts_sha256: Option>, /// Backend requirement. `None` = universal (runs on every backend — /// f16/bf16 weights and bf16-scale q4/q8 bundles). `Some("cuda")` / /// `Some("metal")` restricts resolution to clients of that backend @@ -215,6 +222,17 @@ impl Catalog { if e.quant.is_empty() { anyhow::bail!("{}: empty quant", e.id); } + if let Some(parts) = &e.parts_sha256 { + if parts.is_empty() { + anyhow::bail!("{}: parts_sha256 is present but empty", e.id); + } + if let Some(bad) = parts + .iter() + .find(|p| p.len() != 64 || !p.bytes().all(|b| b.is_ascii_hexdigit())) + { + anyhow::bail!("{}: parts_sha256 entry {bad:?} is not a sha256", e.id); + } + } // A model may carry a universal row AND a per-backend variant of the // same quant (the resolver prefers the backend-native one); those are // distinguished by `backend`, so the uniqueness key includes it. Two @@ -290,7 +308,10 @@ mod tests { #[test] fn bundled_catalog_is_populated_and_resolves_default_q4() { let cat = Catalog::bundled().unwrap(); - assert!(!cat.models.is_empty(), "bundled catalog should not be empty"); + assert!( + !cat.models.is_empty(), + "bundled catalog should not be empty" + ); // Every entry carries the fields the resolver/installer need. for e in &cat.models { assert!(e.id.starts_with("basecompute/"), "id: {}", e.id); diff --git a/base-convert/crates/base-hub/src/fetch.rs b/base-convert/crates/base-hub/src/fetch.rs index 9d072e5..493ec4c 100644 --- a/base-convert/crates/base-hub/src/fetch.rs +++ b/base-convert/crates/base-hub/src/fetch.rs @@ -10,7 +10,6 @@ use anyhow::{Context, Result}; use hf_hub::progress::{DownloadEvent, FileStatus, Progress, ProgressEvent, ProgressHandler}; -use hf_hub::repository::RepoTreeEntry; use hf_hub::{HFClient, HFClientSync, HFRepositorySync, RepoTypeModel}; use indicatif::{ProgressBar, ProgressStyle}; use std::collections::HashMap; @@ -82,6 +81,24 @@ pub trait Fetcher { None } + /// Pin `revision` to something immutable — the commit it names on the + /// Hub — so that a sequence of requests against a moving branch all + /// observe one publication. Sources with no such notion return the + /// revision unchanged. + fn resolve_revision(&self, repo: &str, revision: &str) -> Result { + let _ = repo; + Ok(revision.to_string()) + } + + /// A stable identifier for the *content* of `filename` at `revision` — + /// the LFS sha256 on the Hub — or `None` when the source has no such + /// notion (fixtures). Lets a multi-file install notice that a file it + /// already consumed has since been replaced under the same name. + fn content_id(&self, repo: &str, revision: &str, filename: &str) -> Result> { + let _ = (repo, revision, filename); + Ok(None) + } + /// Read `range` of `filename` without downloading the rest. /// /// A `.base` header is a 16-byte prefix plus a JSON blob, both at the front @@ -200,6 +217,200 @@ impl ProgressHandler for BarProgress { } } +/// The Hub token, if any: `$HF_TOKEN`, the legacy `$HUGGING_FACE_HUB_TOKEN`, +/// the file `$HF_TOKEN_PATH` names, then the cached login under `$HF_HOME` +/// (default `~/.cache/huggingface`). The same order hf-hub uses, spelled +/// out here because the raw tree listing below is not an hf-hub call. +pub(crate) fn resolve_token() -> Option { + let from_env = |k: &str| std::env::var(k).ok().filter(|s| !s.trim().is_empty()); + if let Some(t) = from_env("HF_TOKEN").or_else(|| from_env("HUGGING_FACE_HUB_TOKEN")) { + return Some(t.trim().to_string()); + } + let path = std::env::var_os("HF_TOKEN_PATH") + .map(PathBuf::from) + .or_else(|| hf_home().map(|h| h.join("token")))?; + std::fs::read_to_string(path) + .ok() + .map(|t| t.trim().to_string()) + .filter(|t| !t.is_empty()) +} + +/// Where the cached login lives, in hf-hub's own order: `$HF_HOME`, then +/// `$XDG_CACHE_HOME/huggingface`, then `~/.cache/huggingface`. +fn hf_home() -> Option { + hf_home_from( + std::env::var_os("HF_HOME").map(PathBuf::from), + std::env::var_os("XDG_CACHE_HOME").map(PathBuf::from), + dirs::home_dir(), + ) +} + +fn hf_home_from( + hf_home: Option, + xdg_cache: Option, + home: Option, +) -> Option { + let nonempty = |p: PathBuf| (!p.as_os_str().is_empty()).then_some(p); + hf_home + .and_then(nonempty) + .or_else(|| xdg_cache.and_then(nonempty).map(|x| x.join("huggingface"))) + .or_else(|| home.map(|h| h.join(".cache").join("huggingface"))) +} + +/// Hub API base. `$HF_ENDPOINT` redirects everything at a mirror, the same +/// variable hf-hub honors for downloads. +pub(crate) fn endpoint() -> String { + std::env::var("HF_ENDPOINT") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "https://huggingface.co".to_string()) +} + +/// One file in a repo tree, as the Hub's own JSON describes it. +/// +/// Read raw rather than through hf-hub's typed listing: that type expects +/// `sha256`/`pointer_size` under `lfs` while the Hub sends `oid`/`pointerSize`, +/// so the LFS hash — the one number that identifies a part's bytes — always +/// came back `None` and routing fell back to the git object id, which is +/// not a hash anything can check a download against. +#[derive(Debug, Clone)] +pub(crate) struct TreeEntry { + pub path: String, + /// Size of the content (the LFS payload where there is one). + pub size: u64, + /// The git object id — a sha1 over the pointer, not the content. + pub git_oid: String, + /// sha256 of the content, for LFS-backed files. + pub lfs_sha256: Option, + /// Served through Xet. + pub xet: bool, +} + +/// Parse one page of `/api/models//tree/` into file entries. +pub(crate) fn parse_tree(body: &str) -> Result> { + let entries: Vec = + serde_json::from_str(body).context("parsing the tree listing")?; + Ok(entries + .iter() + .filter(|e| e.get("type").and_then(|t| t.as_str()) == Some("file")) + .filter_map(|e| { + let path = e.get("path")?.as_str()?.to_string(); + let git_oid = e.get("oid")?.as_str()?.to_string(); + let lfs = e.get("lfs").filter(|l| !l.is_null()); + Some(TreeEntry { + size: lfs + .and_then(|l| l.get("size")) + .or_else(|| e.get("size")) + .and_then(|v| v.as_u64()) + .unwrap_or(0), + lfs_sha256: lfs + .and_then(|l| l.get("oid")) + .and_then(|v| v.as_str()) + .map(str::to_string), + xet: e.get("xetHash").is_some_and(|x| !x.is_null()), + git_oid, + path, + }) + }) + .collect()) +} + +/// What a pinned revision may be: the commit the Hub reports for the +/// requested one, or the requested one itself when it is already a commit +/// id. A branch that resolves to nothing usable is refused rather than +/// handed on as if it were immutable — everything after the pin assumes it +/// cannot move. +pub(crate) fn pinned_revision(requested: &str, sha: Option) -> Result { + let is_commit = |s: &str| s.len() == 40 && s.bytes().all(|b| b.is_ascii_hexdigit()); + match sha { + Some(sha) if is_commit(&sha) => Ok(sha), + _ if is_commit(requested) => Ok(requested.to_string()), + other => anyhow::bail!( + "the Hub did not resolve {requested:?} to a commit (got {other:?}); pass a commit id as the revision" + ), + } +} + +/// `revision` as one URL path segment: a ref like `feature/foo` or +/// `release#1` must not become a sub-path or a fragment. +pub(crate) fn encode_segment(s: &str) -> String { + use percent_encoding::{utf8_percent_encode, AsciiSet, NON_ALPHANUMERIC}; + const KEEP: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'~'); + utf8_percent_encode(s, KEEP).to_string() +} + +/// The whole tree of `repo` at `revision`, following the Hub's `Link` +/// pagination. +pub(crate) fn fetch_tree(repo: &str, revision: &str) -> Result> { + let url = Some(format!( + "{}/api/models/{repo}/tree/{}?recursive=true", + endpoint(), + encode_segment(revision) + )); + let token = resolve_token(); + // The same bounded reads the download client gets: a mirror that sends + // headers and then goes quiet must fail, not hang the pull before it + // starts. + let agent: ureq::Agent = ureq::Agent::config_builder() + .timeout_connect(Some(Duration::from_secs(30))) + .timeout_recv_response(Some(crate::download::read_timeout())) + .timeout_recv_body(Some(crate::download::read_timeout())) + .build() + .into(); + collect_pages(url, |u| { + let mut req = agent.get(u); + if let Some(t) = &token { + req = req.header("Authorization", &format!("Bearer {t}")); + } + let resp = req + .call() + .with_context(|| format!("listing files in {repo}@{revision}"))?; + let next = resp + .headers() + .get("link") + .and_then(|v| v.to_str().ok()) + .and_then(crate::scan::parse_next_link); + let body = resp + .into_body() + .read_to_string() + .with_context(|| format!("reading the file listing for {repo}@{revision}"))?; + Ok((parse_tree(&body)?, next)) + }) +} + +/// Pages a listing may run to before it is treated as a misbehaving +/// mirror rather than a big repo (a page is up to 1000 entries). +const MAX_TREE_PAGES: usize = 100; + +/// Follow `next` links from `first`, `fetch` returning one page's entries +/// and the link after it. A listing still pointing onward at the cap is an +/// error: a truncated tree would be memoized and make files vanish. +fn collect_pages( + first: Option, + mut fetch: impl FnMut(&str) -> Result<(Vec, Option)>, +) -> Result> { + let mut url = first; + let mut out = Vec::new(); + for _ in 0..MAX_TREE_PAGES { + let Some(u) = url.take() else { + return Ok(out); + }; + let (page, next) = fetch(&u)?; + out.extend(page); + url = next; + } + match url { + None => Ok(out), + Some(more) => anyhow::bail!( + "the file listing did not end after {MAX_TREE_PAGES} pages (next: {more}); refusing a truncated tree" + ), + } +} + /// Build one hf-hub client pointed at `staging_root`. /// /// The reqwest client is ours rather than hf-hub's default so it can carry @@ -220,13 +431,9 @@ fn build_client(staging_root: &Path) -> Result { .retry_max_attempts(resolve_max_retries()) .client(http); // hf-hub resolves `$HF_TOKEN` itself but not the legacy - // `$HUGGING_FACE_HUB_TOKEN`, so both are read here and passed explicitly; - // an unset pair leaves hf-hub's own resolution in place. - if let Some(tok) = std::env::var("HF_TOKEN") - .ok() - .or_else(|| std::env::var("HUGGING_FACE_HUB_TOKEN").ok()) - .filter(|s| !s.is_empty()) - { + // `$HUGGING_FACE_HUB_TOKEN`; resolve once here (see `resolve_token`) and + // pass it explicitly so every request agrees on who is asking. + if let Some(tok) = resolve_token() { builder = builder.token(tok); } builder @@ -252,7 +459,7 @@ pub struct HfFetcher { /// several files in the same repo — the artifact, then its sidecars — and /// the listing that answers "how big, and is it Xet?" is the same one /// `list_files` needs. - trees: Mutex>>, + trees: Mutex>>, } /// What routing needs to know about one remote file. @@ -276,18 +483,12 @@ impl HfFetcher { } /// The repo's file tree at `revision`, fetched once and remembered. - fn tree(&self, repo: &str, revision: &str) -> Result> { + fn tree(&self, repo: &str, revision: &str) -> Result> { let key = (repo.to_string(), revision.to_string()); if let Some(hit) = self.trees.lock().unwrap().get(&key) { return Ok(hit.clone()); } - let entries = self - .repo(repo) - .list_tree() - .revision(revision) - .recursive(true) - .send() - .with_context(|| format!("listing files in {repo}@{revision}"))?; + let entries = fetch_tree(repo, revision)?; self.trees.lock().unwrap().insert(key, entries.clone()); Ok(entries) } @@ -305,20 +506,11 @@ impl HfFetcher { Ok(self .tree(repo, revision)? .into_iter() - .find_map(|e| match e { - RepoTreeEntry::File { - path, - size, - oid, - lfs, - xet_hash, - .. - } if path == filename => Some(FileFacts { - size: lfs.as_ref().and_then(|l| l.size).unwrap_or(size), - xet: xet_hash.is_some(), - key: lfs.and_then(|l| l.sha256).unwrap_or(oid), - }), - _ => None, + .find(|e| e.path == filename) + .map(|e| FileFacts { + size: e.size, + xet: e.xet, + key: e.lfs_sha256.unwrap_or(e.git_oid), })) } @@ -398,10 +590,7 @@ impl Fetcher for HfFetcher { Ok(self .tree(repo, revision)? .into_iter() - .filter_map(|e| match e { - RepoTreeEntry::File { path, .. } => Some(path), - _ => None, - }) + .map(|e| e.path) .collect()) } @@ -415,6 +604,27 @@ impl Fetcher for HfFetcher { ) } + fn resolve_revision(&self, repo: &str, revision: &str) -> Result { + let info = self + .repo(repo) + .info() + .revision(revision.to_string()) + .send() + .with_context(|| format!("looking up {repo}@{revision}"))?; + pinned_revision(revision, info.sha) + } + + fn content_id(&self, repo: &str, revision: &str, filename: &str) -> Result> { + // Only the LFS hash is a hash of the bytes. A file git stores inline + // has a sha1 over its header and content, which nothing downstream + // can compare a download against, so it reports no id at all. + Ok(self + .tree(repo, revision)? + .into_iter() + .find(|e| e.path == filename) + .and_then(|e| e.lfs_sha256)) + } + fn read_range( &self, repo: &str, @@ -530,6 +740,102 @@ impl Fetcher for MockFetcher { mod tests { use super::*; + #[test] + fn the_cached_login_is_looked_for_where_hf_hub_puts_it() { + let p = |s: &str| Some(PathBuf::from(s)); + assert_eq!(hf_home_from(p("/hf"), p("/xdg"), p("/home/u")), p("/hf")); + assert_eq!( + hf_home_from(None, p("/xdg"), p("/home/u")), + p("/xdg/huggingface") + ); + assert_eq!( + hf_home_from(p(""), p("/xdg"), p("/home/u")), + p("/xdg/huggingface") + ); + assert_eq!( + hf_home_from(None, None, p("/home/u")), + p("/home/u/.cache/huggingface") + ); + assert_eq!(hf_home_from(None, None, None), None); + } + + #[test] + fn a_listing_that_never_ends_is_refused_not_truncated() { + // Every page points onward: exhausting the cap is an error, and + // nothing partial is returned. + let mut pages = 0; + let err = collect_pages(Some("p0".to_string()), |_| { + pages += 1; + Ok((vec![pages], Some(format!("p{pages}")))) + }) + .unwrap_err() + .to_string(); + assert!(err.contains("truncated tree"), "{err}"); + assert_eq!(pages, MAX_TREE_PAGES); + + // A listing that ends is returned whole, however many pages. + let got = collect_pages(Some("p0".to_string()), |u| { + let n: usize = u[1..].parse().unwrap(); + Ok((vec![n], (n < 3).then(|| format!("p{}", n + 1)))) + }) + .unwrap(); + assert_eq!(got, vec![0, 1, 2, 3]); + let none: Vec = collect_pages(None, |_| -> Result<(Vec, Option)> { + unreachable!() + }) + .unwrap(); + assert!(none.is_empty()); + } + + #[test] + fn a_pin_is_a_commit_or_nothing() { + let sha = "7c73ace2115ad5a838277152d555ec1229281c46".to_string(); + assert_eq!(pinned_revision("main", Some(sha.clone())).unwrap(), sha); + // Asked for a commit already: it stands on its own. + assert_eq!(pinned_revision(&sha, None).unwrap(), sha); + // A branch the Hub cannot resolve is not quietly kept mutable. + let err = pinned_revision("main", None).unwrap_err().to_string(); + assert!(err.contains("did not resolve"), "{err}"); + let err = pinned_revision("main", Some("not-a-sha".into())) + .unwrap_err() + .to_string(); + assert!(err.contains("did not resolve"), "{err}"); + } + + #[test] + fn a_revision_is_one_path_segment() { + assert_eq!(encode_segment("main"), "main"); + assert_eq!(encode_segment("feature/foo"), "feature%2Ffoo"); + assert_eq!(encode_segment("release#1"), "release%231"); + assert_eq!(encode_segment("v1.2-rc_3~x"), "v1.2-rc_3~x"); + } + + #[test] + fn tree_listing_keeps_the_lfs_hash_the_typed_client_drops() { + // Verbatim shape of the Hub's answer: `lfs.oid` is the sha256 of the + // content, top-level `oid` the git object id, and small files have + // no `lfs` block at all. + let body = r#"[ + {"type":"file","oid":"7bc52451a2e2576b266715c7898b9d79dbe0bb25","size":134, + "lfs":{"oid":"cc36eece6e94329331b5c4abe4f0d31c06d56658c4360ebb1c4b8a974ed20bfe","size":48318382080,"pointerSize":134}, + "path":"parts/GLM-5.2-Q4.base.part-000"}, + {"type":"file","oid":"abc","size":12,"path":"README.md"}, + {"type":"directory","oid":"def","path":"parts"} + ]"#; + let got = parse_tree(body).unwrap(); + assert_eq!(got.len(), 2); + assert_eq!(got[0].path, "parts/GLM-5.2-Q4.base.part-000"); + assert_eq!(got[0].size, 48318382080); + assert_eq!( + got[0].lfs_sha256.as_deref(), + Some("cc36eece6e94329331b5c4abe4f0d31c06d56658c4360ebb1c4b8a974ed20bfe") + ); + assert_eq!(got[0].git_oid, "7bc52451a2e2576b266715c7898b9d79dbe0bb25"); + assert!(!got[0].xet); + assert_eq!(got[1].size, 12); + assert_eq!(got[1].lfs_sha256, None); + } + /// Fetcher that owns an hf-hub-style staging tree: /// `/models----/blobs/` with /// `snapshots//` symlinks pointing at the blobs — the layout diff --git a/base-convert/crates/base-hub/src/gen.rs b/base-convert/crates/base-hub/src/gen.rs index 8f0d9cf..f9f9dcb 100644 --- a/base-convert/crates/base-hub/src/gen.rs +++ b/base-convert/crates/base-hub/src/gen.rs @@ -118,6 +118,7 @@ pub fn entry_from_header( quant, size: Some(size), sha256: Some(sha256), + parts_sha256: None, backend, }) } diff --git a/base-convert/crates/base-hub/src/lib.rs b/base-convert/crates/base-hub/src/lib.rs index c2eef0b..28e8d41 100644 --- a/base-convert/crates/base-hub/src/lib.rs +++ b/base-convert/crates/base-hub/src/lib.rs @@ -11,6 +11,7 @@ pub mod catalog; pub mod download; pub mod fetch; pub mod gen; +pub mod parts; pub mod registry; pub mod scan; diff --git a/base-convert/crates/base-hub/src/parts.rs b/base-convert/crates/base-hub/src/parts.rs new file mode 100644 index 0000000..a29797d --- /dev/null +++ b/base-convert/crates/base-hub/src/parts.rs @@ -0,0 +1,1862 @@ +//! Split `.base` bundles: recognizing a part set in a repo listing and +//! reassembling one into a single artifact on install. +//! +//! The Hub caps a single file at 50 GB. A bundle past that ships as +//! `.base.part-000`, `.base.part-001`, … (a plain byte split, +//! the same convention as multi-part GGUFs), and the logical artifact is +//! `.base`. Nothing else about the bundle changes: part 000 opens with +//! the ordinary header, so a ranged read of it still answers "what is this". +//! +//! Reassembly is streamed into `.partial` and renamed into place only +//! when every part has landed, so a half-built bundle never looks installed. +//! Each staged part is deleted as soon as it has been appended, which keeps +//! peak disk at one bundle plus one part rather than two bundles. A record +//! beside the partial says how many parts it holds, so a pull killed between +//! parts resumes at the next part instead of at byte zero. +//! +//! What makes the stitched result trustworthy, given that nothing on the Hub +//! carries a whole-file checksum for it: +//! +//! * a manifest beside the parts (`.base.manifest.json`, see +//! [`Manifest`]) says how many parts there are, how long each is, the +//! sha256 of each, and the sha256 and length of the whole. Nothing in the +//! part filenames says how many there should be, and a bundle may end in +//! extension slots the header does not flag, so the manifest is the one +//! thing that can prove the tail was not lost. It is required; +//! * the revision is pinned to a commit before the first byte moves, so a +//! mutable branch advancing mid-pull cannot mix two publications; +//! * every part is hashed as it is appended and checked against the +//! manifest, and the Hub's own per-part hashes must agree with it; +//! * the finished file is parsed as a `.base` as a second line of defence: +//! its weights blob and any slot section have to fit; +//! * one process at a time works on a destination, held by a lock file. + +use crate::fetch::{install_file, Fetcher}; +use anyhow::{bail, Context, Result}; +use base_format::{BLOB_ALIGNMENT, FORMAT_VERSION, MAGIC, PREFIX_LEN}; +use indicatif::ProgressBar; +use sha2::{Digest, Sha256}; +use std::collections::BTreeMap; +use std::fs::{File, OpenOptions}; +use std::io::{BufReader, Read, Write}; +use std::os::unix::io::AsRawFd; +use std::path::{Path, PathBuf}; + +/// A `.base` bundle as a repo hosts it: one file, or an ordered part set. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Artifact { + /// The logical `.base` path (`parts/GLM-5.2-Q4.base`), which is also the + /// real path when `parts` is empty. + pub name: String, + /// Part paths in byte order; empty for a whole file. + pub parts: Vec, +} + +impl Artifact { + pub fn whole(name: impl Into) -> Self { + Self { + name: name.into(), + parts: Vec::new(), + } + } + + pub fn is_split(&self) -> bool { + !self.parts.is_empty() + } +} + +/// `.base.part-` → (`.base`, index). Anything else is +/// not a part. +pub fn split_part_name(path: &str) -> Option<(String, u32)> { + let (name, idx) = path.rsplit_once(".part-")?; + if !name.ends_with(".base") || idx.is_empty() || !idx.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + Some((name.to_string(), idx.parse().ok()?)) +} + +/// A repo listing sorted into what can be pulled and what cannot. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct Grouped { + pub artifacts: Vec, + /// Logical names of part sets that are not usable, with the reason — + /// a gap in the indices, typically a quant still being uploaded. Kept + /// apart so one bad set does not take a repo's other bundles with it. + pub malformed: Vec<(String, String)>, +} + +/// Group a repo listing into artifacts: every whole `.base` file, plus one +/// artifact per part set. A set with a gap in its indices is reported as +/// malformed rather than accepted as a shorter bundle; a set that simply +/// stops early cannot be told apart here and is caught by the completeness +/// check on install. When a repo ships both a whole file and parts under +/// the same name, the whole file wins. +pub fn group(files: impl IntoIterator) -> Grouped { + let mut whole = Vec::new(); + let mut sets: BTreeMap> = BTreeMap::new(); + // Two spellings of one index (`part-000` and `part-0000`) are two files + // claiming the same slot; neither can be picked over the other. + let mut doubled: BTreeMap = BTreeMap::new(); + for f in files { + if let Some((name, idx)) = split_part_name(&f) { + if let Some(other) = sets.entry(name.clone()).or_default().insert(idx, f.clone()) { + doubled + .entry(name) + .or_insert_with(|| format!("index {idx} is listed twice ({other} and {f})")); + } + } else if f.ends_with(".base") { + whole.push(f); + } + } + let mut out = Grouped { + artifacts: whole.into_iter().map(Artifact::whole).collect(), + malformed: Vec::new(), + }; + for (name, parts) in sets { + if out.artifacts.iter().any(|a| a.name == name) { + continue; + } + if let Some(why) = doubled.remove(&name) { + out.malformed.push((name, why)); + continue; + } + let gap = (0u32..) + .zip(parts.keys()) + .find(|(want, have)| *have != want) + .map(|(want, _)| want); + if let Some(want) = gap { + out.malformed.push(( + name, + format!( + "part set is missing part {want:03} (found {} parts)", + parts.len() + ), + )); + continue; + } + out.artifacts.push(Artifact { + name, + parts: parts.into_values().collect(), + }); + } + out +} + +/// The artifact `name` denotes in `repo`: the part set behind it when the +/// listing shows one, else the whole file. A name the listing does not know +/// is returned as a whole file so the download itself reports the miss; a +/// name that is a malformed part set is an error saying why. +pub fn find(fetcher: &dyn Fetcher, repo: &str, revision: &str, name: &str) -> Result { + let files = fetcher + .list_files(repo, revision) + .with_context(|| format!("listing files in {repo}@{revision}"))?; + let grouped = group(files); + if let Some((_, why)) = grouped.malformed.iter().find(|(n, _)| n == name) { + bail!("{repo}/{name}: {why}"); + } + Ok(grouped + .artifacts + .into_iter() + .find(|a| a.name == name) + .unwrap_or_else(|| Artifact::whole(name))) +} + +/// What a publisher states about a split bundle, in `.manifest.json` +/// beside the parts. +/// +/// ```json +/// {"size": 432223, "sha256": "…whole file…", +/// "parts": [{"name": "X.base.part-000", "size": 45, "sha256": "…"}, …]} +/// ``` +/// +/// `name` is the part's basename; parts are listed in byte order. Written +/// by `basert catalog-manifest` from the local parts, without reassembling +/// them: the whole-file hash is one sha256 run across the parts in order. +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct Manifest { + pub size: u64, + pub sha256: String, + pub parts: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct ManifestPart { + pub name: String, + pub size: u64, + pub sha256: String, +} + +/// `parts/X.base` → `parts/X.base.manifest.json`. +pub fn manifest_name(name: &str) -> String { + format!("{name}.manifest.json") +} + +impl Manifest { + /// The per-part hashes, in order. + pub fn part_ids(&self) -> Vec { + self.parts.iter().map(|p| p.sha256.clone()).collect() + } + + /// Does this manifest describe exactly `parts` (paths in byte order)? + /// Count, names and the size sum all have to line up, and every hash + /// has to be one — a malformed whole-file hash would otherwise ride + /// into the catalog and fail every pull only after the last byte. + pub fn check_against(&self, parts: &[String]) -> Result<()> { + if !is_sha256(&self.sha256) { + bail!("the manifest's whole-file sha256 is not a sha256"); + } + if self.parts.len() != parts.len() { + bail!( + "the manifest lists {} parts but the repo has {}", + self.parts.len(), + parts.len() + ); + } + for (m, p) in self.parts.iter().zip(parts) { + let base = p.rsplit('/').next().unwrap_or(p); + if m.name != base { + bail!("the manifest names {} where the repo has {base}", m.name); + } + if !is_sha256(&m.sha256) { + bail!("the manifest's sha256 for {} is not a sha256", m.name); + } + } + let sum = self + .parts + .iter() + .try_fold(0u64, |a, p| a.checked_add(p.size)) + .context("manifest part sizes overflow")?; + if sum != self.size { + bail!( + "the manifest's parts sum to {sum} bytes but it says the whole is {}", + self.size + ); + } + Ok(()) + } + + /// Do the per-part sizes match what the source reports for each part? + /// Two wrong sizes can cancel out in the sum; they cannot here. + pub fn check_sizes(&self, listed: &[u64]) -> Result<()> { + if listed.len() != self.parts.len() { + bail!( + "the manifest lists {} parts but {} sizes were given", + self.parts.len(), + listed.len() + ); + } + for (m, have) in self.parts.iter().zip(listed) { + if m.size != *have { + bail!( + "the manifest says {} is {} bytes but the repo holds {have}", + m.name, + m.size + ); + } + } + Ok(()) + } +} + +/// Build a manifest from local part files, in the order given, hashing +/// each part and the whole in one pass. +pub fn build_manifest(parts: &[PathBuf]) -> Result { + if parts.is_empty() { + bail!("no parts to describe"); + } + let mut whole = Sha256::new(); + let mut out = Manifest { + size: 0, + sha256: String::new(), + parts: Vec::with_capacity(parts.len()), + }; + let mut buf = vec![0u8; 8 << 20]; + for path in parts { + let name = path + .file_name() + .and_then(|n| n.to_str()) + .with_context(|| format!("{}: not a file name", path.display()))? + .to_string(); + let mut reader = BufReader::with_capacity(8 << 20, File::open(path)?); + let mut one = Sha256::new(); + let mut size = 0u64; + loop { + let n = reader.read(&mut buf)?; + if n == 0 { + break; + } + one.update(&buf[..n]); + whole.update(&buf[..n]); + size += n as u64; + } + out.size += size; + out.parts.push(ManifestPart { + name, + size, + sha256: hex(one.finalize().as_slice()), + }); + } + out.sha256 = hex(whole.finalize().as_slice()); + Ok(out) +} + +fn hex(bytes: &[u8]) -> String { + bytes.iter().map(|b| format!("{b:02x}")).collect() +} + +/// A manifest is a few KB; anything claiming to be larger is not one, and +/// is refused before it is allocated. Shared with the catalog scan. +pub const MAX_MANIFEST_LEN: u64 = 1024 * 1024; + +/// The manifest published beside `name`'s parts. Its absence is an error +/// that says what to publish: without it a lost tail is undetectable. +fn read_manifest( + fetcher: &dyn Fetcher, + repo: &str, + revision: &str, + name: &str, +) -> Result { + let mname = manifest_name(name); + let path = fetcher.get_file(repo, revision, &mname).with_context(|| { + format!( + "{repo}/{mname}: a split bundle needs its manifest beside the parts (publish one with `basert catalog-manifest`)" + ) + })?; + let len = std::fs::metadata(&path) + .with_context(|| format!("sizing {}", path.display()))? + .len(); + if len == 0 || len >= MAX_MANIFEST_LEN { + bail!("{repo}/{mname}: implausible manifest size {len}"); + } + let bytes = std::fs::read(&path).with_context(|| format!("reading {}", path.display()))?; + serde_json::from_slice(&bytes).with_context(|| format!("parsing {repo}/{mname}")) +} + +/// Download `artifact` and install it at `dst`, reassembling a part set. +/// +/// `expected_ids` is what a catalog row pinned for the parts (their sha256s, +/// in order). When given, the Hub's listing has to agree with it before any +/// byte is fetched, and each part is verified against it as it lands; a +/// disagreement means the bundle was republished since the row was written. +pub fn install( + fetcher: &dyn Fetcher, + repo: &str, + revision: &str, + artifact: &Artifact, + dst: &Path, + expected_ids: Option<&[String]>, +) -> Result<()> { + // Pin the revision first: every later request names the commit, not + // the branch, so a publish landing mid-pull cannot be half-observed. + let pinned = fetcher + .resolve_revision(repo, revision) + .with_context(|| format!("resolving {repo}@{revision}"))?; + if pinned != revision { + eprintln!( + " pinned: {revision} → {}", + &pinned[..pinned.len().min(12)] + ); + } + // The caller's artifact came from whatever the branch pointed at when + // it listed; look again at the pinned commit before deciding anything + // — whether it is whole or split included — so the parts, the ids, and + // the downloads all describe one publication. + let artifact = find(fetcher, repo, &pinned, &artifact.name)?; + if !artifact.is_split() { + let src = fetcher.get_file(repo, &pinned, &artifact.name)?; + return install_file(fetcher, repo, &src, dst); + } + let manifest = read_manifest(fetcher, repo, &pinned, &artifact.name)?; + manifest.check_against(&artifact.parts).with_context(|| { + format!( + "{repo}/{}: manifest disagrees with the listing", + artifact.name + ) + })?; + let ids = manifest.part_ids(); + // The Hub's own per-part hashes, where it has them, have to agree with + // the manifest; and so does what the catalog pinned. + for (part, id) in artifact.parts.iter().zip(&ids) { + if let Some(listed) = fetcher.content_id(repo, &pinned, part)? { + if !listed.eq_ignore_ascii_case(id) { + bail!( + "{repo}/{part} is {listed} on the Hub but the manifest says {id}: the manifest is stale, republish it" + ); + } + } + } + if let Some(expected) = expected_ids { + reconcile_ids(&artifact.parts, &ids, expected)?; + } + + let lock = Lock::acquire(dst)?; + Reassembly::open(dst, repo, &pinned, &artifact.parts, ids, manifest)? + .run(fetcher, repo, &pinned)?; + lock.release(); + Ok(()) +} + +/// What the catalog row pinned has to be what the manifest says now; a +/// difference means the bundle was republished since the row was written. +fn reconcile_ids(parts: &[String], manifest: &[String], expected: &[String]) -> Result<()> { + if expected.len() != parts.len() { + bail!( + "the catalog pins {} parts for this bundle but the Hub lists {}: the bundle was republished, refresh the catalog", + expected.len(), + parts.len() + ); + } + for ((part, now), expected) in parts.iter().zip(manifest).zip(expected) { + if !now.eq_ignore_ascii_case(expected) { + bail!( + "{part} is now {now}, not the {expected} the catalog pinned: the bundle was republished, refresh the catalog" + ); + } + } + Ok(()) +} + +/// Where a part set is stitched together before it becomes `dst`. +pub fn partial_path(dst: &Path) -> PathBuf { + with_suffix(dst, ".partial") +} + +fn record_path(dst: &Path) -> PathBuf { + with_suffix(dst, ".partial.json") +} + +fn lock_path(dst: &Path) -> PathBuf { + with_suffix(dst, ".partial.lock") +} + +fn with_suffix(p: &Path, suffix: &str) -> PathBuf { + let mut s = p.as_os_str().to_owned(); + s.push(suffix); + PathBuf::from(s) +} + +/// Exclusive hold on one destination for the length of a reassembly, so two +/// pulls of the same variant cannot both append to one partial. Taken +/// without waiting: the second pull is told, not queued behind a 400 GB +/// download it would only duplicate. +struct Lock { + file: File, + path: PathBuf, +} + +impl Lock { + fn acquire(dst: &Path) -> Result { + let path = lock_path(dst); + let file = OpenOptions::new() + .create(true) + .write(true) + .truncate(false) + .open(&path) + .with_context(|| format!("opening {}", path.display()))?; + // Safety: a valid fd for the lifetime of the call; flock has no + // memory-safety preconditions beyond that. + let rc = unsafe { libc::flock(file.as_raw_fd(), libc::LOCK_EX | libc::LOCK_NB) }; + if rc != 0 { + let err = std::io::Error::last_os_error(); + if err.raw_os_error() == Some(libc::EWOULDBLOCK) { + bail!( + "another basert pull is already assembling {}; wait for it or stop it first", + dst.display() + ); + } + return Err(err).with_context(|| format!("locking {}", path.display())); + } + Ok(Self { file, path }) + } + + /// Remove the lock file once the install is in place. The lock is still + /// held while the file is unlinked, so a pull arriving in that window + /// creates a fresh lock file rather than sharing this one. + fn release(self) { + let _ = std::fs::remove_file(&self.path); + drop(self.file); + } +} + +/// What the partial holds, persisted after every part so a restart can pick +/// up where the last one stopped. +/// +/// Resuming is only safe when the remaining parts continue the same bytes +/// the partial already holds, so the record names where they came from: +/// the repo, the pinned commit, the part paths, and each part's content id. +/// A mutable revision republished between two attempts pins to a different +/// commit, and the partial is rebuilt rather than spliced. +#[derive(Debug, serde::Serialize, serde::Deserialize, PartialEq, Eq)] +struct Record { + repo: String, + revision: String, + parts: Vec, + /// Per-part sha256s, from the manifest. + ids: Vec, + /// Parts fully appended, in order. + appended: usize, + /// Byte length of the partial once those parts were appended. Anything + /// past it is a torn write from an interrupted append and is cut off. + len: u64, +} + +struct Reassembly { + dst: PathBuf, + partial: PathBuf, + record_path: PathBuf, + record: Record, + manifest: Manifest, + /// Some of the partial predates this process: its bytes were checked + /// by an earlier attempt, not this one. + resumed: bool, + /// Running sha256 of everything this process appended, so a fresh + /// install can compare the whole with the manifest without a second + /// read of the file. + whole: Sha256, +} + +impl Reassembly { + fn open( + dst: &Path, + repo: &str, + revision: &str, + parts: &[String], + ids: Vec, + manifest: Manifest, + ) -> Result { + let partial = partial_path(dst); + let record_path = record_path(dst); + let fresh = Record { + repo: repo.to_string(), + revision: revision.to_string(), + parts: parts.to_vec(), + ids, + appended: 0, + len: 0, + }; + let same_source = |r: &Record| { + r.repo == fresh.repo + && r.revision == fresh.revision + && r.parts == fresh.parts + && r.ids == fresh.ids + }; + let record = match std::fs::read(&record_path) { + Ok(bytes) => serde_json::from_slice::(&bytes) + .ok() + .filter(same_source) + .filter(|r| std::fs::metadata(&partial).is_ok_and(|m| m.len() >= r.len)) + .unwrap_or(fresh), + Err(_) => fresh, + }; + if record.appended > 0 { + eprintln!( + " resume: {} of {} parts already assembled", + record.appended, + record.parts.len() + ); + // Cut off a torn tail from an interrupted append. + OpenOptions::new() + .write(true) + .open(&partial) + .and_then(|f| f.set_len(record.len)) + .with_context(|| format!("truncating {}", partial.display()))?; + } else { + let _ = std::fs::remove_file(&partial); + let _ = std::fs::remove_file(&record_path); + } + let resumed = record.appended > 0; + Ok(Self { + dst: dst.to_path_buf(), + partial, + record_path, + record, + manifest, + resumed, + whole: Sha256::new(), + }) + } + + fn run(mut self, fetcher: &dyn Fetcher, repo: &str, revision: &str) -> Result<()> { + let total = self.record.parts.len(); + for i in self.record.appended..total { + let part = self.record.parts[i].clone(); + let src = fetcher.get_file(repo, revision, &part)?; + let (digest, got) = append( + &src, + &self.partial, + &format!("assemble {}/{total}", i + 1), + &mut self.whole, + )?; + let want = &self.manifest.parts[i]; + let problem = if got != want.size { + Some(format!( + "{part} is {got} bytes but the manifest says {}", + want.size + )) + } else if !digest.eq_ignore_ascii_case(&want.sha256) { + Some(format!( + "{part} hashed to {digest} but the manifest says {}", + want.sha256 + )) + } else { + None + }; + if let Some(problem) = problem { + // Neither the staged part nor the partial can be trusted + // past this point; the next attempt rebuilds. The staged + // bytes go only if they are ours to drop. + self.discard(); + discard_staged(fetcher, repo, &src); + bail!("{problem}: the download was corrupted or the part was republished; run the pull again"); + } + discard_staged(fetcher, repo, &src); + self.record.appended = i + 1; + self.record.len = std::fs::metadata(&self.partial) + .with_context(|| format!("sizing {}", self.partial.display()))? + .len(); + self.save()?; + } + + // Every part matched the manifest, so the whole is the manifest's + // size by construction; say so explicitly rather than trust it. + if self.record.len != self.manifest.size { + self.discard(); + bail!( + "{total} parts reassemble to {} bytes but the manifest says {}", + self.record.len, + self.manifest.size + ); + } + // The whole has to be what the manifest says it is, not only each + // part: a manifest with right parts and a wrong whole would + // otherwise pass here and fail a catalog pull after the last byte. + // A fresh install has hashed everything it appended; a resumed one + // inherited bytes an earlier attempt verified, which anything that + // touched the partial since (a torn write past the record, a disk + // fault, a stray tool) could have changed, so it reads them again. + let (kind, got) = if self.resumed { + ("resumed", sha256_file(&self.partial, "verify")?) + } else { + ("reassembled", hex(self.whole.clone().finalize().as_slice())) + }; + if !got.eq_ignore_ascii_case(&self.manifest.sha256) { + self.discard(); + bail!( + "the {kind} file hashed to {got} but the manifest says {}: the manifest is wrong or the partial was damaged; run the pull again", + self.manifest.sha256 + ); + } + // Second line of defence, against a manifest written from a broken + // set: the file is parsed as a `.base` and its weights blob and any + // slot section have to fit. + if let Err(e) = check_complete(&self.partial, self.record.len) { + self.discard(); + return Err(e.context(format!( + "{total} parts reassemble to {} bytes but the bundle is not complete: the part set on the Hub is missing its tail", + self.record.len + ))); + } + + std::fs::rename(&self.partial, &self.dst).with_context(|| { + format!( + "installing {} as {}", + self.partial.display(), + self.dst.display() + ) + })?; + let _ = std::fs::remove_file(&self.record_path); + Ok(()) + } + + fn save(&self) -> Result<()> { + let bytes = serde_json::to_vec(&self.record)?; + let tmp = with_suffix(&self.record_path, ".tmp"); + let mut f = File::create(&tmp)?; + f.write_all(&bytes)?; + f.sync_all()?; + std::fs::rename(&tmp, &self.record_path) + .with_context(|| format!("writing {}", self.record_path.display()))?; + Ok(()) + } + + /// Throw the partial away: what it holds has been shown to be wrong, so + /// there is no resume value in it. + fn discard(&self) { + let _ = std::fs::remove_file(&self.partial); + let _ = std::fs::remove_file(&self.record_path); + } +} + +fn is_sha256(id: &str) -> bool { + id.len() == 64 && id.bytes().all(|b| b.is_ascii_hexdigit()) +} + +/// Append `src` to `dst` in bounded memory, hashing the bytes on the way +/// through — into `whole` too, the running hash of the file being built; +/// returns the hex sha256 of `src` and how many bytes it was. Shows a bar +/// in the same style as the download that preceded it. +fn append(src: &Path, dst: &Path, label: &str, whole: &mut Sha256) -> Result<(String, u64)> { + let len = std::fs::metadata(src)?.len(); + let bar = ProgressBar::new(len); + bar.set_style(crate::fetch::bar_style()); + bar.set_message(label.to_string()); + let mut reader = BufReader::with_capacity(8 << 20, File::open(src)?); + let mut out = OpenOptions::new() + .append(true) + .create(true) + .open(dst) + .with_context(|| format!("opening {} for append", dst.display()))?; + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; 8 << 20]; + let mut total = 0u64; + loop { + let n = reader.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + whole.update(&buf[..n]); + out.write_all(&buf[..n])?; + total += n as u64; + bar.inc(n as u64); + } + out.sync_all()?; + bar.finish_and_clear(); + Ok((hex(hasher.finalize().as_slice()), total)) +} + +/// sha256 of a whole file, in bounded memory, with a progress bar. +fn sha256_file(path: &Path, label: &str) -> Result { + let len = std::fs::metadata(path)?.len(); + let bar = ProgressBar::new(len); + bar.set_style(crate::fetch::bar_style()); + bar.set_message(label.to_string()); + let mut reader = BufReader::with_capacity(8 << 20, File::open(path)?); + let mut hasher = Sha256::new(); + let mut buf = vec![0u8; 8 << 20]; + loop { + let n = reader.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + bar.inc(n as u64); + } + bar.finish_and_clear(); + Ok(hex(hasher.finalize().as_slice())) +} + +/// Drop a staged part that has been appended, so the staging tree never +/// holds more than the part in flight. Fixtures and shared caches the +/// fetcher does not own are left alone. +fn discard_staged(fetcher: &dyn Fetcher, repo: &str, src: &Path) { + let owned = fetcher + .staging_dir(repo) + .is_some_and(|dir| src.starts_with(&dir)); + if !owned { + return; + } + // hf-hub's snapshot path is a symlink into `blobs/`; remove both so the + // bytes actually go away. + if let Ok(real) = std::fs::canonicalize(src) { + let _ = std::fs::remove_file(&real); + } + let _ = std::fs::remove_file(src); +} + +/// Headers past this are corrupt or hostile: refuse rather than allocate. +/// The same bound the remote header scanner applies. +const MAX_HEADER_LEN: u64 = 256 * 1024 * 1024; + +/// Flags that promise an extension-slot section after the weights blob. +const SLOT_FLAGS: base_format::HeaderFlags = base_format::HeaderFlags::HAS_LORA + .union(base_format::HeaderFlags::HAS_SPECULATOR) + .union(base_format::HeaderFlags::HAS_COMPUTE_GRAPH) + .union(base_format::HeaderFlags::HAS_KV_WARMUP) + .union(base_format::HeaderFlags::HAS_TRACE_REF) + .union(base_format::HeaderFlags::ROPE_PRECOMPUTED); + +/// Is a stitched `.base` of `len` bytes everything its header describes? +/// +/// Two things have to hold. The weights blob must fit: the file is at least +/// the blob start plus the furthest tensor extent. And the slot section, +/// which follows the blob, must be whole: every slot record the count +/// promises lies inside the file, and when the header's flags advertise +/// slots there is at least one. A file cut exactly at the blob end with +/// slots the header never flagged is the one case nothing here can see. +/// +/// The header comes from an arbitrary repository, so its numbers are not +/// trusted: the allocation is bounded and every sum is checked. +fn check_complete(path: &Path, len: u64) -> Result<()> { + let mut f = File::open(path)?; + let mut prefix = [0u8; PREFIX_LEN as usize]; + f.read_exact(&mut prefix) + .context("shorter than a .base prefix")?; + if prefix[0..4] != MAGIC { + bail!("not a .base file (bad magic)"); + } + let version = u32::from_le_bytes(prefix[4..8].try_into().unwrap()); + if version != FORMAT_VERSION { + bail!("unsupported .base format version {version}"); + } + let header_len = u64::from_le_bytes(prefix[8..16].try_into().unwrap()); + if header_len == 0 || header_len >= MAX_HEADER_LEN { + bail!("implausible header length {header_len}"); + } + let mut json = vec![0u8; header_len as usize]; + f.read_exact(&mut json) + .context("file ends inside the header")?; + let header = base_format::Header::from_json_bytes(&json).context("parsing the header")?; + + let overflow = || anyhow::anyhow!("tensor extents overflow: the header is corrupt"); + let blob_offset = (PREFIX_LEN + header_len) + .checked_next_multiple_of(BLOB_ALIGNMENT) + .ok_or_else(overflow)?; + let mut blob_end = blob_offset; + let mmproj = header.mmproj.iter().flat_map(|m| m.tensors.iter()); + for t in header.tensors.iter().chain(mmproj) { + let regions = [ + (Some(t.offset), Some(t.length)), + (t.scale_offset, t.scale_length), + (t.bias_offset, t.bias_length), + (t.awq_scale_offset, t.awq_scale_length), + ]; + for (offset, length) in regions { + let (Some(o), Some(l)) = (offset, length) else { + continue; + }; + let end = o + .checked_add(l) + .and_then(|e| e.checked_add(blob_offset)) + .ok_or_else(overflow)?; + blob_end = blob_end.max(end); + } + } + if len < blob_end { + bail!("the weights blob needs {blob_end} bytes, the file has {len}"); + } + + let slots = walk_slots(&mut f, blob_end, len).context("reading the extension slots")?; + if header.flags.intersects(SLOT_FLAGS) && slots == 0 { + bail!( + "the header advertises extension slots (flags {:?}) but the file ends at the weights blob", + header.flags & SLOT_FLAGS + ); + } + Ok(()) +} + +/// Count the slot records after the blob, reading only their headers and +/// stepping over the payloads, so a multi-GB LoRA or speculator costs a +/// few seeks. Zero when nothing follows the blob. An error when the count +/// promises records the file does not hold. +fn walk_slots(f: &mut File, blob_end: u64, len: u64) -> Result { + use std::io::{Seek, SeekFrom}; + // The section starts at the first 8-byte boundary after the blob. + let start = blob_end + .checked_next_multiple_of(8) + .context("slot offset overflow")?; + if start >= len { + return Ok(0); + } + f.seek(SeekFrom::Start(start))?; + let mut buf4 = [0u8; 4]; + f.read_exact(&mut buf4).context("truncated slot count")?; + let n = u32::from_le_bytes(buf4); + let mut pos = start + 4; + for i in 0..n { + // u16 kind, u16 flags, u64 payload_length, u64 xxh64. + let mut rec = [0u8; 20]; + f.read_exact(&mut rec) + .with_context(|| format!("slot {i} of {n}: truncated record header"))?; + let payload_len = u64::from_le_bytes(rec[4..12].try_into().unwrap()); + let payload_end = pos + .checked_add(20) + .and_then(|p| p.checked_add(payload_len)) + .context("slot payload overflow")?; + if payload_end > len { + bail!("slot {i} of {n}: payload runs to byte {payload_end}, the file has {len}"); + } + // The writer pads every payload to 8 bytes, the last one included, + // and the canonical reader consumes that padding with read_exact. + pos = payload_end + .checked_next_multiple_of(8) + .context("slot padding overflow")?; + if pos > len { + bail!("slot {i} of {n}: padding runs to byte {pos}, the file has {len}"); + } + f.seek(SeekFrom::Start(pos))?; + } + Ok(n) +} + +/// A real, tiny `.base`: a valid header describing one f32 tensor of +/// `blob_len` bytes, then that blob. Test support for anything that has to +/// feed [`install`] bytes it will accept, which is why it is public: the +/// pull command's tests live in another crate. +#[doc(hidden)] +pub fn synthetic_bundle(blob_len: u64) -> Vec { + synthetic_bundle_with_slot(blob_len, None, false) +} + +/// [`synthetic_bundle`] with, optionally, one extension slot holding +/// `slot` after the blob, and the header flagged as carrying LoRA when +/// `flagged` (whether or not a slot is actually written). +#[doc(hidden)] +pub fn synthetic_bundle_with_slot(blob_len: u64, slot: Option<&[u8]>, flagged: bool) -> Vec { + let header = serde_json::json!({ + "schema": 1, + "arch": "test", + "quant_scheme": "base_q4", + "min_hw": "apple_m1", + "created": "2026-09-08T00:00:00Z", + "baserT_version": "0.2.4", + "source": { "format": "test", "sha256": "0".repeat(64), "filename": "x" }, + "tokenizer": {}, + "config": {}, + "flags": if flagged { "0x00000010" } else { "0x00000000" }, + "tensors": [{ + "name": "w", "dtype": "f32", "shape": [blob_len / 4], + "offset": 0, "length": blob_len + }] + }) + .to_string(); + let mut v = MAGIC.to_vec(); + v.extend_from_slice(&FORMAT_VERSION.to_le_bytes()); + v.extend_from_slice(&(header.len() as u64).to_le_bytes()); + v.extend_from_slice(header.as_bytes()); + let blob_offset = (v.len() as u64).div_ceil(BLOB_ALIGNMENT) * BLOB_ALIGNMENT; + v.resize(blob_offset as usize, 0); + v.extend((0..blob_len).map(|i| (i % 251) as u8)); + if let Some(payload) = slot { + // Slot section: pad to 8, u32 count, then one record (kind, flags, + // payload length, xxh64 of zero = unchecked) and its payload. + v.resize((v.len() as u64).div_ceil(8) as usize * 8, 0); + v.extend_from_slice(&1u32.to_le_bytes()); + v.extend_from_slice(&0x0001u16.to_le_bytes()); + v.extend_from_slice(&0u16.to_le_bytes()); + v.extend_from_slice(&(payload.len() as u64).to_le_bytes()); + v.extend_from_slice(&0u64.to_le_bytes()); + v.extend_from_slice(payload); + v.resize((v.len() as u64).div_ceil(8) as usize * 8, 0); + } + v +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::fetch::MockFetcher; + + fn names(v: &[&str]) -> Vec { + v.iter().map(|s| s.to_string()).collect() + } + + #[test] + fn part_names_parse_and_reject() { + assert_eq!( + split_part_name("parts/GLM-5.2-Q4.base.part-007"), + Some(("parts/GLM-5.2-Q4.base".to_string(), 7)) + ); + assert_eq!(split_part_name("m-Q4.base"), None); + assert_eq!(split_part_name("m-Q4.base.part-"), None); + assert_eq!(split_part_name("m-Q4.base.part-1a"), None); + assert_eq!(split_part_name("m-Q4.safetensors.part-001"), None); + // The manifest sits beside the parts and is neither. + assert_eq!(split_part_name("m-Q4.base.manifest.json"), None); + } + + #[test] + fn group_orders_parts_and_keeps_whole_files() { + // Listing order is not byte order; the group must be. + let got = group(names(&[ + "README.md", + "parts/g.base.part-002", + "m-Q8.base", + "parts/g.base.part-000", + "parts/g.base.manifest.json", + "parts/g.base.part-001", + ])); + assert_eq!( + got.artifacts, + vec![ + Artifact::whole("m-Q8.base"), + Artifact { + name: "parts/g.base".into(), + parts: names(&[ + "parts/g.base.part-000", + "parts/g.base.part-001", + "parts/g.base.part-002" + ]), + }, + ] + ); + assert!(got.malformed.is_empty()); + } + + #[test] + fn group_isolates_a_gapped_set_and_prefers_a_whole_file() { + // A quant mid-upload must not take the repo's other bundles with it. + let got = group(names(&[ + "g.base.part-000", + "g.base.part-002", + "m-Q8.base", + "h.base.part-000", + ])); + assert_eq!( + got.artifacts, + vec![ + Artifact::whole("m-Q8.base"), + Artifact { + name: "h.base".into(), + parts: names(&["h.base.part-000"]), + }, + ] + ); + assert_eq!(got.malformed.len(), 1); + assert_eq!(got.malformed[0].0, "g.base"); + assert!( + got.malformed[0].1.contains("missing part 001"), + "{}", + got.malformed[0].1 + ); + + let got = group(names(&["g.base", "g.base.part-000", "g.base.part-001"])); + assert_eq!(got.artifacts, vec![Artifact::whole("g.base")]); + assert!(got.malformed.is_empty()); + } + + #[test] + fn group_rejects_two_spellings_of_one_index() { + let got = group(names(&[ + "g.base.part-000", + "g.base.part-0000", + "g.base.part-001", + "m-Q8.base", + ])); + assert_eq!(got.artifacts, vec![Artifact::whole("m-Q8.base")]); + assert_eq!(got.malformed.len(), 1); + assert_eq!(got.malformed[0].0, "g.base"); + assert!( + got.malformed[0].1.contains("index 0 is listed twice"), + "{}", + got.malformed[0].1 + ); + } + + #[test] + fn find_names_the_malformed_set_it_was_asked_for() { + let tmp = tempfile::tempdir().unwrap(); + let repo_dir = tmp.path().join("org").join("m"); + std::fs::create_dir_all(&repo_dir).unwrap(); + for f in ["m-Q4.base.part-000", "m-Q4.base.part-002", "m-Q8.base"] { + std::fs::write(repo_dir.join(f), b"x").unwrap(); + } + let fetcher = MockFetcher::new(tmp.path()); + let err = find(&fetcher, "org/m", "main", "m-Q4.base") + .unwrap_err() + .to_string(); + assert!(err.contains("missing part 001"), "{err}"); + // The sibling is unaffected. + assert_eq!( + find(&fetcher, "org/m", "main", "m-Q8.base").unwrap(), + Artifact::whole("m-Q8.base") + ); + } + + /// Cut `bytes` into `n` parts of roughly equal size. + fn split(bytes: &[u8], n: usize) -> Vec> { + let each = bytes.len().div_ceil(n); + bytes.chunks(each).map(<[u8]>::to_vec).collect() + } + + fn sha(bytes: &[u8]) -> String { + hex(Sha256::digest(bytes).as_slice()) + } + + /// A mock repo `org/m` holding `m-Q4.base.part-NNN` for each of `parts`, + /// with a manifest built from exactly those files — what an honest + /// publisher ships. + fn fixture_repo(tmp: &Path, parts: &[Vec]) -> (MockFetcher, PathBuf) { + let (fetcher, repo_dir) = fixture_repo_bare(tmp, parts); + let paths: Vec = (0..parts.len()) + .map(|i| repo_dir.join(format!("m-Q4.base.part-{i:03}"))) + .collect(); + let manifest = build_manifest(&paths).unwrap(); + write_manifest(&repo_dir, &manifest); + (fetcher, repo_dir) + } + + /// The same repo with no manifest at all. + fn fixture_repo_bare(tmp: &Path, parts: &[Vec]) -> (MockFetcher, PathBuf) { + let repo_dir = tmp.join("org").join("m"); + std::fs::create_dir_all(&repo_dir).unwrap(); + for (i, bytes) in parts.iter().enumerate() { + std::fs::write(repo_dir.join(format!("m-Q4.base.part-{i:03}")), bytes).unwrap(); + } + (MockFetcher::new(tmp), repo_dir) + } + + fn write_manifest(repo_dir: &Path, manifest: &Manifest) { + std::fs::write( + repo_dir.join(manifest_name("m-Q4.base")), + serde_json::to_vec_pretty(manifest).unwrap(), + ) + .unwrap(); + } + + fn dst_in(tmp: &Path) -> PathBuf { + let dst = tmp.join("out").join("model.base"); + std::fs::create_dir_all(dst.parent().unwrap()).unwrap(); + dst + } + + fn record_for(art: &Artifact, manifest: &Manifest, appended: usize, len: u64) -> Record { + Record { + repo: "org/m".into(), + revision: "main".into(), + parts: art.parts.clone(), + ids: manifest.part_ids(), + appended, + len, + } + } + + fn manifest_of(parts: &[Vec], tmp: &Path) -> Manifest { + let dir = tmp.join("manifest-src"); + std::fs::create_dir_all(&dir).unwrap(); + let paths: Vec = parts + .iter() + .enumerate() + .map(|(i, b)| { + let p = dir.join(format!("m-Q4.base.part-{i:03}")); + std::fs::write(&p, b).unwrap(); + p + }) + .collect(); + build_manifest(&paths).unwrap() + } + + #[test] + fn build_manifest_hashes_each_part_and_the_whole_in_one_pass() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let m = manifest_of(&parts, tmp.path()); + assert_eq!(m.size, whole.len() as u64); + assert_eq!(m.sha256, sha(&whole)); + assert_eq!(m.parts.len(), 3); + for (i, p) in parts.iter().enumerate() { + assert_eq!(m.parts[i].name, format!("m-Q4.base.part-{i:03}")); + assert_eq!(m.parts[i].size, p.len() as u64); + assert_eq!(m.parts[i].sha256, sha(p)); + } + m.check_against(&names(&[ + "parts/m-Q4.base.part-000", + "parts/m-Q4.base.part-001", + "parts/m-Q4.base.part-002", + ])) + .unwrap(); + let err = m + .check_against(&names(&["m-Q4.base.part-000", "m-Q4.base.part-001"])) + .unwrap_err() + .to_string(); + assert!(err.contains("lists 3 parts but the repo has 2"), "{err}"); + + // A whole-file hash that is not one is refused up front. + let mut bad = m.clone(); + bad.sha256 = "not-a-digest".into(); + let err = bad + .check_against(&names(&[ + "m-Q4.base.part-000", + "m-Q4.base.part-001", + "m-Q4.base.part-002", + ])) + .unwrap_err() + .to_string(); + assert!(err.contains("whole-file sha256"), "{err}"); + + // Sizes are compared one by one, not only as a sum. + let sizes: Vec = parts.iter().map(|p| p.len() as u64).collect(); + m.check_sizes(&sizes).unwrap(); + let mut swapped = sizes.clone(); + swapped[0] += 1; + swapped[1] -= 1; + let err = m.check_sizes(&swapped).unwrap_err().to_string(); + assert!(err.contains("bytes but the repo holds"), "{err}"); + } + + #[test] + fn install_reassembles_parts_in_order() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let (fetcher, repo_dir) = fixture_repo(tmp.path(), &split(&whole, 3)); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + assert_eq!(art.parts.len(), 3); + + let dst = dst_in(tmp.path()); + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + + assert_eq!(std::fs::read(&dst).unwrap(), whole); + assert!(!partial_path(&dst).exists(), "partial must be renamed away"); + assert!(!record_path(&dst).exists(), "record must be cleared"); + assert!(!lock_path(&dst).exists(), "lock must be cleared"); + // Fixtures are not owned by the mock fetcher: never deleted. + assert!(repo_dir.join("m-Q4.base.part-002").exists()); + } + + #[test] + fn install_needs_a_manifest() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let (fetcher, _) = fixture_repo_bare(tmp.path(), &split(&whole, 2)); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let err = format!( + "{:#}", + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap_err() + ); + assert!(err.contains("needs its manifest"), "{err}"); + assert!(!dst.exists()); + } + + #[test] + fn install_refuses_an_oversized_manifest_before_reading_it() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let (fetcher, repo_dir) = fixture_repo(tmp.path(), &split(&whole, 2)); + let big = vec![b' '; MAX_MANIFEST_LEN as usize + 1]; + std::fs::write(repo_dir.join(manifest_name("m-Q4.base")), &big).unwrap(); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let err = format!( + "{:#}", + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap_err() + ); + assert!(err.contains("implausible manifest size"), "{err}"); + } + + #[test] + fn install_resumes_from_the_recorded_part() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let (fetcher, _) = fixture_repo(tmp.path(), &parts); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let manifest = manifest_of(&parts, tmp.path()); + let dst = dst_in(tmp.path()); + + // Two parts landed, then the append of the third tore mid-way. + let two: Vec = parts[..2].concat(); + let mut torn = two.clone(); + torn.extend_from_slice(b"XX"); + std::fs::write(partial_path(&dst), &torn).unwrap(); + let rec = record_for(&art, &manifest, 2, two.len() as u64); + std::fs::write(record_path(&dst), serde_json::to_vec(&rec).unwrap()).unwrap(); + + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + #[test] + fn install_checks_the_whole_hash_on_a_fresh_install_too() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let (fetcher, repo_dir) = fixture_repo(tmp.path(), &parts); + // Right parts, wrong whole: a manifest edited by hand, say. + let mut m = manifest_of(&parts, tmp.path()); + m.sha256 = sha(b"not the whole"); + write_manifest(&repo_dir, &m); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let err = install(&fetcher, "org/m", "main", &art, &dst, None) + .unwrap_err() + .to_string(); + assert!(err.contains("reassembled file hashed to"), "{err}"); + assert!(!dst.exists()); + assert!(!partial_path(&dst).exists()); + } + + #[test] + fn install_rehashes_a_resumed_partial_against_the_manifest() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let (fetcher, _) = fixture_repo(tmp.path(), &parts); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let manifest = manifest_of(&parts, tmp.path()); + let dst = dst_in(tmp.path()); + + // Two parts were appended by an earlier attempt; since then one + // byte inside them changed without changing the length, which the + // record alone cannot see. + let mut two: Vec = parts[..2].concat(); + two[100] ^= 0xff; + std::fs::write(partial_path(&dst), &two).unwrap(); + let rec = record_for(&art, &manifest, 2, two.len() as u64); + std::fs::write(record_path(&dst), serde_json::to_vec(&rec).unwrap()).unwrap(); + + let err = install(&fetcher, "org/m", "main", &art, &dst, None) + .unwrap_err() + .to_string(); + assert!(err.contains("resumed file hashed to"), "{err}"); + assert!(!dst.exists()); + assert!( + !partial_path(&dst).exists(), + "a damaged partial is not kept" + ); + + // The next attempt starts clean and succeeds. + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + #[test] + fn install_restarts_when_the_record_describes_other_parts() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 2); + let (fetcher, _) = fixture_repo(tmp.path(), &parts); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let manifest = manifest_of(&parts, tmp.path()); + let dst = dst_in(tmp.path()); + + std::fs::write(partial_path(&dst), b"stale").unwrap(); + let rec = Record { + parts: names(&["other.base.part-000", "other.base.part-001"]), + ..record_for(&art, &manifest, 1, 5) + }; + std::fs::write(record_path(&dst), serde_json::to_vec(&rec).unwrap()).unwrap(); + + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + /// Mock whose parts carry content ids and whose branch resolves to a + /// commit, like the Hub. + struct HubLikeFetcher { + inner: MockFetcher, + ids: Vec>, + commit: String, + } + + impl Fetcher for HubLikeFetcher { + fn get_file(&self, repo: &str, revision: &str, filename: &str) -> Result { + assert_eq!( + revision, self.commit, + "files must be fetched at the pinned commit" + ); + self.inner.get_file(repo, revision, filename) + } + fn list_files(&self, repo: &str, revision: &str) -> Result> { + self.inner.list_files(repo, revision) + } + fn resolve_revision(&self, _: &str, _: &str) -> Result { + Ok(self.commit.clone()) + } + fn content_id(&self, _: &str, revision: &str, filename: &str) -> Result> { + assert_eq!(revision, self.commit); + let (_, idx) = split_part_name(filename).unwrap(); + Ok(self.ids[idx as usize].clone()) + } + } + + fn hub_like(tmp: &Path, parts: &[Vec]) -> HubLikeFetcher { + let (inner, _) = fixture_repo(tmp, parts); + HubLikeFetcher { + inner, + ids: parts.iter().map(|p| Some(sha(p))).collect(), + commit: "c0ffee".into(), + } + } + + #[test] + fn install_pins_the_revision_and_verifies_each_part() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let fetcher = hub_like(tmp.path(), &parts); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + /// Mock whose branch listing and pinned-commit listing differ: the + /// branch gained a part after the caller listed it. + struct AdvancingFetcher { + inner: MockFetcher, + commit: String, + } + + impl Fetcher for AdvancingFetcher { + fn get_file(&self, repo: &str, revision: &str, filename: &str) -> Result { + assert_eq!(revision, self.commit); + self.inner.get_file(repo, revision, filename) + } + fn list_files(&self, repo: &str, revision: &str) -> Result> { + let all = self.inner.list_files(repo, revision)?; + if revision == self.commit { + Ok(all) + } else { + // The branch, as the caller saw it: one part short. + Ok(all + .into_iter() + .filter(|f| !f.ends_with("part-002")) + .collect()) + } + } + fn resolve_revision(&self, _: &str, _: &str) -> Result { + Ok(self.commit.clone()) + } + } + + /// Mock whose branch listing shows a whole file where the pinned + /// commit holds a part set: the publication changed shape after the + /// caller listed. + struct ResplitFetcher { + inner: MockFetcher, + commit: String, + } + + impl Fetcher for ResplitFetcher { + fn get_file(&self, repo: &str, revision: &str, filename: &str) -> Result { + assert_eq!(revision, self.commit); + self.inner.get_file(repo, revision, filename) + } + fn list_files(&self, repo: &str, revision: &str) -> Result> { + if revision == self.commit { + self.inner.list_files(repo, revision) + } else { + Ok(vec!["m-Q4.base".to_string()]) + } + } + fn resolve_revision(&self, _: &str, _: &str) -> Result { + Ok(self.commit.clone()) + } + } + + #[test] + fn install_decides_whole_or_split_at_the_pinned_commit() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let (inner, _) = fixture_repo(tmp.path(), &split(&whole, 3)); + let fetcher = ResplitFetcher { + inner, + commit: "c0ffee".into(), + }; + // The caller saw a whole file on the branch; the commit has parts. + let seen = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + assert!(!seen.is_split()); + let dst = dst_in(tmp.path()); + install(&fetcher, "org/m", "main", &seen, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + #[test] + fn install_relists_the_parts_at_the_pinned_commit() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let (inner, _) = fixture_repo(tmp.path(), &split(&whole, 3)); + let fetcher = AdvancingFetcher { + inner, + commit: "c0ffee".into(), + }; + // Listed from the branch: two parts. Installed from the pinned + // commit: all three, so the bundle is whole. + let stale = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + assert_eq!(stale.parts.len(), 2); + let dst = dst_in(tmp.path()); + install(&fetcher, "org/m", "main", &stale, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + #[test] + fn install_refuses_a_manifest_the_hub_disagrees_with() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let mut fetcher = hub_like(tmp.path(), &parts); + // The Hub's LFS hash for part 1 is not what the manifest says: a + // part was republished without republishing the manifest. + fetcher.ids[1] = Some(sha(b"something else")); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let err = install(&fetcher, "org/m", "main", &art, &dst, None) + .unwrap_err() + .to_string(); + assert!(err.contains("manifest is stale"), "{err}"); + assert!(!dst.exists()); + assert!(!partial_path(&dst).exists(), "refused before any byte"); + } + + #[test] + fn install_rejects_a_part_whose_bytes_do_not_match_the_manifest() { + // Corruption the listing cannot see (the mock reports no ids): + // the bytes that arrive differ from what the manifest promised. + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let (fetcher, repo_dir) = fixture_repo(tmp.path(), &parts); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + + // Same length, one byte flipped. + let mut flipped = parts[1].clone(); + flipped[10] ^= 0xff; + std::fs::write(repo_dir.join("m-Q4.base.part-001"), &flipped).unwrap(); + let err = install(&fetcher, "org/m", "main", &art, &dst, None) + .unwrap_err() + .to_string(); + assert!(err.contains("hashed to"), "{err}"); + assert!( + !partial_path(&dst).exists(), + "a bad partial has no resume value" + ); + + // Shorter than promised. + std::fs::write(repo_dir.join("m-Q4.base.part-001"), &parts[1][..100]).unwrap(); + let err = install(&fetcher, "org/m", "main", &art, &dst, None) + .unwrap_err() + .to_string(); + assert!(err.contains("bytes but the manifest says"), "{err}"); + assert!(!dst.exists()); + } + + #[test] + fn install_restarts_when_a_part_was_republished_under_the_same_name() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let fetcher = hub_like(tmp.path(), &parts); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let manifest = manifest_of(&parts, tmp.path()); + let dst = dst_in(tmp.path()); + + // Two parts were assembled from the previous publication, whose + // part 000 had different bytes. Splicing the new tail onto them + // would install a bundle nobody published. + let mut old_ids = manifest.part_ids(); + old_ids[0] = sha(b"previous part 0"); + std::fs::write(partial_path(&dst), b"OLD-HEAD").unwrap(); + let rec = Record { + revision: fetcher.commit.clone(), + ids: old_ids, + ..record_for(&art, &manifest, 2, 8) + }; + std::fs::write(record_path(&dst), serde_json::to_vec(&rec).unwrap()).unwrap(); + + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + + // Same ids and commit: the record is honored and the tail appended. + std::fs::remove_file(&dst).unwrap(); + let two: Vec = parts[..2].concat(); + std::fs::write(partial_path(&dst), &two).unwrap(); + let rec = Record { + ids: manifest.part_ids(), + len: two.len() as u64, + ..rec + }; + std::fs::write(record_path(&dst), serde_json::to_vec(&rec).unwrap()).unwrap(); + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + #[test] + fn install_checks_the_manifest_against_a_catalog_row() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + let fetcher = hub_like(tmp.path(), &parts); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let pinned: Vec = parts.iter().map(|p| sha(p)).collect(); + + // Agreeing row: installs. + install(&fetcher, "org/m", "main", &art, &dst, Some(&pinned)).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + std::fs::remove_file(&dst).unwrap(); + + // The catalog pinned a different part 2: refused before any byte. + let mut stale = pinned.clone(); + stale[2] = sha(b"older publication"); + let err = install(&fetcher, "org/m", "main", &art, &dst, Some(&stale)) + .unwrap_err() + .to_string(); + assert!(err.contains("republished"), "{err}"); + assert!(!partial_path(&dst).exists()); + + // A different part count is the same story. + let err = install(&fetcher, "org/m", "main", &art, &dst, Some(&pinned[..2])) + .unwrap_err() + .to_string(); + assert!(err.contains("pins 2 parts"), "{err}"); + } + + #[test] + fn install_refuses_a_part_set_missing_its_tail() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + // The publisher wrote the manifest for three parts and uploaded + // two: gap-free, and nothing in the names says a third was meant + // to exist. The manifest does. + let (fetcher, repo_dir) = fixture_repo(tmp.path(), &parts); + std::fs::remove_file(repo_dir.join("m-Q4.base.part-002")).unwrap(); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + assert_eq!(art.parts.len(), 2); + let dst = dst_in(tmp.path()); + let err = format!( + "{:#}", + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap_err() + ); + assert!(err.contains("lists 3 parts but the repo has 2"), "{err}"); + assert!(!dst.exists(), "a truncated bundle must not be installed"); + assert!(!partial_path(&dst).exists()); + } + + #[test] + fn install_refuses_a_slot_only_tail_the_header_cannot_see() { + // A bundle ending in an unflagged slot (calibration data, say), + // cut so the last part is exactly that slot section, and that part + // never uploaded. The header is silent about the slot; only the + // manifest knows a third part existed. + let tmp = tempfile::tempdir().unwrap(); + let payload = vec![7u8; 300]; + let whole = synthetic_bundle_with_slot(1000, Some(&payload), false); + let no_slot = synthetic_bundle_with_slot(1000, None, false); + let mut parts = split(&no_slot, 2); + parts.push(whole[no_slot.len()..].to_vec()); + assert_eq!(parts.concat(), whole); + let (fetcher, repo_dir) = fixture_repo(tmp.path(), &parts); + std::fs::remove_file(repo_dir.join("m-Q4.base.part-002")).unwrap(); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let err = format!( + "{:#}", + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap_err() + ); + assert!(err.contains("lists 3 parts but the repo has 2"), "{err}"); + assert!(!dst.exists()); + } + + #[test] + fn install_refuses_a_bundle_whose_advertised_slots_were_in_the_lost_tail() { + // The header-level defence, for a manifest written from an already + // broken set: the flags promise slots and the file ends at the blob. + let tmp = tempfile::tempdir().unwrap(); + let no_slot = synthetic_bundle_with_slot(1000, None, true); + let (fetcher, _) = fixture_repo(tmp.path(), &split(&no_slot, 2)); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let err = format!( + "{:#}", + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap_err() + ); + assert!(err.contains("advertises extension slots"), "{err}"); + assert!(!dst.exists()); + + // With the slot present, the same bundle installs. + let tmp = tempfile::tempdir().unwrap(); + let payload = vec![7u8; 300]; + let whole = synthetic_bundle_with_slot(1000, Some(&payload), true); + let (fetcher, _) = fixture_repo(tmp.path(), &split(&whole, 3)); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + #[test] + fn install_refuses_a_slot_section_missing_only_its_padding() { + let tmp = tempfile::tempdir().unwrap(); + // A 301-byte payload is followed by 3 pad bytes; lose just those. + let payload = vec![7u8; 301]; + let whole = synthetic_bundle_with_slot(1000, Some(&payload), false); + assert_eq!(whole.len() % 8, 0); + let cut = &whole[..whole.len() - 3]; + let (fetcher, _) = fixture_repo(tmp.path(), &split(cut, 2)); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let err = format!( + "{:#}", + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap_err() + ); + assert!(err.contains("padding runs to byte"), "{err}"); + assert!(!dst.exists()); + } + + #[test] + fn install_refuses_a_slot_section_cut_short_even_when_unflagged() { + let tmp = tempfile::tempdir().unwrap(); + let payload = vec![7u8; 300]; + let whole = synthetic_bundle_with_slot(1000, Some(&payload), false); + // Lose the last 100 bytes: inside the slot payload. + let cut = &whole[..whole.len() - 100]; + let (fetcher, _) = fixture_repo(tmp.path(), &split(cut, 2)); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + let err = format!( + "{:#}", + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap_err() + ); + assert!(err.contains("payload runs to byte"), "{err}"); + assert!(!dst.exists()); + } + + #[test] + fn check_complete_bounds_the_header_and_checks_its_sums() { + let tmp = tempfile::tempdir().unwrap(); + // A prefix claiming a 1 GiB header on a 16-byte file: refused + // before any allocation. + let mut v = MAGIC.to_vec(); + v.extend_from_slice(&FORMAT_VERSION.to_le_bytes()); + v.extend_from_slice(&(1u64 << 30).to_le_bytes()); + let p = tmp.path().join("huge.base"); + std::fs::write(&p, &v).unwrap(); + let err = check_complete(&p, v.len() as u64).unwrap_err().to_string(); + assert!(err.contains("implausible header length"), "{err}"); + + // A tensor whose offset + length wraps u64: reported, not wrapped + // into a small requirement. + let mut whole = synthetic_bundle(64); + let json_start = PREFIX_LEN as usize; + let header_len = u64::from_le_bytes(whole[8..16].try_into().unwrap()) as usize; + let json = String::from_utf8(whole[json_start..json_start + header_len].to_vec()).unwrap(); + let bad = json.replace("\"offset\":0", &format!("\"offset\":{}", u64::MAX - 8)); + assert_ne!(bad, json); + let mut rebuilt = whole[..8].to_vec(); + rebuilt.extend_from_slice(&(bad.len() as u64).to_le_bytes()); + rebuilt.extend_from_slice(bad.as_bytes()); + let blob_offset = (rebuilt.len() as u64).div_ceil(BLOB_ALIGNMENT) * BLOB_ALIGNMENT; + rebuilt.resize(blob_offset as usize, 0); + rebuilt.extend_from_slice(&whole.split_off(whole.len() - 64)); + let p = tmp.path().join("wrap.base"); + std::fs::write(&p, &rebuilt).unwrap(); + let err = check_complete(&p, rebuilt.len() as u64) + .unwrap_err() + .to_string(); + assert!(err.contains("overflow"), "{err}"); + } + + #[test] + fn install_is_exclusive_per_destination() { + let tmp = tempfile::tempdir().unwrap(); + let whole = synthetic_bundle(1000); + let (fetcher, _) = fixture_repo(tmp.path(), &split(&whole, 2)); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = dst_in(tmp.path()); + + // Another pull holds the destination. + let other = Lock::acquire(&dst).unwrap(); + let err = install(&fetcher, "org/m", "main", &art, &dst, None) + .unwrap_err() + .to_string(); + assert!(err.contains("another basert pull"), "{err}"); + other.release(); + + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), whole); + } + + #[test] + fn install_of_a_whole_file_is_a_plain_copy() { + let tmp = tempfile::tempdir().unwrap(); + let repo_dir = tmp.path().join("org").join("m"); + std::fs::create_dir_all(&repo_dir).unwrap(); + std::fs::write(repo_dir.join("m-Q4.base"), b"whole").unwrap(); + let fetcher = MockFetcher::new(tmp.path()); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + assert!(!art.is_split()); + let dst = tmp.path().join("model.base"); + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + assert_eq!(std::fs::read(&dst).unwrap(), b"whole"); + } + + /// Owned staging tree in hf-hub's shape, so the space-bounding deletes + /// can be observed. + struct StagedFetcher { + staging: PathBuf, + } + + impl StagedFetcher { + fn repo_dir(&self, repo: &str) -> PathBuf { + self.staging + .join(format!("models--{}", repo.replace('/', "--"))) + } + + fn stage(&self, repo: &str, revision: &str, filename: &str, bytes: &[u8]) { + let rdir = self.repo_dir(repo); + let blobs = rdir.join("blobs"); + let snap = rdir.join("snapshots").join(revision); + std::fs::create_dir_all(&blobs).unwrap(); + std::fs::create_dir_all(&snap).unwrap(); + let blob = blobs.join(format!("etag-{filename}")); + std::fs::write(&blob, bytes).unwrap(); + std::os::unix::fs::symlink(&blob, snap.join(filename)).unwrap(); + } + } + + impl Fetcher for StagedFetcher { + fn get_file(&self, repo: &str, revision: &str, filename: &str) -> Result { + let p = self + .repo_dir(repo) + .join("snapshots") + .join(revision) + .join(filename); + anyhow::ensure!(p.exists(), "not staged: {}", p.display()); + Ok(p) + } + + fn list_files(&self, repo: &str, revision: &str) -> Result> { + let dir = self.repo_dir(repo).join("snapshots").join(revision); + let mut out = Vec::new(); + for e in std::fs::read_dir(dir)? { + out.push(e?.file_name().to_string_lossy().into_owned()); + } + Ok(out) + } + + fn staging_dir(&self, repo: &str) -> Option { + Some(self.repo_dir(repo)) + } + } + + #[test] + fn install_from_staging_frees_each_part_as_it_lands() { + let tmp = tempfile::tempdir().unwrap(); + let fetcher = StagedFetcher { + staging: tmp.path().join("staging"), + }; + let whole = synthetic_bundle(1000); + let parts = split(&whole, 3); + for (i, b) in parts.iter().enumerate() { + fetcher.stage("org/m", "main", &format!("m-Q4.base.part-{i:03}"), b); + } + let manifest = manifest_of(&parts, tmp.path()); + fetcher.stage( + "org/m", + "main", + &manifest_name("m-Q4.base"), + &serde_json::to_vec(&manifest).unwrap(), + ); + let art = find(&fetcher, "org/m", "main", "m-Q4.base").unwrap(); + let dst = tmp.path().join("model.base"); + install(&fetcher, "org/m", "main", &art, &dst, None).unwrap(); + + assert_eq!(std::fs::read(&dst).unwrap(), whole); + let blobs = fetcher.repo_dir("org/m").join("blobs"); + let left: Vec = std::fs::read_dir(&blobs) + .unwrap() + .flatten() + .map(|e| e.file_name().to_string_lossy().into_owned()) + .filter(|n| !n.contains("manifest")) + .collect(); + assert!( + left.is_empty(), + "every staged part must be consumed, found {left:?}" + ); + } +} diff --git a/base-convert/crates/base-hub/src/registry.rs b/base-convert/crates/base-hub/src/registry.rs index 87cd815..1d8ebcb 100644 --- a/base-convert/crates/base-hub/src/registry.rs +++ b/base-convert/crates/base-hub/src/registry.rs @@ -63,6 +63,7 @@ pub enum ModelRef { arch: Option, size: Option, sha256: Option, + parts_sha256: Option>, }, /// Arbitrary HF repo of source safetensors — download and convert locally. HuggingFace { @@ -273,6 +274,7 @@ impl CatalogRegistry { arch: e.arch.clone(), size: e.size, sha256: e.sha256.clone(), + parts_sha256: e.parts_sha256.clone(), } } @@ -325,7 +327,11 @@ impl CatalogRegistry { /// fall through to a convert-on-pull of a bundle this client can't run. When /// the id+quant simply isn't published, both fields are `(None, false)` and /// the caller may convert-on-pull from the source repo. - pub fn resolve_with_status(&self, id: &str, want_quant: Option<&str>) -> (Option, bool) { + pub fn resolve_with_status( + &self, + id: &str, + want_quant: Option<&str>, + ) -> (Option, bool) { let is_exact = |e: &crate::catalog::CatalogEntry| e.id == id; let is_id = |e: &crate::catalog::CatalogEntry| e.id == id || e.id.eq_ignore_ascii_case(id); @@ -351,13 +357,20 @@ impl CatalogRegistry { Some(w) => Self::quant_matches(&e.quant, w), }; - let matched: Vec<&crate::catalog::CatalogEntry> = - self.catalog.models.iter().filter(|e| is_id(e) && quant_ok(e)).collect(); + let matched: Vec<&crate::catalog::CatalogEntry> = self + .catalog + .models + .iter() + .filter(|e| is_id(e) && quant_ok(e)) + .collect(); if matched.is_empty() { return (None, false); // this id+quant isn't published — convert-on-pull } - let mut runnable: Vec<&crate::catalog::CatalogEntry> = - matched.iter().copied().filter(|e| Self::backend_ok(e)).collect(); + let mut runnable: Vec<&crate::catalog::CatalogEntry> = matched + .iter() + .copied() + .filter(|e| Self::backend_ok(e)) + .collect(); if runnable.is_empty() { // Published, but only for a backend this client can't run. if let Some(e) = matched.first() { @@ -551,7 +564,10 @@ mod tests { let tmp = tempfile::tempdir().unwrap(); let reg = MergedRegistry::new(tmp.path(), catalog_with_one()); - match reg.resolve("basecompute/demo", "main", None, false).unwrap() { + match reg + .resolve("basecompute/demo", "main", None, false) + .unwrap() + { ModelRef::Catalog { hf_repo, variant, .. } => { @@ -590,7 +606,8 @@ mod tests { // Asking for q4 → the installed artifact. assert!(matches!( - reg.resolve("basecompute/demo", "main", Some("q4"), false).unwrap(), + reg.resolve("basecompute/demo", "main", Some("q4"), false) + .unwrap(), ModelRef::Local { .. } )); // Asking for q8 must NOT return the installed q4 — quant-aware catalog @@ -610,8 +627,10 @@ mod tests { // Build a CatalogRegistry from inline JSON (test helper). fn catalog_json(models: &str) -> CatalogRegistry { CatalogRegistry::from_catalog( - Catalog::from_json(&format!(r#"{{"schema":1,"updated":"x","models":[{models}]}}"#)) - .unwrap(), + Catalog::from_json(&format!( + r#"{{"schema":1,"updated":"x","models":[{models}]}}"# + )) + .unwrap(), ) } @@ -651,7 +670,10 @@ mod tests { for rows in [format!("{uni},{nat}"), format!("{nat},{uni}")] { let reg = catalog_json(&rows); reg.catalog.validate().unwrap(); // both rows coexist (backend in key) - match reg.resolve_variant("basecompute/hybrid", Some("q4")).unwrap() { + match reg + .resolve_variant("basecompute/hybrid", Some("q4")) + .unwrap() + { ModelRef::Catalog { file, variant, .. } => { assert_eq!(file, format!("hybrid-Q4-{be}.base")); assert_eq!(variant, format!("{be}-q4")); // distinct cache dir @@ -672,23 +694,37 @@ mod tests { } // A quant that isn't published: absent (not backend-locked) → the merged // resolver may convert-on-pull. - assert_eq!(reg.resolve_with_status("basecompute/demo", Some("q8")), (None, false)); + assert_eq!( + reg.resolve_with_status("basecompute/demo", Some("q8")), + (None, false) + ); } #[test] fn resolve_backend_locked_reports_status_not_absent() { - let foreign = if CatalogRegistry::client_backend() == "cuda" { "metal" } else { "cuda" }; + let foreign = if CatalogRegistry::client_backend() == "cuda" { + "metal" + } else { + "cuda" + }; let reg = catalog_json(&format!( r#"{{"id":"basecompute/locked","hf_repo":"basecompute/locked","file":"locked-Q4.base","arch":"llama","quant":"{foreign}-q4","backend":"{foreign}"}}"# )); // Published for this id+quant, but only for a foreign backend → refuse, // and flag it distinctly from "absent" so the caller errors (below). - assert_eq!(reg.resolve_with_status("basecompute/locked", Some("q4")), (None, true)); + assert_eq!( + reg.resolve_with_status("basecompute/locked", Some("q4")), + (None, true) + ); } #[test] fn merged_resolve_refuses_backend_locked_no_hf_fallthrough() { - let foreign = if CatalogRegistry::client_backend() == "cuda" { "metal" } else { "cuda" }; + let foreign = if CatalogRegistry::client_backend() == "cuda" { + "metal" + } else { + "cuda" + }; let tmp = tempfile::tempdir().unwrap(); let reg = MergedRegistry::new( tmp.path(), @@ -698,7 +734,9 @@ mod tests { ); // Must ERROR (backend-locked), not fall through to a raw HF download of a // bundle this client can't run. - assert!(reg.resolve("basecompute/locked", "main", Some("q4"), false).is_err()); + assert!(reg + .resolve("basecompute/locked", "main", Some("q4"), false) + .is_err()); } #[test] @@ -719,7 +757,10 @@ mod tests { let uni = cache::variant_dir(tmp.path(), "basecompute/h", "default-q4").unwrap(); std::fs::create_dir_all(&uni).unwrap(); std::fs::write(cache::base_artifact_path(&uni), b"universal").unwrap(); - match reg.resolve("basecompute/h", "main", Some("q4"), false).unwrap() { + match reg + .resolve("basecompute/h", "main", Some("q4"), false) + .unwrap() + { ModelRef::Catalog { variant, .. } => assert_eq!(variant, format!("{be}-q4")), other => panic!("cached universal must not shadow the native pick, got {other:?}"), } @@ -760,7 +801,10 @@ mod tests { r#"{"id":"basecompute/CamelModel","hf_repo":"basecompute/alias","file":"alias-Q4.base","arch":"llama","quant":"default-q4"}, {"id":"basecompute/camelmodel","hf_repo":"basecompute/exact","file":"exact-Q4.base","arch":"llama","quant":"default-q4"}"#, ); - match reg.resolve_variant("basecompute/camelmodel", Some("q4")).unwrap() { + match reg + .resolve_variant("basecompute/camelmodel", Some("q4")) + .unwrap() + { ModelRef::Catalog { hf_repo, .. } => assert_eq!(hf_repo, "basecompute/exact"), other => panic!("exact id must win, got {other:?}"), } @@ -800,7 +844,10 @@ mod tests { std::fs::create_dir_all(&vdir).unwrap(); std::fs::write(cache::base_artifact_path(&vdir), b"not a real base").unwrap(); - match reg.resolve("basecompute/demo", "main", None, false).unwrap() { + match reg + .resolve("basecompute/demo", "main", None, false) + .unwrap() + { ModelRef::Local { variant, .. } => assert_eq!(variant, "default-q4"), other => panic!("expected Local, got {other:?}"), } diff --git a/base-convert/crates/base-hub/src/scan.rs b/base-convert/crates/base-hub/src/scan.rs index 59ee228..401871b 100644 --- a/base-convert/crates/base-hub/src/scan.rs +++ b/base-convert/crates/base-hub/src/scan.rs @@ -48,54 +48,33 @@ pub struct RemoteFile { /// bundle would look unpinnable. Reading the tree directly is both correct and /// one less layer. pub trait RepoIndex { - fn files(&self, repo: &str) -> Result>; + /// The files of `repo` at `revision` — a commit, once the scan has + /// pinned one, so the listing, the manifest and the header all describe + /// one publication. + fn files(&self, repo: &str, revision: &str) -> Result>; } /// The real Hub. pub struct HubApi; impl RepoIndex for HubApi { - fn files(&self, repo: &str) -> Result> { - let url = format!("{}/api/models/{repo}/tree/main?recursive=true", endpoint()); - let body = ureq::get(&url) - .call() - .with_context(|| format!("listing files in {repo}"))? - .into_body() - .read_to_string() - .with_context(|| format!("reading the file listing for {repo}"))?; - let entries: Vec = - serde_json::from_str(&body).with_context(|| format!("parsing {repo}'s tree"))?; - Ok(entries - .iter() - .filter(|e| e.get("type").and_then(|t| t.as_str()) == Some("file")) - .filter_map(|e| { - let path = e.get("path")?.as_str()?.to_string(); - let lfs = e.get("lfs"); - Some(RemoteFile { - size: lfs - .and_then(|l| l.get("size")) - .or_else(|| e.get("size")) - .and_then(|v| v.as_u64()) - .unwrap_or(0), - sha256: lfs - .and_then(|l| l.get("oid")) - .and_then(|v| v.as_str()) - .map(str::to_string), - path, - }) + fn files(&self, repo: &str, revision: &str) -> Result> { + // The same listing the pull uses: paginated, authenticated, timed, + // and keeping the LFS hash. A tree read in one shot would stop at + // the first page and make a bundle on the next one vanish from the + // regenerated catalog. + Ok(crate::fetch::fetch_tree(repo, revision)? + .into_iter() + .map(|e| RemoteFile { + path: e.path, + size: e.size, + sha256: e.lfs_sha256, }) .collect()) } } -/// Hub API base. `$HF_ENDPOINT` redirects the whole scan at a mirror, the same -/// variable hf-hub honors for downloads. -fn endpoint() -> String { - std::env::var("HF_ENDPOINT") - .ok() - .filter(|s| !s.is_empty()) - .unwrap_or_else(|| "https://huggingface.co".to_string()) -} +use crate::fetch::endpoint; /// What a scan found, including what it deliberately did not publish. #[derive(Debug, Default)] @@ -177,7 +156,7 @@ pub fn list_org_repos(org: &str) -> Result> { } /// Pull the `rel="next"` target out of a `Link` header. -fn parse_next_link(link: &str) -> Option { +pub(crate) fn parse_next_link(link: &str) -> Option { link.split(',').find_map(|part| { if !part.contains("rel=\"next\"") { return None; @@ -232,13 +211,56 @@ pub fn read_remote_header( revision: &str, filename: &str, ) -> Result
{ - let prefix = fetcher.read_range(repo, revision, filename, 0..PREFIX_LEN)?; + read_header_via( + &|range| fetcher.read_range(repo, revision, filename, range), + filename, + ) +} + +/// One logical byte range of a split bundle, read from whichever parts it +/// falls in. A header can straddle part 000 when the split is small or the +/// tokenizer large; the installer accepts any cut, so the scan has to. +pub fn read_split_range( + fetcher: &dyn Fetcher, + repo: &str, + revision: &str, + parts: &[(String, u64)], + range: std::ops::Range, +) -> Result> { + let mut out = Vec::with_capacity((range.end - range.start) as usize); + let mut start = 0u64; + for (path, size) in parts { + let end = start.checked_add(*size).context("part sizes overflow")?; + let lo = range.start.max(start); + let hi = range.end.min(end); + if lo < hi { + out.extend(fetcher.read_range(repo, revision, path, lo - start..hi - start)?); + } + start = end; + if start >= range.end { + break; + } + } + Ok(out) +} + +/// The header of a `.base`, given a way to read byte ranges of it. +fn read_header_via( + read: &dyn Fn(std::ops::Range) -> Result>, + filename: &str, +) -> Result
{ + let prefix = read(0..PREFIX_LEN)?; anyhow::ensure!( prefix.len() == PREFIX_LEN as usize, "{filename}: short prefix ({} bytes)", prefix.len() ); anyhow::ensure!(&prefix[0..4] == b"BASE", "{filename}: not a .base file"); + let version = u32::from_le_bytes(prefix[4..8].try_into().unwrap()); + anyhow::ensure!( + version == base_format::FORMAT_VERSION, + "{filename}: unsupported .base format version {version}" + ); let header_len = u64::from_le_bytes(prefix[8..16].try_into().unwrap()); // A header claiming to be enormous is corrupt or hostile; refuse rather // than allocate it. 256MB is far past any real tokenizer. @@ -246,15 +268,111 @@ pub fn read_remote_header( header_len > 0 && header_len < 256 * 1024 * 1024, "{filename}: implausible header length {header_len}" ); - let json = fetcher.read_range( - repo, - revision, - filename, - PREFIX_LEN..PREFIX_LEN + header_len, - )?; + let json = read(PREFIX_LEN..PREFIX_LEN + header_len)?; + anyhow::ensure!( + json.len() as u64 == header_len, + "{filename}: header truncated ({} of {header_len} bytes)", + json.len() + ); Header::from_json_bytes(&json).with_context(|| format!("parsing {filename}'s header")) } +/// One bundle as the scan sees it: the logical file, where its header can +/// be read from, and how many parts it is split into (0 for a whole file). +struct Bundle { + file: RemoteFile, + /// Where the header's bytes live: the file itself, or the parts in + /// order with their sizes. + header_from: Vec<(String, u64)>, + parts: usize, + /// The Hub's sha256 for each part, in order, `None` where it reports + /// none (an inline part, a mirror without LFS metadata); empty for a + /// whole file. Every value that is there gets compared. + part_ids: Vec>, + /// The part paths, in order; empty for a whole file. + part_paths: Vec, + /// The Hub's size for each part, in order; empty for a whole file. + part_sizes: Vec, + /// The manifest published beside the parts, when there is one. + manifest: Option, +} + +/// Fold a repo listing into bundles. A part set becomes one logical file +/// whose size is the sum of its parts, with no sha256 of its own yet (the +/// manifest supplies that later) and its header behind part 000. +fn group_bundles(files: Vec) -> (Vec, Vec<(String, String)>) { + let by_path: std::collections::HashMap = + files.into_iter().map(|f| (f.path.clone(), f)).collect(); + let grouped = crate::parts::group(by_path.keys().cloned()); + let mut malformed = grouped.malformed; + let mut out = Vec::with_capacity(grouped.artifacts.len()); + for a in grouped.artifacts { + if a.is_split() { + let manifest = by_path.get(&crate::parts::manifest_name(&a.name)).cloned(); + let files: Vec<&RemoteFile> = a.parts.iter().filter_map(|p| by_path.get(p)).collect(); + let part_sizes: Vec = files.iter().map(|f| f.size).collect(); + // A listing is untrusted input: sizes that do not add up are a + // malformed set, not a panic that ends the scan. + let Some(size) = part_sizes + .iter() + .try_fold(0u64, |acc, s| acc.checked_add(*s)) + else { + malformed.push((a.name, "listed part sizes overflow".to_string())); + continue; + }; + let part_ids: Vec> = files.iter().map(|f| f.sha256.clone()).collect(); + out.push(Bundle { + file: RemoteFile { + path: a.name, + size, + sha256: None, + }, + header_from: a + .parts + .iter() + .cloned() + .zip(part_sizes.iter().copied()) + .collect(), + parts: a.parts.len(), + part_ids, + part_paths: a.parts, + part_sizes, + manifest, + }); + } else if let Some(f) = by_path.get(&a.name) { + out.push(Bundle { + header_from: vec![(f.path.clone(), f.size)], + file: f.clone(), + parts: 0, + part_ids: Vec::new(), + part_paths: Vec::new(), + part_sizes: Vec::new(), + manifest: None, + }); + } + } + // Listings are unordered maps here; keep the report deterministic. + out.sort_by(|x, y| x.file.path.cmp(&y.file.path)); + (out, malformed) +} + +/// Fetch and parse a split bundle's manifest with one ranged read. +fn read_remote_manifest( + fetcher: &dyn Fetcher, + repo: &str, + revision: &str, + file: &RemoteFile, +) -> Result { + anyhow::ensure!( + file.size > 0 && file.size < crate::parts::MAX_MANIFEST_LEN, + "{}: implausible manifest size {}", + file.path, + file.size + ); + let bytes = fetcher.read_range(repo, revision, &file.path, 0..file.size)?; + serde_json::from_slice(&bytes).with_context(|| format!("parsing {}", file.path)) +} + /// Build catalog rows for every `.base` bundle published under `org`. /// /// `known` supplies rows to reuse: when a file's sha256 is unchanged, its @@ -270,7 +388,19 @@ pub fn scan_org( ) -> Result { let mut report = ScanReport::default(); for repo in repos { - let files = match index.files(repo) { + // One commit for everything read about this repo, so a publish + // landing mid-scan cannot pair one publication's manifest with + // another's header. + let pinned = match fetcher.resolve_revision(repo, "main") { + Ok(p) => p, + Err(e) => { + report + .skipped + .push((repo.clone(), format!("resolving main failed: {e:#}"))); + continue; + } + }; + let files = match index.files(repo, &pinned) { Ok(f) => f, Err(e) => { report @@ -279,15 +409,93 @@ pub fn scan_org( continue; } }; - let bundles: Vec<_> = files - .into_iter() - .filter(|f| f.path.ends_with(".base")) - .collect(); + let (bundles, malformed) = group_bundles(files); + // A part set with a gap (a quant mid-upload, say) is reported on + // its own; the repo's other bundles are scanned as usual. + for (name, why) in malformed { + report.skipped.push((format!("{repo}/{name}"), why)); + } if bundles.is_empty() { report.empty_repos.push(repo.clone()); continue; } - for f in bundles { + for Bundle { + file: mut f, + header_from, + parts, + part_ids, + part_paths, + part_sizes, + manifest, + } in bundles + { + // A split bundle has a sha256 per part on the Hub and none for + // the whole. Its manifest supplies the whole-file hash and size, + // and is believed only when the Hub's per-part hashes are the + // ones it lists — a republished part with a stale manifest + // would otherwise pin a hash no download can match. + let mut parts_sha256 = None; + if parts > 0 { + let Some(m) = manifest else { + report.skipped.push(( + format!("{repo}/{}", f.path), + format!( + "split into {parts} parts with no {} beside them; publish one with `basert catalog-manifest`", + crate::parts::manifest_name(&f.path) + ), + )); + continue; + }; + let manifest = match read_remote_manifest(fetcher, repo, &pinned, &m) { + Ok(m) => m, + Err(e) => { + report + .skipped + .push((format!("{repo}/{}", f.path), format!("{e:#}"))); + continue; + } + }; + // The same check a pull runs, so a row is never written for + // a manifest every pull would then refuse. + if let Err(e) = manifest + .check_against(&part_paths) + .and_then(|()| manifest.check_sizes(&part_sizes)) + { + report.skipped.push(( + format!("{repo}/{}", f.path), + format!("its manifest disagrees with the listing: {e:#}; republish the manifest"), + )); + continue; + } + let ids = manifest.part_ids(); + // Every hash the Hub does have must agree with the manifest; + // hex case is not content, so the compare ignores it, as the + // install's does. + let same_ids = ids.len() == part_ids.len() + && ids + .iter() + .zip(&part_ids) + .all(|(a, b)| b.as_deref().is_none_or(|b| a.eq_ignore_ascii_case(b))); + let hub_ids: Vec<&str> = part_ids + .iter() + .map(|p| p.as_deref().unwrap_or("?")) + .collect(); + if manifest.parts.len() != parts || !same_ids { + report.skipped.push(( + format!("{repo}/{}", f.path), + format!( + "its manifest describes {} parts [{}] but the Hub holds {parts} [{}]: republish the manifest", + manifest.parts.len(), + ids.join(", "), + hub_ids.join(", ") + ), + )); + continue; + } + f.sha256 = Some(manifest.sha256.clone()); + f.size = manifest.size; + parts_sha256 = Some(ids); + } let Some(sha256) = f.sha256.clone() else { report.skipped.push(( format!("{repo}/{}", f.path), @@ -309,11 +517,22 @@ pub fn scan_org( )); continue; } - report.entries.push(prev.clone()); + // The same bytes may have been re-split at other boundaries + // since the row was written, which changes the parts without + // changing the whole: the split metadata is today's. + let mut row = prev.clone(); + if parts_sha256.is_some() { + row.parts_sha256 = parts_sha256; + row.size = Some(f.size); + } + report.entries.push(row); report.reused += 1; continue; } - let header = match read_remote_header(fetcher, repo, "main", &f.path) { + let header = match read_header_via( + &|range| read_split_range(fetcher, repo, &pinned, &header_from, range), + &f.path, + ) { Ok(h) => h, Err(e) => { report @@ -351,6 +570,7 @@ pub fn scan_org( // so publishing a new quant does not quietly drop it from // the ones already curated. entry.source_repo = inherited_source_repo(known, repo); + entry.parts_sha256 = parts_sha256; report.entries.push(entry) } Err(e) => report @@ -425,7 +645,7 @@ mod tests { } impl RepoIndex for BytesFetcher { - fn files(&self, _: &str) -> Result> { + fn files(&self, _: &str, _: &str) -> Result> { Ok(self .files .iter() @@ -495,6 +715,7 @@ mod tests { quant: "default-q4".into(), size: Some(19), sha256: Some("abc123".into()), + parts_sha256: None, backend: None, }], }; @@ -559,6 +780,7 @@ mod tests { quant: "default-q4".into(), size: Some(1), sha256: Some("q4sha".into()), + parts_sha256: None, backend: None, }, CatalogEntry { @@ -571,6 +793,7 @@ mod tests { quant: "default-q4".into(), size: Some(1), sha256: Some("othersha".into()), + parts_sha256: None, backend: None, }, ], @@ -585,6 +808,333 @@ mod tests { assert_eq!(inherited_source_repo(&known, "basecompute/unknown"), None); } + #[test] + fn a_split_bundle_is_catalogued_from_its_manifest() { + use crate::parts::{synthetic_bundle, Manifest, ManifestPart}; + // Fake but well-formed hashes: the manifest check wants 64 hex. + fn h(tag: &str) -> String { + format!( + "{:0>64}", + tag.bytes().map(|b| format!("{b:02x}")).collect::() + ) + } + // Part 000 carries a real header; the scan reads it from there. + let whole = synthetic_bundle(1000); + let (p0, p1) = whole.split_at(whole.len() / 2); + let manifest = Manifest { + size: whole.len() as u64, + sha256: h("whole-sha"), + parts: vec![ + ManifestPart { + name: "g.base.part-000".into(), + size: p0.len() as u64, + sha256: h("p0"), + }, + ManifestPart { + name: "g.base.part-001".into(), + size: p1.len() as u64, + sha256: h("p1"), + }, + ], + }; + let repo = |manifest: Option<&Manifest>| { + let mut files = vec![ + ("parts/g.base.part-000".into(), p0.to_vec(), Some(h("p0"))), + ("parts/g.base.part-001".into(), p1.to_vec(), Some(h("p1"))), + ]; + if let Some(m) = manifest { + files.push(( + "parts/g.base.manifest.json".into(), + serde_json::to_vec(m).unwrap(), + None, + )); + } + BytesFetcher { files } + }; + let none = Catalog { + schema: 1, + updated: String::new(), + models: vec![], + }; + let scan = |f: &BytesFetcher, known: &Catalog| { + scan_org(f, f, "basecompute", &["basecompute/g".to_string()], known).unwrap() + }; + + // No manifest: the repo is not "empty", and the skip says what to + // publish. + let r = scan(&repo(None), &none); + assert!(r.empty_repos.is_empty()); + assert!(r.entries.is_empty()); + assert_eq!(r.skipped.len(), 1, "{:?}", r.skipped); + assert_eq!(r.skipped[0].0, "basecompute/g/parts/g.base"); + assert!( + r.skipped[0].1.contains("catalog-manifest"), + "{}", + r.skipped[0].1 + ); + + // With one: a row, whole-file hash and size from the manifest, the + // per-part hashes pinned on it, header facts from part 000. + let r = scan(&repo(Some(&manifest)), &none); + assert!(r.skipped.is_empty(), "{:?}", r.skipped); + assert_eq!(r.entries.len(), 1); + let e = &r.entries[0]; + assert_eq!(e.file, "parts/g.base"); + assert_eq!(e.sha256, Some(h("whole-sha"))); + assert_eq!(e.size, Some(whole.len() as u64)); + assert_eq!(e.parts_sha256, Some(vec![h("p0"), h("p1")])); + assert_eq!(e.arch.as_deref(), Some("test")); + + // A manifest with the right hashes but a wrong part name or size + // would be accepted here and refused by every pull: it is refused + // here instead. + let mut misnamed = manifest.clone(); + misnamed.parts[1].name = "g.base.part-01".into(); + let r = scan(&repo(Some(&misnamed)), &none); + assert!(r.entries.is_empty()); + assert_eq!(r.skipped.len(), 1, "{:?}", r.skipped); + assert!( + r.skipped[0].1.contains("disagrees with the listing"), + "{}", + r.skipped[0].1 + ); + let mut wrong_size = manifest.clone(); + wrong_size.parts[1].size += 1; + let r = scan(&repo(Some(&wrong_size)), &none); + assert!(r.entries.is_empty(), "{:?}", r.entries); + assert_eq!(r.skipped.len(), 1, "{:?}", r.skipped); + // Two wrong sizes that cancel in the sum are still wrong. + let mut offset = manifest.clone(); + offset.parts[0].size += 1; + offset.parts[1].size -= 1; + let r = scan(&repo(Some(&offset)), &none); + assert!(r.entries.is_empty(), "{:?}", r.entries); + assert_eq!(r.skipped.len(), 1, "{:?}", r.skipped); + assert!( + r.skipped[0].1.contains("bytes but the repo holds"), + "{}", + r.skipped[0].1 + ); + + // A part the Hub has no hash for does not blind the scan to the + // others: part 1's Hub hash still has to match. + let mut half_known = repo(Some(&manifest)); + half_known.files[0].2 = None; + half_known.files[1].2 = Some(h("p1-hub")); + let r = scan(&half_known, &none); + assert!(r.entries.is_empty(), "{:?}", r.entries); + assert_eq!(r.skipped.len(), 1, "{:?}", r.skipped); + assert!( + r.skipped[0].1.contains("republish the manifest"), + "{}", + r.skipped[0].1 + ); + // And with only part 0 unknown and part 1 agreeing, it is fine. + let mut half_known = repo(Some(&manifest)); + half_known.files[0].2 = None; + let r = scan(&half_known, &none); + assert!(r.skipped.is_empty(), "{:?}", r.skipped); + + // A parseable header of another format version is refused here, + // not after a pull has reassembled the whole bundle. + let mut v2 = repo(Some(&manifest)); + v2.files[0].1[4..8].copy_from_slice(&2u32.to_le_bytes()); + let r = scan(&v2, &none); + assert!(r.entries.is_empty()); + assert_eq!(r.skipped.len(), 1, "{:?}", r.skipped); + assert!( + r.skipped[0] + .1 + .contains("unsupported .base format version 2"), + "{}", + r.skipped[0].1 + ); + + // A cut inside the prefix: the header is read across the parts. + let (t0, t1) = whole.split_at(10); + let tiny = Manifest { + parts: vec![ + ManifestPart { + name: "g.base.part-000".into(), + size: t0.len() as u64, + sha256: h("t0"), + }, + ManifestPart { + name: "g.base.part-001".into(), + size: t1.len() as u64, + sha256: h("t1"), + }, + ], + ..manifest.clone() + }; + let straddling = BytesFetcher { + files: vec![ + ("parts/g.base.part-000".into(), t0.to_vec(), Some(h("t0"))), + ("parts/g.base.part-001".into(), t1.to_vec(), Some(h("t1"))), + ( + "parts/g.base.manifest.json".into(), + serde_json::to_vec(&tiny).unwrap(), + None, + ), + ], + }; + let r = scan(&straddling, &none); + assert!(r.skipped.is_empty(), "{:?}", r.skipped); + assert_eq!(r.entries.len(), 1); + assert_eq!(r.entries[0].arch.as_deref(), Some("test")); + + // Uppercase hex in the manifest is the same digest. + let mut upper = manifest.clone(); + for p in &mut upper.parts { + p.sha256 = p.sha256.to_ascii_uppercase(); + } + let r = scan(&repo(Some(&upper)), &none); + assert!(r.skipped.is_empty(), "{:?}", r.skipped); + assert_eq!(r.entries.len(), 1); + + // A manifest the Hub disagrees with (part 1 republished): refused. + let mut stale = manifest.clone(); + stale.parts[1].sha256 = "p1-old".into(); + let r = scan(&repo(Some(&stale)), &none); + assert!(r.entries.is_empty()); + assert_eq!(r.skipped.len(), 1, "{:?}", r.skipped); + assert!( + r.skipped[0].1.contains("republish the manifest"), + "{}", + r.skipped[0].1 + ); + + // An unchanged bundle reuses its row without a header read. + let known = Catalog { + schema: 1, + updated: String::new(), + models: scan(&repo(Some(&manifest)), &none).entries, + }; + let r = scan(&repo(Some(&manifest)), &known); + assert_eq!(r.reused, 1); + assert_eq!(r.entries.len(), 1); + + // The same bytes re-split at other boundaries: the whole-file hash + // is unchanged, so the row is reused, but its part hashes are + // today's — a pull checks the manifest against the row. + let (q0, q1) = whole.split_at(whole.len() / 3); + let resplit = Manifest { + parts: vec![ + ManifestPart { + name: "g.base.part-000".into(), + size: q0.len() as u64, + sha256: h("q0"), + }, + ManifestPart { + name: "g.base.part-001".into(), + size: q1.len() as u64, + sha256: h("q1"), + }, + ], + ..manifest.clone() + }; + let f = BytesFetcher { + files: vec![ + ("parts/g.base.part-000".into(), q0.to_vec(), Some(h("q0"))), + ("parts/g.base.part-001".into(), q1.to_vec(), Some(h("q1"))), + ( + "parts/g.base.manifest.json".into(), + serde_json::to_vec(&resplit).unwrap(), + None, + ), + ], + }; + let r = scan(&f, &known); + assert_eq!(r.reused, 1); + assert_eq!(r.entries[0].parts_sha256, Some(vec![h("q0"), h("q1")])); + } + + #[test] + fn listed_part_sizes_that_overflow_are_a_skip_not_a_panic() { + let files = vec![ + RemoteFile { + path: "g.base.part-000".into(), + size: u64::MAX, + sha256: Some("p0".into()), + }, + RemoteFile { + path: "g.base.part-001".into(), + size: 1, + sha256: Some("p1".into()), + }, + RemoteFile { + path: "m-Q4.base".into(), + size: 3, + sha256: Some("m".into()), + }, + ]; + let (bundles, malformed) = group_bundles(files); + assert_eq!(bundles.len(), 1); + assert_eq!(bundles[0].file.path, "m-Q4.base"); + assert_eq!(malformed.len(), 1); + assert_eq!(malformed[0].0, "g.base"); + assert!(malformed[0].1.contains("overflow"), "{}", malformed[0].1); + } + + #[test] + fn a_gapped_part_set_is_skipped_without_dropping_its_siblings() { + let f = BytesFetcher { + files: vec![ + ( + "m-Q4.base".into(), + b"not a header at all".to_vec(), + Some("abc123".into()), + ), + ( + "m-Q8.base.part-000".into(), + b"x".to_vec(), + Some("p0".into()), + ), + ( + "m-Q8.base.part-002".into(), + b"x".to_vec(), + Some("p2".into()), + ), + ], + }; + let known = Catalog { + schema: 1, + updated: String::new(), + models: vec![CatalogEntry { + id: "basecompute/m".into(), + hf_repo: "basecompute/m".into(), + file: "m-Q4.base".into(), + revision: "main".into(), + source_repo: None, + arch: Some("qwen35".into()), + quant: "default-q4".into(), + size: Some(19), + sha256: Some("abc123".into()), + parts_sha256: None, + backend: None, + }], + }; + let r = scan_org( + &f, + &f, + "basecompute", + &["basecompute/m".to_string()], + &known, + ) + .unwrap(); + // The good row survives; the half-uploaded quant is named. + assert_eq!(r.entries.len(), 1); + assert_eq!(r.entries[0].file, "m-Q4.base"); + assert_eq!(r.skipped.len(), 1, "{:?}", r.skipped); + assert_eq!(r.skipped[0].0, "basecompute/m/m-Q8.base"); + assert!( + r.skipped[0].1.contains("missing part 001"), + "{}", + r.skipped[0].1 + ); + assert!(r.empty_repos.is_empty()); + } + #[test] fn a_repo_with_no_bundles_is_recorded_not_treated_as_failure() { let f = BytesFetcher { diff --git a/base-convert/crates/base-quant/src/base_q4.rs b/base-convert/crates/base-quant/src/base_q4.rs index efbdd80..4e17df6 100644 --- a/base-convert/crates/base-quant/src/base_q4.rs +++ b/base-convert/crates/base-quant/src/base_q4.rs @@ -31,7 +31,10 @@ pub fn pack(weights: &[f32]) -> Packed { /// Pack with an explicit group size (tests use smaller groups for /// hand-verified fixtures). pub fn pack_with_group_size(weights: &[f32], group_size: usize) -> Packed { - assert!(group_size > 0 && group_size % 2 == 0, "group_size must be even"); + assert!( + group_size > 0 && group_size % 2 == 0, + "group_size must be even" + ); assert!( weights.len() % group_size == 0, "weights.len()={} must be a multiple of group_size={}", @@ -66,11 +69,9 @@ pub fn pack_with_group_size(weights: &[f32], group_size: usize) -> Packed { let q = ((val - bias) * inv_scale).round().clamp(0.0, 15.0) as u8; let byte_idx = (g * group_size + i) / 2; if i % 2 == 0 { - packed_weights[byte_idx] = - (packed_weights[byte_idx] & 0xF0) | (q & 0x0F); + packed_weights[byte_idx] = (packed_weights[byte_idx] & 0xF0) | (q & 0x0F); } else { - packed_weights[byte_idx] = - (packed_weights[byte_idx] & 0x0F) | ((q & 0x0F) << 4); + packed_weights[byte_idx] = (packed_weights[byte_idx] & 0x0F) | ((q & 0x0F) << 4); } } } diff --git a/base-convert/crates/base-quant/src/nvfp4.rs b/base-convert/crates/base-quant/src/nvfp4.rs index 5a82f28..e7031a2 100644 --- a/base-convert/crates/base-quant/src/nvfp4.rs +++ b/base-convert/crates/base-quant/src/nvfp4.rs @@ -84,8 +84,8 @@ fn to_e4m3(x: f32) -> (u8, f32) { } let exp = x.log2().floor() as i32; let mantissa_f = x / 2f32.powi(exp); // in [1, 2) - // 3-bit mantissa → 8 bins in [1, 2) spanning mantissa = 1.0..1.875 - // (values are 1 + n/8). + // 3-bit mantissa → 8 bins in [1, 2) spanning mantissa = 1.0..1.875 + // (values are 1 + n/8). let man_q_bits = ((mantissa_f - 1.0) * 8.0).round().clamp(0.0, 7.0) as u32; let man_q = 1.0 + (man_q_bits as f32) / 8.0; diff --git a/base-convert/crates/base-quant/src/profile.rs b/base-convert/crates/base-quant/src/profile.rs index 9d16da3..3046766 100644 --- a/base-convert/crates/base-quant/src/profile.rs +++ b/base-convert/crates/base-quant/src/profile.rs @@ -72,16 +72,15 @@ pub struct ResolvedQuant { impl QuantProfile { /// Parse a profile from JSON bytes. pub fn from_json(bytes: &[u8]) -> Result { - let p: Self = serde_json::from_slice(bytes) - .context("parsing quant profile JSON")?; + let p: Self = serde_json::from_slice(bytes).context("parsing quant profile JSON")?; p.validate()?; Ok(p) } /// Read a profile from a path on disk. pub fn from_path(path: &Path) -> Result { - let bytes = std::fs::read(path) - .with_context(|| format!("reading profile {}", path.display()))?; + let bytes = + std::fs::read(path).with_context(|| format!("reading profile {}", path.display()))?; Self::from_json(&bytes) } @@ -106,9 +105,7 @@ impl QuantProfile { bail!("rule {i}: unbalanced {{}} in pattern {:?}", rule.pattern); } // e4m3 is q8-only per spec. - if rule.scale_dtype == Some(ScaleDtype::E4m3) - && rule.dtype != TensorDtype::BaseQ8 - { + if rule.scale_dtype == Some(ScaleDtype::E4m3) && rule.dtype != TensorDtype::BaseQ8 { bail!( "rule {i}: scale_dtype=e4m3 is only valid for base_q8, not {:?}", rule.dtype @@ -314,18 +311,29 @@ mod tests { #[test] fn alternation_expands() { - let pats = - expand_alternations("a.{b,c,d}.e"); + let pats = expand_alternations("a.{b,c,d}.e"); assert_eq!(pats, vec!["a.b.e", "a.c.e", "a.d.e"]); } #[test] fn alternation_in_match() { let pat = "model.layers.*.self_attn.{q,k,v,o}_proj.weight"; - assert!(pattern_matches(pat, "model.layers.0.self_attn.q_proj.weight")); - assert!(pattern_matches(pat, "model.layers.5.self_attn.k_proj.weight")); - assert!(pattern_matches(pat, "model.layers.5.self_attn.v_proj.weight")); - assert!(pattern_matches(pat, "model.layers.5.self_attn.o_proj.weight")); + assert!(pattern_matches( + pat, + "model.layers.0.self_attn.q_proj.weight" + )); + assert!(pattern_matches( + pat, + "model.layers.5.self_attn.k_proj.weight" + )); + assert!(pattern_matches( + pat, + "model.layers.5.self_attn.v_proj.weight" + )); + assert!(pattern_matches( + pat, + "model.layers.5.self_attn.o_proj.weight" + )); assert!(!pattern_matches( pat, "model.layers.5.self_attn.gate_proj.weight" @@ -373,7 +381,10 @@ mod tests { TensorDtype::Bf16 ); assert_eq!( - profile.resolve("model.layers.0.q_proj.weight").unwrap().dtype, + profile + .resolve("model.layers.0.q_proj.weight") + .unwrap() + .dtype, TensorDtype::BaseQ4 ); } @@ -472,8 +483,9 @@ mod tests { let p = QuantProfile::from_path(&path).unwrap(); // MoE expert: q4 / gs=64. - let expert = - p.resolve("model.layers.0.mlp.experts.0.gate_proj.weight").unwrap(); + let expert = p + .resolve("model.layers.0.mlp.experts.0.gate_proj.weight") + .unwrap(); assert_eq!(expert.dtype, TensorDtype::BaseQ4); assert_eq!(expert.group_size, 64); @@ -491,10 +503,7 @@ mod tests { // Router stays in fp (kernel reads f16; profile uses f16 // since the runtime's norm/router kernels assume half). let router = p.resolve("model.layers.0.mlp.router.weight").unwrap(); - assert!(matches!( - router.dtype, - TensorDtype::F16 | TensorDtype::Bf16 - )); + assert!(matches!(router.dtype, TensorDtype::F16 | TensorDtype::Bf16)); } /// q3-aggressive profile routes MLP / experts to q3 / gs=32 per spec. @@ -543,7 +552,9 @@ mod tests { // residual-stream noise doesn't compound through 30 layers). for proj in ["gate_proj", "up_proj", "down_proj"] { let name = format!("model.layers.0.mlp.{proj}.weight"); - let r = p.resolve(&name).unwrap_or_else(|| panic!("no rule for {name}")); + let r = p + .resolve(&name) + .unwrap_or_else(|| panic!("no rule for {name}")); assert_eq!(r.dtype, TensorDtype::BaseQ8, "mlp.{proj} should be q8"); assert_eq!(r.group_size, 64, "mlp.{proj} should be gs=64"); } @@ -556,7 +567,9 @@ mod tests { "model.layers.0.mlp.router.weight", "model.layers.0.ffn_gate_inp.weight", ] { - let r = p.resolve(name).unwrap_or_else(|| panic!("no rule for {name}")); + let r = p + .resolve(name) + .unwrap_or_else(|| panic!("no rule for {name}")); assert_eq!(r.dtype, TensorDtype::BaseQ8, "{name} should be q8"); assert_eq!(r.group_size, 64, "{name} should be gs=64"); } @@ -568,7 +581,9 @@ mod tests { "model.layers.0.ffn_gate_up_exps.weight", "model.layers.0.ffn_down_exps.weight", ] { - let r = p.resolve(name).unwrap_or_else(|| panic!("no rule for {name}")); + let r = p + .resolve(name) + .unwrap_or_else(|| panic!("no rule for {name}")); assert_eq!(r.dtype, TensorDtype::BaseQ4, "{name} should be q4"); assert_eq!(r.group_size, 64, "{name} should be gs=64"); } @@ -576,8 +591,14 @@ mod tests { // Attention projections — q4/gs=64. for proj in ["q_proj", "k_proj", "v_proj", "o_proj"] { let name = format!("model.layers.0.self_attn.{proj}.weight"); - let r = p.resolve(&name).unwrap_or_else(|| panic!("no rule for {name}")); - assert_eq!(r.dtype, TensorDtype::BaseQ4, "self_attn.{proj} should be q4"); + let r = p + .resolve(&name) + .unwrap_or_else(|| panic!("no rule for {name}")); + assert_eq!( + r.dtype, + TensorDtype::BaseQ4, + "self_attn.{proj} should be q4" + ); } // Embed_tokens — q4/gs=64 to match MLX (also drops lm_head with @@ -596,7 +617,9 @@ mod tests { "model.layers.0.pre_feedforward_layernorm.weight", "model.layers.0.pre_feedforward_layernorm_2.weight", ] { - let r = p.resolve(name).unwrap_or_else(|| panic!("no rule for {name}")); + let r = p + .resolve(name) + .unwrap_or_else(|| panic!("no rule for {name}")); assert_eq!(r.dtype, TensorDtype::F16, "{name} should be f16"); } @@ -609,7 +632,9 @@ mod tests { "model.layers.0.router.scale", "model.layers.0.ffn_gate_inp.scale", ] { - let r = p.resolve(name).unwrap_or_else(|| panic!("no rule for {name}")); + let r = p + .resolve(name) + .unwrap_or_else(|| panic!("no rule for {name}")); assert_eq!(r.dtype, TensorDtype::F16, "{name} should be f16"); } } @@ -646,7 +671,9 @@ mod tests { "encoder.blocks.0.mlp.0.weight", "decoder.blocks.2.mlp.2.weight", ] { - let r = p.resolve(name).unwrap_or_else(|| panic!("{file}: no rule for {name}")); + let r = p + .resolve(name) + .unwrap_or_else(|| panic!("{file}: no rule for {name}")); assert_eq!(r.dtype, dtype, "{file}: {name}"); assert_eq!(r.group_size, gs, "{file}: {name}"); assert_eq!(r.scale_dtype, ScaleDtype::Bf16, "{file}: {name}"); @@ -670,7 +697,9 @@ mod tests { "encoder.blocks.0.mlp.0.bias", "decoder.ln.bias", ] { - let r = p.resolve(name).unwrap_or_else(|| panic!("{file}: no rule for {name}")); + let r = p + .resolve(name) + .unwrap_or_else(|| panic!("{file}: no rule for {name}")); assert_eq!(r.dtype, TensorDtype::F16, "{file}: {name} must stay f16"); } } diff --git a/base-convert/crates/base-quant/src/rtn.rs b/base-convert/crates/base-quant/src/rtn.rs index d47b295..c27c432 100644 --- a/base-convert/crates/base-quant/src/rtn.rs +++ b/base-convert/crates/base-quant/src/rtn.rs @@ -64,7 +64,8 @@ pub fn pack(weights: &[f32], cfg: RtnConfig) -> Packed { let group_size = cfg.group_size as usize; let n_groups = weights.len() / group_size; let mut q_lanes: Vec = vec![0; weights.len()]; - let mut scales_bytes = Vec::with_capacity(n_groups * cfg.scale_dtype.bytes_per_group() as usize); + let mut scales_bytes = + Vec::with_capacity(n_groups * cfg.scale_dtype.bytes_per_group() as usize); let mut biases_bytes = Vec::with_capacity(if cfg.symmetric { 0 } else { @@ -118,7 +119,9 @@ pub fn pack(weights: &[f32], cfg: RtnConfig) -> Packed { let inv_scale = 1.0 / scale_rt; for (i, &val) in group.iter().enumerate() { let q = if cfg.symmetric { - (val * inv_scale).round().clamp(q_min_sym as f32, q_max_sym as f32) as i32 + (val * inv_scale) + .round() + .clamp(q_min_sym as f32, q_max_sym as f32) as i32 } else { ((val - bias_rt) * inv_scale) .round() @@ -147,6 +150,111 @@ pub fn pack(weights: &[f32], cfg: RtnConfig) -> Packed { } } +/// Importance-weighted RTN ("imatrix-lite"): per-group affine fit that +/// minimizes Σ w·(x − (q·s + b))² instead of plain min/max range fitting, +/// where `w` is a per-INPUT-CHANNEL importance (activation second moment +/// from the calibration sidecar). Groups run along the input dim, so group +/// g covers channels [(g·gs) mod in_features, +gs). The packed layout, +/// scale dtype round-trip and value space are identical to `pack` — the +/// runtime cannot tell the difference; only WHERE the quantization error +/// lands changes (away from salient channels). Asymmetric only. +/// +/// This is the runtime-free counterpart of AWQ: same intuition (spend the +/// grid on channels the network actually feels), no folded activation +/// scale to apply at inference. +pub fn pack_weighted( + weights: &[f32], + cfg: RtnConfig, + channel_w: &[f32], + in_features: usize, +) -> Packed { + assert!(!cfg.symmetric, "pack_weighted supports asymmetric only"); + assert!(cfg.group_size > 0 && weights.len() % cfg.group_size as usize == 0); + assert!( + in_features > 0 && in_features % cfg.group_size as usize == 0, + "in_features {in_features} must be a multiple of group_size {}", + cfg.group_size + ); + assert_eq!( + channel_w.len(), + in_features, + "one importance per input channel" + ); + + let group_size = cfg.group_size as usize; + let n_groups = weights.len() / group_size; + let mut q_lanes: Vec = vec![0; weights.len()]; + let mut scales_bytes = + Vec::with_capacity(n_groups * cfg.scale_dtype.bytes_per_group() as usize); + let mut biases_bytes = + Vec::with_capacity(n_groups * cfg.scale_dtype.bytes_per_group() as usize); + let q_max: f32 = ((1u32 << cfg.bits) - 1) as f32; + + for g in 0..n_groups { + let group = &weights[g * group_size..(g + 1) * group_size]; + let ch0 = (g * group_size) % in_features; + let w = &channel_w[ch0..ch0 + group_size]; + + // Seed from plain min/max, then a few weighted-Lloyd rounds: + // assign under (s, b), refit (s, b) by weighted least squares. + let (mut mn, mut mx) = (f32::INFINITY, f32::NEG_INFINITY); + for &x in group { + mn = mn.min(x); + mx = mx.max(x); + } + let mut s = (mx - mn) / q_max; + if s == 0.0 { + s = 1.0; + } + let mut b = mn; + let wsum: f64 = w.iter().map(|&v| v.max(1e-8) as f64).sum(); + for _ in 0..8 { + // Weighted moments of (q, x) under the current assignment. + let (mut sq, mut sx, mut sqq, mut sqx) = (0f64, 0f64, 0f64, 0f64); + for (i, &x) in group.iter().enumerate() { + let wi = w[i].max(1e-8) as f64; + let q = (((x - b) / s).round()).clamp(0.0, q_max) as f64; + sq += wi * q; + sx += wi * x as f64; + sqq += wi * q * q; + sqx += wi * q * x as f64; + } + let qm = sq / wsum; + let xm = sx / wsum; + let var = sqq / wsum - qm * qm; + if var > 1e-12 { + let cov = sqx / wsum - qm * xm; + let s_new = (cov / var) as f32; + if s_new.abs() > 1e-12 { + s = s_new; + b = (xm - qm * (s as f64)) as f32; + } + } + } + + // Final encode: round (s, b) through the scale dtype FIRST, then + // assign q under the rounded pair (matches `pack`'s contract that + // pack-side rounding equals dequant exactly). + let (scale_rt, scale_enc) = round_trip_scale(s, cfg.scale_dtype); + scales_bytes.extend_from_slice(&scale_enc); + let (bias_rt, bias_enc) = round_trip_scale(b, cfg.scale_dtype); + biases_bytes.extend_from_slice(&bias_enc); + let inv = 1.0 / scale_rt; + for (i, &x) in group.iter().enumerate() { + let q = (((x - bias_rt) * inv).round()).clamp(0.0, q_max) as u32; + q_lanes[g * group_size + i] = q; + } + } + + Packed { + packed_weights: pack_lanes(&q_lanes, cfg.bits), + scales: scales_bytes, + biases: biases_bytes, + group_size: cfg.group_size, + scale_dtype: Some(cfg.scale_dtype), + } +} + /// Dequantize an RTN-packed tensor. Inverse of `pack`. Used for tests /// and the `--validate` trace gate. pub fn unpack(packed: &Packed, total_values: usize, cfg: RtnConfig) -> Vec { @@ -179,11 +287,7 @@ pub fn unpack(packed: &Packed, total_values: usize, cfg: RtnConfig) -> Vec for i in 0..group_size { let flat = g * group_size + i; let q = q_lanes[flat] as i32; - let q_real = if cfg.symmetric { - q - q_offset_sym - } else { - q - }; + let q_real = if cfg.symmetric { q - q_offset_sym } else { q }; out[flat] = q_real as f32 * scale + bias; } } @@ -285,8 +389,7 @@ fn unpack_q3(bytes: &[u8], total: usize) -> Vec { assert!(total % 8 == 0); let mut out = Vec::with_capacity(total); for chunk in bytes.chunks_exact(3) { - let acc = - (chunk[0] as u32) | ((chunk[1] as u32) << 8) | ((chunk[2] as u32) << 16); + let acc = (chunk[0] as u32) | ((chunk[1] as u32) << 8) | ((chunk[2] as u32) << 16); for i in 0..8 { out.push((acc >> (i * 3)) & 0x7); } @@ -353,8 +456,7 @@ fn unpack_q6(bytes: &[u8], total: usize) -> Vec { assert!(total % 4 == 0); let mut out = Vec::with_capacity(total); for chunk in bytes.chunks_exact(3) { - let acc = - (chunk[0] as u32) | ((chunk[1] as u32) << 8) | ((chunk[2] as u32) << 16); + let acc = (chunk[0] as u32) | ((chunk[1] as u32) << 8) | ((chunk[2] as u32) << 16); for i in 0..4 { out.push((acc >> (i * 6)) & 0x3F); } @@ -447,7 +549,7 @@ fn encode_e4m3_approx(x: f32) -> u8 { let mag = x.abs(); let log2 = mag.log2().round() as i32; let exp = (log2 + 7).clamp(0, 15) as u8; // 4-bit exponent, bias 7 - // Mantissa: 3 bits of fraction beyond the implicit leading 1. + // Mantissa: 3 bits of fraction beyond the implicit leading 1. let frac = mag / 2f32.powi(log2) - 1.0; let mantissa = (frac * 8.0).round().clamp(0.0, 7.0) as u8; sign | (exp << 3) | (mantissa & 0x7) @@ -483,10 +585,7 @@ mod tests { let r = unpack(&p, weights.len(), cfg); let levels = (1 << cfg.bits) as f32; let mn = weights.iter().cloned().fold(f32::INFINITY, f32::min); - let mx = weights - .iter() - .cloned() - .fold(f32::NEG_INFINITY, f32::max); + let mx = weights.iter().cloned().fold(f32::NEG_INFINITY, f32::max); let range = (mx - mn).max(1e-6); let step = range / levels; let tol = step * 0.6; @@ -559,7 +658,9 @@ mod tests { let mut xs = Vec::with_capacity(1024); let mut s = 0xCAFEBABEu64; for _ in 0..1024 { - s = s.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + s = s + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); let t = ((s >> 33) & 0x7FFFFFFF) as f32 / 0x7FFFFFFF as f32; xs.push(-1.5 + 3.0 * t); } @@ -571,10 +672,27 @@ mod tests { }; let new = pack(&xs, cfg); let old = crate::base_q4::pack(&xs); - assert_eq!(new.packed_weights.len(), old.packed_weights.len(), "weight byte count mismatch"); - assert_eq!(new.scales.len(), old.scales.len(), "scale byte count mismatch"); - assert_eq!(new.biases.len(), old.biases.len(), "bias byte count mismatch"); - for (i, (a, b)) in new.packed_weights.iter().zip(old.packed_weights.iter()).enumerate() { + assert_eq!( + new.packed_weights.len(), + old.packed_weights.len(), + "weight byte count mismatch" + ); + assert_eq!( + new.scales.len(), + old.scales.len(), + "scale byte count mismatch" + ); + assert_eq!( + new.biases.len(), + old.biases.len(), + "bias byte count mismatch" + ); + for (i, (a, b)) in new + .packed_weights + .iter() + .zip(old.packed_weights.iter()) + .enumerate() + { if a != b { panic!("weight bytes diverge at index {i}: new=0x{a:02x} old=0x{b:02x}"); } @@ -613,7 +731,11 @@ mod tests { let p = pack(&xs, cfg); assert!(p.biases.is_empty(), "symmetric must have no biases"); let r = unpack(&p, xs.len(), cfg); - let max_err = xs.iter().zip(r.iter()).map(|(a, b)| (a - b).abs()).fold(0f32, f32::max); + let max_err = xs + .iter() + .zip(r.iter()) + .map(|(a, b)| (a - b).abs()) + .fold(0f32, f32::max); assert!(max_err < 4.5, "symmetric q4 range max err {max_err}"); } diff --git a/base-convert/crates/base-quant/tests/canonical_pipeline.rs b/base-convert/crates/base-quant/tests/canonical_pipeline.rs index 4c2134a..9137ef9 100644 --- a/base-convert/crates/base-quant/tests/canonical_pipeline.rs +++ b/base-convert/crates/base-quant/tests/canonical_pipeline.rs @@ -64,8 +64,12 @@ fn make_header() -> Header { sha256: "0".repeat(64), filename: "synthetic-fp16.safetensors".into(), }, - tokenizer: TokenizerBlob { fields: BTreeMap::new() }, - config: ModelConfig { fields: BTreeMap::new() }, + tokenizer: TokenizerBlob { + fields: BTreeMap::new(), + }, + config: ModelConfig { + fields: BTreeMap::new(), + }, metadata: BTreeMap::new(), target_backend: TargetBackend::Metal, quant_profile: "smoke-q4-q8".into(), @@ -75,11 +79,17 @@ fn make_header() -> Header { tensors: vec![], mmproj: None, calibration: None, + provenance: None, sig: None, } } -fn make_entry(name: &str, dtype: TensorDtype, shape: Vec, group_size: Option) -> TensorEntry { +fn make_entry( + name: &str, + dtype: TensorDtype, + shape: Vec, + group_size: Option, +) -> TensorEntry { TensorEntry { name: name.into(), dtype, @@ -261,7 +271,10 @@ fn canonical_pipeline_round_trip_smoke() { let recon = unpack_rtn(&packed, orig_weights.len(), *cfg); let levels = (1u32 << cfg.bits) as f32; - let max = orig_weights.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let max = orig_weights + .iter() + .cloned() + .fold(f32::NEG_INFINITY, f32::max); let min = orig_weights.iter().cloned().fold(f32::INFINITY, f32::min); let step = (max - min) / levels; // Tolerance: one full step + small absolute fudge. diff --git a/base-convert/crates/base-readers/examples/gguf_inspect.rs b/base-convert/crates/base-readers/examples/gguf_inspect.rs index 668b328..4d80ee3 100644 --- a/base-convert/crates/base-readers/examples/gguf_inspect.rs +++ b/base-convert/crates/base-readers/examples/gguf_inspect.rs @@ -5,7 +5,9 @@ use base_readers::gguf::{ggml_type_name, GgufFile}; fn main() -> anyhow::Result<()> { - let path = std::env::args().nth(1).expect("usage: gguf_inspect "); + let path = std::env::args() + .nth(1) + .expect("usage: gguf_inspect "); let f = GgufFile::open(&path)?; println!("gguf v{}", f.version); println!("arch: {:?}", f.arch()); @@ -17,7 +19,11 @@ fn main() -> anyhow::Result<()> { base_readers::gguf::KvValue::Array(a) => format!("Array[{}]", a.len()), other => format!("{:?}", other), }; - let short = if short.len() > 80 { format!("{}...", &short[..80]) } else { short }; + let short = if short.len() > 80 { + format!("{}...", &short[..80]) + } else { + short + }; println!(" {:60} {}", k, short); } return Ok(()); diff --git a/base-convert/crates/base-readers/src/gguf/dequant.rs b/base-convert/crates/base-readers/src/gguf/dequant.rs index 2ffe9ac..b2814c1 100644 --- a/base-convert/crates/base-readers/src/gguf/dequant.rs +++ b/base-convert/crates/base-readers/src/gguf/dequant.rs @@ -89,11 +89,11 @@ impl GgmlType { F16 | BF16 | I16 => (1, 2), I8 => (1, 1), I64 | F64 => (1, 8), - Q4_0 => (32, 18), // fp16 d + 16 bytes (32 × 4-bit) - Q4_1 => (32, 20), // fp16 d + fp16 m + 16 bytes - Q5_0 => (32, 22), // fp16 d + 4-byte qh + 16 bytes (32 × 5-bit) + Q4_0 => (32, 18), // fp16 d + 16 bytes (32 × 4-bit) + Q4_1 => (32, 20), // fp16 d + fp16 m + 16 bytes + Q5_0 => (32, 22), // fp16 d + 4-byte qh + 16 bytes (32 × 5-bit) Q5_1 => (32, 24), - Q8_0 => (32, 34), // fp16 d + 32 × i8 + Q8_0 => (32, 34), // fp16 d + 32 × i8 Q8_1 => (32, 36), Q2K => (256, 84), Q3K => (256, 110), @@ -169,7 +169,11 @@ pub fn dequant_to_f32(info: &TensorInfo, bytes: &[u8]) -> Result> { fn f32_from_bytes(bytes: &[u8], n: usize) -> Result> { if bytes.len() < n * 4 { - bail!("F32 byte length mismatch: got {}, need {}", bytes.len(), n * 4); + bail!( + "F32 byte length mismatch: got {}, need {}", + bytes.len(), + n * 4 + ); } let mut out = Vec::with_capacity(n); for i in 0..n { @@ -183,7 +187,11 @@ fn f32_from_bytes(bytes: &[u8], n: usize) -> Result> { fn f16_from_bytes(bytes: &[u8], n: usize) -> Result> { if bytes.len() < n * 2 { - bail!("F16 byte length mismatch: got {}, need {}", bytes.len(), n * 2); + bail!( + "F16 byte length mismatch: got {}, need {}", + bytes.len(), + n * 2 + ); } let mut out = Vec::with_capacity(n); for i in 0..n { @@ -194,7 +202,11 @@ fn f16_from_bytes(bytes: &[u8], n: usize) -> Result> { fn bf16_from_bytes(bytes: &[u8], n: usize) -> Result> { if bytes.len() < n * 2 { - bail!("BF16 byte length mismatch: got {}, need {}", bytes.len(), n * 2); + bail!( + "BF16 byte length mismatch: got {}, need {}", + bytes.len(), + n * 2 + ); } let mut out = Vec::with_capacity(n); for i in 0..n { @@ -651,8 +663,8 @@ fn dequant_q6_k(bytes: &[u8], n: usize) -> Result> { let sc: &[i8] = unsafe { std::slice::from_raw_parts(bytes[base + 128 + 64..].as_ptr() as *const i8, 16) }; - let d = f16::from_le_bytes([bytes[base + 128 + 64 + 16], bytes[base + 128 + 64 + 17]]) - .to_f32(); + let d = + f16::from_le_bytes([bytes[base + 128 + 64 + 16], bytes[base + 128 + 64 + 17]]).to_f32(); let out_base = b * QK_K; // Each block processes two 128-value halves. @@ -663,17 +675,14 @@ fn dequant_q6_k(bytes: &[u8], n: usize) -> Result> { let y_off = half * 128; for l in 0..32 { let is = l / 16; - let q1 = ((ql[ql_off + l] & 0x0F) - | ((qh[qh_off + l] & 0x03) << 4)) as i32 - - 32; - let q2 = ((ql[ql_off + l + 32] & 0x0F) - | (((qh[qh_off + l] >> 2) & 0x03) << 4)) as i32 - - 32; - let q3 = ((ql[ql_off + l] >> 4) - | (((qh[qh_off + l] >> 4) & 0x03) << 4)) as i32 + let q1 = ((ql[ql_off + l] & 0x0F) | ((qh[qh_off + l] & 0x03) << 4)) as i32 - 32; + let q2 = ((ql[ql_off + l + 32] & 0x0F) | (((qh[qh_off + l] >> 2) & 0x03) << 4)) + as i32 - 32; - let q4 = ((ql[ql_off + l + 32] >> 4) - | (((qh[qh_off + l] >> 6) & 0x03) << 4)) as i32 + let q3 = + ((ql[ql_off + l] >> 4) | (((qh[qh_off + l] >> 4) & 0x03) << 4)) as i32 - 32; + let q4 = ((ql[ql_off + l + 32] >> 4) | (((qh[qh_off + l] >> 6) & 0x03) << 4)) + as i32 - 32; out[out_base + y_off + l] = d * sc[sc_off + is] as f32 * q1 as f32; out[out_base + y_off + l + 32] = d * sc[sc_off + is + 2] as f32 * q2 as f32; @@ -700,7 +709,8 @@ mod tests { -128, -64, 0, 64, 127, 1, -1, 100, 10, 20, 30, 40, 50, 60, 70, 80, 90, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ]; - bytes.extend_from_slice(unsafe { std::slice::from_raw_parts(qs.as_ptr() as *const u8, 32) }); + bytes + .extend_from_slice(unsafe { std::slice::from_raw_parts(qs.as_ptr() as *const u8, 32) }); let out = dequant_q8_0(&bytes, 32).unwrap(); assert_eq!(out[0], -256.0); assert_eq!(out[1], -128.0); diff --git a/base-convert/crates/base-readers/src/gguf/parse.rs b/base-convert/crates/base-readers/src/gguf/parse.rs index ea1c86e..5ab6184 100644 --- a/base-convert/crates/base-readers/src/gguf/parse.rs +++ b/base-convert/crates/base-readers/src/gguf/parse.rs @@ -48,8 +48,8 @@ pub struct GgufFile { impl GgufFile { pub fn open>(path: P) -> Result { - let file = File::open(path.as_ref()) - .with_context(|| format!("opening {:?}", path.as_ref()))?; + let file = + File::open(path.as_ref()).with_context(|| format!("opening {:?}", path.as_ref()))?; let mmap = unsafe { Mmap::map(&file)? }; Self::from_mmap(mmap) } @@ -292,4 +292,16 @@ impl KvValue { _ => None, } } + + pub fn as_bool(&self) -> Option { + match self { + KvValue::Bool(b) => Some(*b), + // Some GGUF writers store boolean flags as small integers. + KvValue::U8(n) => Some(*n != 0), + KvValue::I8(n) => Some(*n != 0), + KvValue::U32(n) => Some(*n != 0), + KvValue::I32(n) => Some(*n != 0), + _ => None, + } + } } diff --git a/base-convert/crates/base-readers/src/hf.rs b/base-convert/crates/base-readers/src/hf.rs index fd09b92..a5f8678 100644 --- a/base-convert/crates/base-readers/src/hf.rs +++ b/base-convert/crates/base-readers/src/hf.rs @@ -55,8 +55,15 @@ impl HfDir { } let config_path = dir.join("config.json"); - let config_bytes = std::fs::read(&config_path) - .with_context(|| format!("reading {:?}", config_path))?; + let config_bytes = + std::fs::read(&config_path).with_context(|| format!("reading {:?}", config_path))?; + let (config_bytes, n_nonfinite) = sanitize_python_json(&config_bytes); + if n_nonfinite > 0 { + eprintln!( + " note: config.json contains {n_nonfinite} bare Infinity/NaN literal(s) \ + (Python's json.dump emits these; they are not valid JSON) — read as null" + ); + } let config: serde_json::Value = serde_json::from_slice(&config_bytes).context("parsing config.json")?; @@ -72,10 +79,7 @@ impl HfDir { let chat_template_jinja = { let p = dir.join("chat_template.jinja"); if p.exists() { - Some( - std::fs::read_to_string(&p) - .with_context(|| format!("reading {:?}", p))?, - ) + Some(std::fs::read_to_string(&p).with_context(|| format!("reading {:?}", p))?) } else { tokenizer_config .as_ref() @@ -87,47 +91,44 @@ impl HfDir { // Discover shards. Prefer the index.json path if present. let index_path = dir.join("model.safetensors.index.json"); - let (shard_paths, routing): (Vec, Option>) = if index_path - .exists() - { - let idx_bytes = std::fs::read(&index_path)?; - let idx: ShardIndex = - serde_json::from_slice(&idx_bytes).context("parsing shard index")?; - let mut shards: Vec = idx - .weight_map - .values() - .cloned() - .collect::>() - .into_iter() - .map(|name| dir.join(name)) - .collect(); - shards.sort(); - (shards, Some(idx.weight_map)) - } else { - // Glob for model*.safetensors. - let mut paths: Vec = std::fs::read_dir(&dir)? - .filter_map(|e| e.ok().map(|e| e.path())) - .filter(|p| { - p.extension().and_then(|e| e.to_str()) == Some("safetensors") - && p.file_name() - .and_then(|n| n.to_str()) - .map(|n| n.starts_with("model")) - .unwrap_or(false) - }) - .collect(); - paths.sort(); - if paths.is_empty() { - bail!("no .safetensors shards in {:?}", dir); - } - (paths, None) - }; + let (shard_paths, routing): (Vec, Option>) = + if index_path.exists() { + let idx_bytes = std::fs::read(&index_path)?; + let idx: ShardIndex = + serde_json::from_slice(&idx_bytes).context("parsing shard index")?; + let mut shards: Vec = idx + .weight_map + .values() + .cloned() + .collect::>() + .into_iter() + .map(|name| dir.join(name)) + .collect(); + shards.sort(); + (shards, Some(idx.weight_map)) + } else { + // Glob for model*.safetensors. + let mut paths: Vec = std::fs::read_dir(&dir)? + .filter_map(|e| e.ok().map(|e| e.path())) + .filter(|p| { + p.extension().and_then(|e| e.to_str()) == Some("safetensors") + && p.file_name() + .and_then(|n| n.to_str()) + .map(|n| n.starts_with("model")) + .unwrap_or(false) + }) + .collect(); + paths.sort(); + if paths.is_empty() { + bail!("no .safetensors shards in {:?}", dir); + } + (paths, None) + }; let mut shards = Vec::with_capacity(shard_paths.len()); for p in &shard_paths { - shards.push( - SafetensorsFile::open(p) - .with_context(|| format!("opening shard {:?}", p))?, - ); + shards + .push(SafetensorsFile::open(p).with_context(|| format!("opening shard {:?}", p))?); } // Build a name → (shard_idx, tensor_idx) lookup. @@ -237,5 +238,118 @@ fn read_optional_json(path: &Path) -> Result> { return Ok(None); } let bytes = std::fs::read(path)?; + let (bytes, _) = sanitize_python_json(&bytes); Ok(Some(serde_json::from_slice(&bytes)?)) } + +/// Rewrite Python's non-standard JSON literals (`Infinity`, `-Infinity`, +/// `NaN`) to `null` so a strict parser accepts the document. +/// +/// `json.dump` emits these by default and plenty of published checkpoints +/// carry them — NVIDIA's Nemotron-H config has +/// `"time_step_limit": [0.0, Infinity]`. They are not valid JSON (RFC 8259 +/// has no non-finite numbers) and `serde_json` rejects the whole file, so +/// without this the model cannot be read at all. +/// +/// `null` rather than a large finite float: consumers read this config as +/// a `serde_json::Value` and query keys with `.as_f64()`, so null reads as +/// "absent" and a consumer that actually needs the value sees nothing +/// rather than a silently wrong number. +/// +/// Replacement happens only outside string literals, so a key or value +/// whose *text* contains `NaN` is untouched. Returns the rewritten bytes +/// and how many literals were replaced. +pub fn sanitize_python_json(bytes: &[u8]) -> (Vec, usize) { + const LITERALS: [&[u8]; 3] = [b"-Infinity", b"Infinity", b"NaN"]; + + // Fast path: nothing to do for the overwhelming majority of files. + if !bytes.windows(3).any(|w| w == b"NaN") && !bytes.windows(8).any(|w| w == b"Infinity") { + return (bytes.to_vec(), 0); + } + + let mut out = Vec::with_capacity(bytes.len()); + let mut count = 0usize; + let mut i = 0usize; + let mut in_string = false; + while i < bytes.len() { + let b = bytes[i]; + if in_string { + out.push(b); + if b == b'\\' && i + 1 < bytes.len() { + // Copy the escaped character verbatim so an escaped quote + // does not look like the end of the string. + out.push(bytes[i + 1]); + i += 2; + continue; + } + if b == b'"' { + in_string = false; + } + i += 1; + continue; + } + if b == b'"' { + in_string = true; + out.push(b); + i += 1; + continue; + } + if let Some(lit) = LITERALS.iter().find(|lit| bytes[i..].starts_with(lit)) { + out.extend_from_slice(b"null"); + count += 1; + i += lit.len(); + continue; + } + out.push(b); + i += 1; + } + (out, count) +} + +#[cfg(test)] +mod sanitize_tests { + use super::sanitize_python_json; + + fn s(input: &str) -> (String, usize) { + let (bytes, n) = sanitize_python_json(input.as_bytes()); + (String::from_utf8(bytes).unwrap(), n) + } + + #[test] + fn rewrites_nemotron_time_step_limit() { + let (out, n) = s(r#"{"time_step_limit": [0.0, Infinity]}"#); + assert_eq!(out, r#"{"time_step_limit": [0.0, null]}"#); + assert_eq!(n, 1); + serde_json::from_str::(&out).unwrap(); + } + + #[test] + fn rewrites_negative_infinity_and_nan() { + let (out, n) = s(r#"{"a": -Infinity, "b": NaN}"#); + assert_eq!(out, r#"{"a": null, "b": null}"#); + assert_eq!(n, 2); + } + + #[test] + fn leaves_strings_alone() { + // The words appear inside string literals, where they are data. + let (out, n) = s(r#"{"note": "NaN and Infinity", "x": Infinity}"#); + assert_eq!(out, r#"{"note": "NaN and Infinity", "x": null}"#); + assert_eq!(n, 1); + } + + #[test] + fn escaped_quote_does_not_end_the_string() { + let (out, n) = s(r#"{"note": "a \" NaN", "x": 1}"#); + assert_eq!(out, r#"{"note": "a \" NaN", "x": 1}"#); + assert_eq!(n, 0); + } + + #[test] + fn valid_json_is_untouched() { + let src = r#"{"a": 1.0, "b": [1, 2], "c": "text"}"#; + let (out, n) = s(src); + assert_eq!(out, src); + assert_eq!(n, 0); + } +} diff --git a/base-convert/crates/base-readers/src/lib.rs b/base-convert/crates/base-readers/src/lib.rs index b94c6eb..f65ad1a 100644 --- a/base-convert/crates/base-readers/src/lib.rs +++ b/base-convert/crates/base-readers/src/lib.rs @@ -30,8 +30,15 @@ pub fn detect_format(path: &std::path::Path) -> anyhow::Result { anyhow::bail!("{:?} is a directory without config.json", path); } let bytes = std::fs::read(&cfg_path)?; + // Python's `json.dump` writes bare Infinity/NaN, which strict JSON + // has no syntax for; published checkpoints do carry them. + let (bytes, _) = crate::hf::sanitize_python_json(&bytes); let cfg: serde_json::Value = serde_json::from_slice(&bytes)?; - if cfg.get("quantization").and_then(|v| v.get("bits")).is_some() { + if cfg + .get("quantization") + .and_then(|v| v.get("bits")) + .is_some() + { Ok(SourceFormat::MlxSafetensors) } else { Ok(SourceFormat::HfSafetensors) diff --git a/base-convert/crates/base-readers/src/mlx.rs b/base-convert/crates/base-readers/src/mlx.rs index 14f0fd1..8f642ab 100644 --- a/base-convert/crates/base-readers/src/mlx.rs +++ b/base-convert/crates/base-readers/src/mlx.rs @@ -48,6 +48,21 @@ pub struct MlxDir { pub quant: MlxQuant, } +/// Borrowed view of an MLX-quantized tensor's raw storage, for +/// passthrough conversion (see [`MlxDir::tensor_packed`]). +pub struct MlxPackedTensor<'a> { + /// Packed weight bytes (nibble stream, low-nibble first). + pub packed: &'a [u8], + /// Per-group scales, stored as `scale_dtype` (BF16 on current + /// checkpoints, F16 on pre-0.20 ones). + pub scales: &'a [u8], + /// Per-group biases, same dtype as scales. + pub biases: &'a [u8], + pub group_size: u32, + pub bits: u32, + pub scale_dtype: StDtype, +} + impl MlxDir { pub fn open>(dir: P) -> Result { let hf = HfDir::open(dir)?; @@ -183,10 +198,9 @@ impl MlxDir { .tensor_bytes(&scales_name) .with_context(|| format!("tensor_bytes({scales_name}) missing after tensor_info()"))?; let biases_bytes = match biases_dtype { - Some(_) => self - .hf - .tensor_bytes(&biases_name) - .with_context(|| format!("tensor_bytes({biases_name}) missing after tensor_info()"))?, + Some(_) => self.hf.tensor_bytes(&biases_name).with_context(|| { + format!("tensor_bytes({biases_name}) missing after tensor_info()") + })?, None => &[], }; @@ -228,8 +242,7 @@ impl MlxDir { word |= (packed_bytes[byte0 + 1] as u32) << 8; } let q = (word >> sh) & mask; - out[o_base + i * in_features + j] = - (q as f32) * scale + bias; + out[o_base + i * in_features + j] = (q as f32) * scale + bias; } } } @@ -237,6 +250,70 @@ impl MlxDir { Ok(out) } + /// Raw packed payload of an MLX-quantized tensor, verbatim. + /// + /// MLX's `U32 [.., in/(32/bits)]` little-endian nibble packing is + /// byte-identical to `base_q4`'s two-per-byte low-nibble-first + /// stream (verified against `mx.dequantize` at 0.0 difference), so + /// a passthrough conversion can reuse these bytes without the + /// dequant→requant round trip that costs ~4.4% of weight RMS. + /// + /// Returns `Ok(None)` for unquantized tensors (no `.scales` + /// sibling). Errors if the tensor's (bits, group_size) differ from + /// `expect_bits`/`expect_group_size` — a passthrough caller must + /// fail loudly rather than silently requantize, or the "weights + /// are bit-identical to the source" contract breaks. + pub fn tensor_packed( + &self, + name: &str, + expect_bits: u32, + expect_group_size: u32, + ) -> Result>> { + let Some(scales_name) = quant_sibling(name, "scales") else { + return Ok(None); + }; + if self.hf.tensor_info(&scales_name).is_none() { + return Ok(None); + } + let q = self.quant_for_tensor(name); + if q.bits != expect_bits || q.group_size != expect_group_size { + bail!( + "MLX tensor {name:?} is {}-bit gs={} — passthrough expects {}-bit gs={}", + q.bits, + q.group_size, + expect_bits, + expect_group_size + ); + } + let scales_info = self + .hf + .tensor_info(&scales_name) + .with_context(|| format!("scales {scales_name} missing"))?; + let scale_dtype = scales_info.dtype; + let biases_name = quant_sibling(name, "biases") + .with_context(|| format!("expected `.weight`-suffixed name, got {name}"))?; + let packed = self + .hf + .tensor_bytes(name) + .with_context(|| format!("tensor_bytes({name}) missing"))?; + let scales = self + .hf + .tensor_bytes(&scales_name) + .with_context(|| format!("tensor_bytes({scales_name}) missing"))?; + let biases = self + .hf + .tensor_bytes(&biases_name) + .with_context(|| format!("MLX affine tensor {name:?} has scales but no biases"))?; + Ok(Some(MlxPackedTensor { + packed, + scales, + biases, + group_size: q.group_size, + bits: q.bits, + scale_dtype, + })) + } + /// Logical shape of an MLX-packed tensor (unpacking the last dim). /// Returns `Some(shape)` if the tensor is packed, `None` otherwise. /// Resolves bits per-tensor — Gemma 4 26B-A4B 4-bit checkpoints @@ -246,8 +323,7 @@ impl MlxDir { /// logical in_features the runtime expects. pub fn unpacked_shape(&self, name: &str) -> Option> { let info = self.hf.tensor_info(name)?; - quant_sibling(name, "scales") - .and_then(|sn| self.hf.tensor_info(&sn))?; + quant_sibling(name, "scales").and_then(|sn| self.hf.tensor_info(&sn))?; if info.shape.len() < 2 { return None; } @@ -264,6 +340,169 @@ impl MlxDir { Some(shape) } + /// Zero-loss transplant of an MLX affine-quantized tensor into + /// `base_q4`'s on-disk layout. + /// + /// `base_q4` and MLX-affine 4-bit are the *same* scheme: INT4 + /// asymmetric, `value = q * scale + bias`, one f16 scale and f16 + /// bias per group of 64. At 4 bits MLX's little-endian bitstream is + /// byte-for-byte `base_q4`'s low-nibble-first packing + /// (`byte = (q[2i+1] << 4) | q[2i]`), so the weight bytes transplant + /// verbatim and only the scale/bias regions need re-laying-out. + /// + /// Taking this path instead of dequant → requant matters for more + /// than speed: re-deriving `scale = (max - min) / 15` from already + /// quantized values lands on a *different* grid whenever a group's + /// codes don't span the full 0..15 range, so the round trip is not + /// the identity. Transplanting reproduces the reference engine's + /// weights exactly, which is what makes an MLX-vs-baseRT numerical + /// comparison attributable to engine math. + /// + /// Returns `None` (rather than an error) whenever the tensor is not + /// an exact match for the target scheme — different bits, a + /// different group size, a symmetric tensor with no `.biases`, or a + /// plain unquantized tensor. Callers fall back to dequant → requant. + pub fn packed_base_q4(&self, name: &str, group_size: u32) -> Result> { + let Some(scales_name) = quant_sibling(name, "scales") else { + return Ok(None); + }; + let Some(scales_info) = self.hf.tensor_info(&scales_name) else { + return Ok(None); // not MLX-packed + }; + let q = self.quant_for_tensor(name); + if q.bits != 4 || q.group_size != group_size { + return Ok(None); // 8-bit override, or a group size base_q4 can't express + } + let Some(biases_name) = quant_sibling(name, "biases") else { + return Ok(None); + }; + let Some(biases_info) = self.hf.tensor_info(&biases_name) else { + return Ok(None); // symmetric tensor — base_q4 is asymmetric + }; + let packed = self + .hf + .tensor_info(name) + .with_context(|| format!("packed tensor {name} missing"))?; + if packed.shape.len() < 2 { + return Ok(None); + } + + let group_size_usize = group_size as usize; + let (batch_dims, packed_last) = packed.shape.split_at(packed.shape.len() - 1); + let packed_in = packed_last[0] as usize; + let (batch_dims_split, out_dim_slice) = batch_dims.split_at(batch_dims.len() - 1); + let out_features = out_dim_slice[0] as usize; + // 4-bit: 8 codes per u32. + let in_features = packed_in * 8; + if in_features % group_size_usize != 0 { + return Ok(None); + } + let batch = (batch_dims_split.iter().product::() as usize).max(1); + let total_values = batch * out_features * in_features; + let n_groups = total_values / group_size_usize; + + let packed_bytes = self + .hf + .tensor_bytes(name) + .with_context(|| format!("tensor_bytes({name}) missing after tensor_info()"))?; + // The transplant is only sound if the source is exactly as large + // as the layout implies — a short/long buffer means our shape + // arithmetic disagrees with the file, and copying it verbatim + // would silently produce a corrupt bundle. + if packed_bytes.len() != total_values / 2 { + bail!( + "MLX packed tensor {:?}: {} weight bytes but shape {:?} implies {} \ + (4-bit codes, 2 per byte)", + name, + packed_bytes.len(), + packed.shape, + total_values / 2 + ); + } + let scales_bytes = self + .hf + .tensor_bytes(&scales_name) + .with_context(|| format!("tensor_bytes({scales_name}) missing after tensor_info()"))?; + let biases_bytes = self + .hf + .tensor_bytes(&biases_name) + .with_context(|| format!("tensor_bytes({biases_name}) missing after tensor_info()"))?; + if scales_bytes.len() != n_groups * 2 || biases_bytes.len() != n_groups * 2 { + bail!( + "MLX packed tensor {:?}: scales/biases are {}/{} bytes but shape {:?} implies \ + {} groups of {} (2 bytes each)", + name, + scales_bytes.len(), + biases_bytes.len(), + packed.shape, + n_groups, + group_size + ); + } + + // base_q4 stores f16 scales and biases. F16 sources copy + // verbatim; bf16 sources (mlx-lm ≳ 0.20) are widened to f32 and + // renarrowed — exact in the mantissa (bf16 has 8 bits, f16 has + // 11) but bf16's wider exponent range can overflow to inf or + // flush to zero, so report how many did. + let (scales, scales_narrowed) = narrow_to_f16_le(scales_bytes, scales_info.dtype)?; + let (biases, biases_narrowed) = narrow_to_f16_le(biases_bytes, biases_info.dtype)?; + + let out_of_f16_range = count_non_finite(&scales) + count_non_finite(&biases); + + Ok(Some(MlxPackedQ4 { + packed_weights: packed_bytes.to_vec(), + scales, + biases, + group_size, + narrowed_from_bf16: scales_narrowed || biases_narrowed, + out_of_f16_range, + })) + } +} + +/// An MLX affine-q4 tensor in `base_q4`'s on-disk layout, ready to write +/// without a dequant → requant round trip. Field names mirror +/// `base_quant::Packed` so the caller's conversion is a plain move. +pub struct MlxPackedQ4 { + /// 4-bit codes, two per byte, low nibble first — MLX's bytes verbatim. + pub packed_weights: Vec, + /// One f16 scale per group, little-endian. + pub scales: Vec, + /// One f16 bias per group, little-endian. + pub biases: Vec, + pub group_size: u32, + /// Source scales/biases were bf16 and had to be renarrowed to f16. + pub narrowed_from_bf16: bool, + /// Scales/biases that left f16's representable range in the process. + pub out_of_f16_range: usize, +} + +/// Re-encode a buffer of f16-or-bf16 halves as little-endian f16. +/// Returns `(bytes, narrowed)` where `narrowed` marks a bf16 source. +/// Public so `--validate` can compare against the bytes the WRITE path +/// actually emitted: a bf16-scaled checkpoint is narrowed on the way in, so a +/// byte-compare against the raw source scales would fail on every tensor. +pub fn narrow_to_f16_le(bytes: &[u8], dtype: StDtype) -> Result<(Vec, bool)> { + match dtype { + StDtype::F16 => Ok((bytes.to_vec(), false)), + StDtype::Bf16 => { + let mut out = Vec::with_capacity(bytes.len()); + for c in bytes.chunks_exact(2) { + let v = bf16::from_le_bytes([c[0], c[1]]).to_f32(); + out.extend_from_slice(&f16::from_f32(v).to_le_bytes()); + } + Ok((out, true)) + } + other => bail!("MLX scales/biases have unsupported dtype {other:?} (expected f16 or bf16)"), + } +} + +fn count_non_finite(f16_le: &[u8]) -> usize { + f16_le + .chunks_exact(2) + .filter(|c| !f16::from_le_bytes([c[0], c[1]]).to_f32().is_finite()) + .count() } fn quant_sibling(name: &str, suffix: &str) -> Option { diff --git a/base-convert/crates/base-readers/src/safetensors.rs b/base-convert/crates/base-readers/src/safetensors.rs index 622b254..e580c3b 100644 --- a/base-convert/crates/base-readers/src/safetensors.rs +++ b/base-convert/crates/base-readers/src/safetensors.rs @@ -44,6 +44,10 @@ pub enum StDtype { U32, U64, Bool, + /// fp8 e4m3 (NVFP4 checkpoints store per-block scales in this dtype). + /// Carried as raw bytes; there is no f32 decode path. + F8E4m3, + F8E5m2, } impl StDtype { @@ -66,6 +70,8 @@ impl StDtype { "U32" => StDtype::U32, "U64" => StDtype::U64, "BOOL" => StDtype::Bool, + "F8_E4M3" => StDtype::F8E4m3, + "F8_E5M2" => StDtype::F8E5m2, other => bail!("unknown safetensors dtype: {other}"), }) } @@ -75,7 +81,7 @@ impl StDtype { StDtype::F32 | StDtype::I32 | StDtype::U32 => 4, StDtype::F16 | StDtype::Bf16 | StDtype::I16 | StDtype::U16 => 2, StDtype::F64 | StDtype::I64 | StDtype::U64 => 8, - StDtype::I8 | StDtype::U8 | StDtype::Bool => 1, + StDtype::I8 | StDtype::U8 | StDtype::Bool | StDtype::F8E4m3 | StDtype::F8E5m2 => 1, } } } @@ -106,8 +112,8 @@ pub struct SafetensorsFile { impl SafetensorsFile { pub fn open>(path: P) -> Result { - let file = File::open(path.as_ref()) - .with_context(|| format!("opening {:?}", path.as_ref()))?; + let file = + File::open(path.as_ref()).with_context(|| format!("opening {:?}", path.as_ref()))?; let mmap = unsafe { Mmap::map(&file)? }; Self::from_mmap(mmap) } diff --git a/base-convert/crates/base-readers/tests/mlx_dequant.rs b/base-convert/crates/base-readers/tests/mlx_dequant.rs index 64880d1..5eeb4e0 100644 --- a/base-convert/crates/base-readers/tests/mlx_dequant.rs +++ b/base-convert/crates/base-readers/tests/mlx_dequant.rs @@ -69,7 +69,11 @@ fn mlx_dir( "model_type": "test", "quantization": { "bits": bits, "group_size": 32 }, }); - std::fs::write(dir.join("config.json"), serde_json::to_vec(&config).unwrap()).unwrap(); + std::fs::write( + dir.join("config.json"), + serde_json::to_vec(&config).unwrap(), + ) + .unwrap(); write_safetensors( &dir.join("model.safetensors"), &[ @@ -175,3 +179,120 @@ fn unsupported_bits_is_a_clear_error() { let err = mlx.tensor_to_f32("layer.weight").unwrap_err().to_string(); assert!(err.contains("unsupported bits=7"), "got: {err}"); } + +/// Decode a `base_q4` payload the way the runtime kernels do: +/// `value = q * scale + bias`, codes two per byte, low nibble first. +/// Deliberately independent of `base-quant` (which this crate does not +/// depend on) so the test checks the *layout contract*, not shared code. +fn decode_base_q4(packed: &[u8], scales: &[u8], biases: &[u8], group_size: usize) -> Vec { + let total = packed.len() * 2; + (0..total) + .map(|i| { + let byte = packed[i / 2]; + let q = if i % 2 == 0 { byte & 0x0F } else { byte >> 4 }; + let g = i / group_size; + let scale = half::f16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]).to_f32(); + let bias = half::f16::from_le_bytes([biases[g * 2], biases[g * 2 + 1]]).to_f32(); + (q as f32) * scale + bias + }) + .collect() +} + +/// The whole point of the transplant: MLX's 4-bit bytes, reinterpreted +/// as `base_q4`, decode to exactly what MLX itself produces. Checked +/// against `B4_EXPECTED` — mlx.core.dequantize's own output — so this +/// pins the nibble order and scale/bias layout against ground truth +/// rather than against our own reader. +#[test] +fn passthrough_decodes_identically_to_mlx_4bit() { + let tmp = tempfile::tempdir().unwrap(); + let mlx = mlx_dir( + tmp.path(), + 4, + &[3, 8], + &[3, 2], + B4_PACKED, + B4_SCALES, + B4_BIASES, + ); + + let p = mlx + .packed_base_q4("layer.weight", 32) + .unwrap() + .expect("4-bit gs=32 tensor should transplant"); + assert_eq!(p.group_size, 32); + assert!(!p.narrowed_from_bf16, "f16 scales copy verbatim"); + assert_eq!(p.out_of_f16_range, 0); + // Verbatim means verbatim: the weight bytes are MLX's own. + assert_eq!(p.packed_weights, u32_bytes(B4_PACKED)); + assert_eq!(p.scales, u16_bytes(B4_SCALES)); + assert_eq!(p.biases, u16_bytes(B4_BIASES)); + + let got = decode_base_q4(&p.packed_weights, &p.scales, &p.biases, 32); + assert_eq!(got.len(), B4_EXPECTED.len()); + for (i, (g, e)) in got.iter().zip(B4_EXPECTED).enumerate() { + assert!( + (g - e).abs() <= TOL, + "element {i}: base_q4 decode {g}, mlx says {e}" + ); + } +} + +/// Anything that is not bit-for-bit the same scheme must decline rather +/// than hand over bytes that would be misread. The caller then falls +/// back to dequant → requant. +#[test] +fn passthrough_declines_on_scheme_mismatch() { + let tmp = tempfile::tempdir().unwrap(); + let mlx = mlx_dir( + tmp.path(), + 4, + &[3, 8], + &[3, 2], + B4_PACKED, + B4_SCALES, + B4_BIASES, + ); + // Right bits, wrong group size: base_q4 at gs=64 cannot express a + // tensor whose scales are per-32. + assert!(mlx.packed_base_q4("layer.weight", 64).unwrap().is_none()); + + // Wrong bits: 6-bit codes cross byte boundaries, so the bytes are + // not a base_q4 nibble stream. + let tmp6 = tempfile::tempdir().unwrap(); + let mlx6 = mlx_dir( + tmp6.path(), + 6, + &[3, 12], + &[3, 2], + B6_PACKED, + B6_SCALES, + B6_BIASES, + ); + assert!(mlx6.packed_base_q4("layer.weight", 32).unwrap().is_none()); +} + +/// A symmetric MLX tensor ships no `.biases`; `base_q4` is asymmetric +/// and has nowhere to get the missing half of the affine pair. +#[test] +fn passthrough_declines_without_biases() { + let tmp = tempfile::tempdir().unwrap(); + let config = serde_json::json!({ + "model_type": "test", + "quantization": { "bits": 4, "group_size": 32 }, + }); + std::fs::write( + tmp.path().join("config.json"), + serde_json::to_vec(&config).unwrap(), + ) + .unwrap(); + write_safetensors( + &tmp.path().join("model.safetensors"), + &[ + ("layer.weight", "U32", &[3, 8], u32_bytes(B4_PACKED)), + ("layer.scales", "F16", &[3, 2], u16_bytes(B4_SCALES)), + ], + ); + let mlx = MlxDir::open(tmp.path()).unwrap(); + assert!(mlx.packed_base_q4("layer.weight", 32).unwrap().is_none()); +} diff --git a/base-convert/crates/base-sign/src/lib.rs b/base-convert/crates/base-sign/src/lib.rs index 056f98b..72d4b7a 100644 --- a/base-convert/crates/base-sign/src/lib.rs +++ b/base-convert/crates/base-sign/src/lib.rs @@ -12,9 +12,7 @@ //! at verify time, as long as the canonicalizer is deterministic. use anyhow::{Context, Result}; -use ed25519_dalek::{ - Signature, Signer, SigningKey, Verifier, VerifyingKey, SECRET_KEY_LENGTH, -}; +use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey, SECRET_KEY_LENGTH}; use sha2::{Digest, Sha256}; /// Sign a payload (canonical JSON || sha256(blob)) with an ed25519 key. @@ -70,8 +68,7 @@ pub fn signing_key_from_bytes(bytes: &[u8]) -> Result { pub fn b64_encode(bytes: &[u8]) -> String { // Minimal base64 implementation to avoid pulling in a crate just // for this. Standard alphabet, with `=` padding. - const ALPHABET: &[u8; 64] = - b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4); let mut i = 0; while i + 3 <= bytes.len() { @@ -226,20 +223,20 @@ pub fn sign_base_file>( pub fn verify_base_file>(path: P, key: &VerifyingKey) -> Result<()> { use base_format::BaseReader; - let reader = BaseReader::open(path.as_ref()) - .with_context(|| format!("opening {:?}", path.as_ref()))?; + let reader = + BaseReader::open(path.as_ref()).with_context(|| format!("opening {:?}", path.as_ref()))?; let Some(recorded_sig) = reader.header().sig.clone() else { return Ok(()); }; if recorded_sig.alg != "ed25519" { - anyhow::bail!( - "unsupported signature algorithm: {:?}", - recorded_sig.alg - ); + anyhow::bail!("unsupported signature algorithm: {:?}", recorded_sig.alg); } let sig_bytes = b64_decode(&recorded_sig.signature)?; if sig_bytes.len() != 64 { - anyhow::bail!("ed25519 signature must be 64 bytes, got {}", sig_bytes.len()); + anyhow::bail!( + "ed25519 signature must be 64 bytes, got {}", + sig_bytes.len() + ); } let mut sig_arr = [0u8; 64]; sig_arr.copy_from_slice(&sig_bytes); diff --git a/base-convert/crates/base-sign/tests/sign_roundtrip.rs b/base-convert/crates/base-sign/tests/sign_roundtrip.rs index bdb8168..10d0461 100644 --- a/base-convert/crates/base-sign/tests/sign_roundtrip.rs +++ b/base-convert/crates/base-sign/tests/sign_roundtrip.rs @@ -35,6 +35,7 @@ fn make_header() -> Header { tensors: vec![], mmproj: None, calibration: None, + provenance: None, sig: None, } } diff --git a/bindings/node/package.json b/bindings/node/package.json index 9f78847..3ea8407 100644 --- a/bindings/node/package.json +++ b/bindings/node/package.json @@ -1,6 +1,6 @@ { "name": "@baseRT/node", - "version": "0.2.3", + "version": "0.2.4", "private": true, "description": "Node.js bindings for BaseRT — LLM inference engine for Apple Silicon (Metal)", "main": "dist/index.js", diff --git a/bindings/node/src/index.ts b/bindings/node/src/index.ts index d4763d5..6464a96 100644 --- a/bindings/node/src/index.ts +++ b/bindings/node/src/index.ts @@ -117,6 +117,12 @@ const ModelConfigC = koffi.struct("BaseRTModelConfig", { gdn_key_head_dim: "uint32", gdn_value_head_dim: "uint32", gdn_conv_kernel: "uint32", + // Nemotron-H hybrid Mamba-2 SSM + ssm_state_size: "uint32", + ssm_conv_kernel: "uint32", + ssm_num_groups: "uint32", + ssm_inner_size: "uint32", + ssm_num_heads: "uint32", // MoE n_experts: "uint32", n_experts_used: "uint32", @@ -125,6 +131,7 @@ const ModelConfigC = koffi.struct("BaseRTModelConfig", { expert_gating: "uint8", norm_topk_prob: "uint8", _moe_pad: koffi.array("uint8", 2), + expert_weights_scale: "float", // Vision tower vision_n_layers: "uint32", vision_dim: "uint32", @@ -186,6 +193,27 @@ const ModelConfigC = koffi.struct("BaseRTModelConfig", { vision_pos_embed_w: "uint32", vision_adapter_dim: "uint32", video_token_id: "uint32", + // GLM 5.2 / glm-dsa + q_lora_rank: "uint32", + kv_lora_rank: "uint32", + qk_nope_head_dim: "uint32", + qk_rope_head_dim: "uint32", + v_head_dim: "uint32", + routed_scaling_factor: "float", + first_k_dense_replace: "uint32", + nextn_predict_layers: "uint32", + indexer_head_count: "uint32", + indexer_key_length: "uint32", + indexer_top_k: "uint32", + // gpt-oss + rope_yarn_beta_fast: "float", + rope_yarn_beta_slow: "float", + swiglu_limit: "float", + swiglu_alpha: "float", + attention_sinks: "uint8", + attention_bias: "uint8", + rope_yarn_truncate: "uint8", + _gptoss_pad: "uint8", }); const SamplingConfigC = koffi.struct("BaseRTSamplingConfig", { diff --git a/bindings/python/baseRT/__init__.py b/bindings/python/baseRT/__init__.py index 1e902bd..78da23e 100644 --- a/bindings/python/baseRT/__init__.py +++ b/bindings/python/baseRT/__init__.py @@ -21,7 +21,7 @@ from pathlib import Path from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union -__version__ = "0.2.3" +__version__ = "0.2.4" # --------------------------------------------------------------------------- # Library loading @@ -138,6 +138,12 @@ class BaseRTModelConfig(ctypes.Structure): ("gdn_key_head_dim", ctypes.c_uint32), ("gdn_value_head_dim", ctypes.c_uint32), ("gdn_conv_kernel", ctypes.c_uint32), + # Nemotron-H hybrid Mamba-2 SSM (0 = not a Mamba-2 hybrid) + ("ssm_state_size", ctypes.c_uint32), + ("ssm_conv_kernel", ctypes.c_uint32), + ("ssm_num_groups", ctypes.c_uint32), + ("ssm_inner_size", ctypes.c_uint32), + ("ssm_num_heads", ctypes.c_uint32), # Mixture-of-Experts (0 = dense) ("n_experts", ctypes.c_uint32), ("n_experts_used", ctypes.c_uint32), @@ -146,6 +152,7 @@ class BaseRTModelConfig(ctypes.Structure): ("expert_gating", ctypes.c_uint8), ("norm_topk_prob", ctypes.c_uint8), ("_moe_pad", ctypes.c_uint8 * 2), + ("expert_weights_scale", ctypes.c_float), # Vision tower ("vision_n_layers", ctypes.c_uint32), ("vision_dim", ctypes.c_uint32), @@ -207,6 +214,27 @@ class BaseRTModelConfig(ctypes.Structure): ("vision_pos_embed_w", ctypes.c_uint32), ("vision_adapter_dim", ctypes.c_uint32), ("video_token_id", ctypes.c_uint32), + # GLM 5.2 / glm-dsa + ("q_lora_rank", ctypes.c_uint32), + ("kv_lora_rank", ctypes.c_uint32), + ("qk_nope_head_dim", ctypes.c_uint32), + ("qk_rope_head_dim", ctypes.c_uint32), + ("v_head_dim", ctypes.c_uint32), + ("routed_scaling_factor", ctypes.c_float), + ("first_k_dense_replace", ctypes.c_uint32), + ("nextn_predict_layers", ctypes.c_uint32), + ("indexer_head_count", ctypes.c_uint32), + ("indexer_key_length", ctypes.c_uint32), + ("indexer_top_k", ctypes.c_uint32), + # gpt-oss + ("rope_yarn_beta_fast", ctypes.c_float), + ("rope_yarn_beta_slow", ctypes.c_float), + ("swiglu_limit", ctypes.c_float), + ("swiglu_alpha", ctypes.c_float), + ("attention_sinks", ctypes.c_uint8), + ("attention_bias", ctypes.c_uint8), + ("rope_yarn_truncate", ctypes.c_uint8), + ("_gptoss_pad", ctypes.c_uint8), ] diff --git a/bindings/python/setup.py b/bindings/python/setup.py index 6ec7381..cb25a04 100644 --- a/bindings/python/setup.py +++ b/bindings/python/setup.py @@ -6,7 +6,7 @@ setup( name="baseRT", - version="0.2.3", + version="0.2.4", description="Python bindings for the BaseRT LLM inference engine (Apple Silicon / Metal)", long_description=long_description, long_description_content_type="text/markdown", diff --git a/bindings/python/tests/test_baseRT.py b/bindings/python/tests/test_baseRT.py index 55f40a8..564027d 100644 --- a/bindings/python/tests/test_baseRT.py +++ b/bindings/python/tests/test_baseRT.py @@ -80,7 +80,7 @@ def test_field_count(self): # Keep this in sync with include/baseRT/types.h. The runtime ABI-size # assertion below catches layout drift; this count catches accidental # omission of same-sized fields from the ctypes mirror. - assert len(BaseRTModelConfig._fields_) == 106 + assert len(BaseRTModelConfig._fields_) == 131 def test_architecture_field_is_char_array(self): # architecture should be a fixed 32-byte char array @@ -567,7 +567,7 @@ def test_model_config_size(self): # Exact sizeof(BaseRTModelConfig) from include/baseRT/types.h; the # library cross-check happens at import via baseRT_model_config_sizeof. # Must match the Rust mirror test (bindings/rust/baseRT-sys). - assert ctypes.sizeof(BaseRTModelConfig) == 1704 + assert ctypes.sizeof(BaseRTModelConfig) == 1792 def test_sampling_config_size(self): size = ctypes.sizeof(BaseRTSamplingConfig) diff --git a/bindings/rust/baseRT-sys/Cargo.toml b/bindings/rust/baseRT-sys/Cargo.toml index f997fcf..350475d 100644 --- a/bindings/rust/baseRT-sys/Cargo.toml +++ b/bindings/rust/baseRT-sys/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "baseRT-sys" -version = "0.2.3" +version = "0.2.4" edition = "2021" description = "Raw FFI bindings for the BaseRT LLM inference engine" license = "Apache-2.0" diff --git a/bindings/rust/baseRT-sys/src/lib.rs b/bindings/rust/baseRT-sys/src/lib.rs index 7cccdbc..bb6c674 100644 --- a/bindings/rust/baseRT-sys/src/lib.rs +++ b/bindings/rust/baseRT-sys/src/lib.rs @@ -70,6 +70,13 @@ pub struct BaseRTModelConfig { pub gdn_value_head_dim: u32, pub gdn_conv_kernel: u32, + // Nemotron-H hybrid Mamba-2 SSM + pub ssm_state_size: u32, + pub ssm_conv_kernel: u32, + pub ssm_num_groups: u32, + pub ssm_inner_size: u32, + pub ssm_num_heads: u32, + // Mixture-of-Experts (0 = dense) pub n_experts: u32, pub n_experts_used: u32, @@ -78,6 +85,7 @@ pub struct BaseRTModelConfig { pub expert_gating: u8, pub norm_topk_prob: u8, pub _moe_pad: [u8; 2], + pub expert_weights_scale: c_float, // Vision tower (all zero = none) pub vision_n_layers: u32, @@ -142,6 +150,27 @@ pub struct BaseRTModelConfig { pub vision_pos_embed_w: u32, pub vision_adapter_dim: u32, pub video_token_id: u32, + // GLM 5.2 / glm-dsa + pub q_lora_rank: u32, + pub kv_lora_rank: u32, + pub qk_nope_head_dim: u32, + pub qk_rope_head_dim: u32, + pub v_head_dim: u32, + pub routed_scaling_factor: c_float, + pub first_k_dense_replace: u32, + pub nextn_predict_layers: u32, + pub indexer_head_count: u32, + pub indexer_key_length: u32, + pub indexer_top_k: u32, + // gpt-oss + pub rope_yarn_beta_fast: f32, + pub rope_yarn_beta_slow: f32, + pub swiglu_limit: f32, + pub swiglu_alpha: f32, + pub attention_sinks: u8, + pub attention_bias: u8, + pub rope_yarn_truncate: u8, + pub _gptoss_pad: u8, } /// Transcription result statistics. @@ -266,6 +295,36 @@ extern "C" { pub fn baseRT_get_config(model: baseRT_model_t) -> BaseRTModelConfig; pub fn baseRT_model_config_sizeof() -> usize; pub fn baseRT_model_memory(model: baseRT_model_t) -> usize; + /// Device memory budget in bytes (Metal: recommended working set; CUDA: + /// total device memory). 0 when no supported device is present. + pub fn baseRT_device_memory_budget() -> usize; + /// A `max_context` sized for this device and this bundle. Metadata-only, + /// so it is safe to call before loading. 0 = could not size. + pub fn baseRT_suggest_max_context( + model_path: *const c_char, + max_batch: c_int, + kv_bits: c_int, + paged_kv: c_int, + ) -> c_int; + /// As above for a set of models that will be resident at the same time: + /// weights and KV pools add up, and the trained-window cap is the + /// shortest among them. + pub fn baseRT_suggest_max_context_multi( + model_paths: *const *const c_char, + n_models: c_int, + max_batch: c_int, + kv_bits: c_int, + paged_kv: c_int, + ) -> c_int; + /// 1 = the window fits the co-resident set, 0 = it does not, -1 = unknown. + pub fn baseRT_context_window_fits( + model_paths: *const *const c_char, + n_models: c_int, + max_batch: c_int, + kv_bits: c_int, + paged_kv: c_int, + window: c_int, + ) -> c_int; pub fn baseRT_get_error() -> *const c_char; // === Tokenization === @@ -465,7 +524,7 @@ mod tests { "config struct unexpectedly small ({})", mem::size_of::() ); - assert_eq!(mem::size_of::(), 1704); + assert_eq!(mem::size_of::(), 1792); } #[test] @@ -549,21 +608,29 @@ mod tests { assert_eq!(&base.attn_output_gate as *const _ as usize - base_ptr, 1232); assert_eq!(&base.linear_attn_layers as *const _ as usize - base_ptr, 1244); assert_eq!(&base.gdn_num_k_heads as *const _ as usize - base_ptr, 1308); - assert_eq!(&base.n_experts as *const _ as usize - base_ptr, 1328); - assert_eq!(&base.vision_n_layers as *const _ as usize - base_ptr, 1348); - assert_eq!(&base.vision_arch as *const _ as usize - base_ptr, 1408); - assert_eq!(&base.audio_n_layers as *const _ as usize - base_ptr, 1424); - assert_eq!(&base.eoa_token_id as *const _ as usize - base_ptr, 1500); - assert_eq!(&base.mrope_section as *const _ as usize - base_ptr, 1504); - assert_eq!(&base.mrope_interleaved as *const _ as usize - base_ptr, 1516); - assert_eq!(&base.rope_scaling_factor as *const _ as usize - base_ptr, 1520); - assert_eq!(&base.rope_orig_max_pos as *const _ as usize - base_ptr, 1532); - assert_eq!(&base.rope_scaling_type as *const _ as usize - base_ptr, 1536); - assert_eq!(&base.qk_scale_factor as *const _ as usize - base_ptr, 1540); - assert_eq!(&base.nope_layers as *const _ as usize - base_ptr, 1552); - assert_eq!(&base.embed_norm_eps as *const _ as usize - base_ptr, 1616); - assert_eq!(&base.vision_window_layers as *const _ as usize - base_ptr, 1620); - assert_eq!(&base.video_token_id as *const _ as usize - base_ptr, 1700); + // Nemotron-H SSM block, inserted between gdn_* and the MoE block — + // 20 bytes that shifted every field below it. + assert_eq!(&base.ssm_state_size as *const _ as usize - base_ptr, 1328); + assert_eq!(&base.ssm_num_heads as *const _ as usize - base_ptr, 1344); + assert_eq!(&base.n_experts as *const _ as usize - base_ptr, 1348); + assert_eq!(&base.expert_weights_scale as *const _ as usize - base_ptr, 1368); + assert_eq!(&base.vision_n_layers as *const _ as usize - base_ptr, 1372); + assert_eq!(&base.vision_arch as *const _ as usize - base_ptr, 1432); + assert_eq!(&base.audio_n_layers as *const _ as usize - base_ptr, 1448); + assert_eq!(&base.eoa_token_id as *const _ as usize - base_ptr, 1524); + assert_eq!(&base.mrope_section as *const _ as usize - base_ptr, 1528); + assert_eq!(&base.mrope_interleaved as *const _ as usize - base_ptr, 1540); + assert_eq!(&base.rope_scaling_factor as *const _ as usize - base_ptr, 1544); + assert_eq!(&base.rope_orig_max_pos as *const _ as usize - base_ptr, 1556); + assert_eq!(&base.rope_scaling_type as *const _ as usize - base_ptr, 1560); + assert_eq!(&base.qk_scale_factor as *const _ as usize - base_ptr, 1564); + assert_eq!(&base.nope_layers as *const _ as usize - base_ptr, 1576); + assert_eq!(&base.embed_norm_eps as *const _ as usize - base_ptr, 1640); + assert_eq!(&base.vision_window_layers as *const _ as usize - base_ptr, 1644); + assert_eq!(&base.video_token_id as *const _ as usize - base_ptr, 1724); + assert_eq!(&base.q_lora_rank as *const _ as usize - base_ptr, 1728); + assert_eq!(&base.indexer_top_k as *const _ as usize - base_ptr, 1768); + assert_eq!(&base.rope_yarn_beta_fast as *const _ as usize - base_ptr, 1772); } #[test] diff --git a/bindings/rust/baseRT/Cargo.toml b/bindings/rust/baseRT/Cargo.toml index 70f89fb..d5f0b9f 100644 --- a/bindings/rust/baseRT/Cargo.toml +++ b/bindings/rust/baseRT/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "baseRT" -version = "0.2.3" +version = "0.2.4" edition = "2021" description = "Safe Rust bindings for the BaseRT LLM inference engine (Apple Silicon)" license = "Apache-2.0" diff --git a/bindings/swift/Sources/CBaseRT/include/baseRT.h b/bindings/swift/Sources/CBaseRT/include/baseRT.h index 0decfaa..1738b64 100644 --- a/bindings/swift/Sources/CBaseRT/include/baseRT.h +++ b/bindings/swift/Sources/CBaseRT/include/baseRT.h @@ -64,7 +64,7 @@ extern "C" { #define BASERT_VERSION_MAJOR 0 #define BASERT_VERSION_MINOR 2 -#define BASERT_VERSION_PATCH 3 +#define BASERT_VERSION_PATCH 4 /// Compile-time version, packed as `(MAJOR<<16) | (MINOR<<8) | PATCH`. /// Useful for `#if BASERT_VERSION >= 0x000200` feature checks. @@ -89,7 +89,11 @@ typedef void *baseRT_model_t; /// binary itself (single-file distributions ship the shared library with the /// kernels linked in, so NULL just works). Named generically so non-Metal /// backends (CUDA/ROCm, future) can reuse the same parameter. -/// max_context: maximum context window (0 = use model default, capped at 4096). +/// max_context: maximum context window. 0 = the model's trained window, capped +/// at the chip's max prefill chunk — a shape default that ignores how much +/// memory this device actually has. A serving front end should instead pass +/// `baseRT_suggest_max_context()`, which derives the window from the device +/// budget, or a number the operator chose. /// Returns NULL on failure. baseRT_model_t baseRT_load_model(const char *model_path, const char *kernel_library_path, int max_context); @@ -123,6 +127,18 @@ enum { /// operator-facing "continuous batching disabled because ..." message), call /// the gated entry point once and read baseRT_get_error(). uint32_t baseRT_capabilities(baseRT_model_t model); +/// Hot-swap the kernel library on an already-loaded model, WITHOUT reloading +/// weights. Weights are mmap-backed Metal buffers independent of the kernel +/// library, so only the compiled pipelines + baked decode dispatch tables are +/// rebuilt against the new metallib (sub-second). Built for the automated +/// kernel-tuning loop: edit a `.metal`, `make shaders`, reload — no ~400 GB +/// model reload per iteration. `metallib_path` NULL uses the default sidecar +/// search (`build/baseRT.metallib`). Picks up `.metal`-only changes; a +/// C++/param/dispatch change still needs a rebuilt binary + restart. +/// Returns BASERT_OK, or an error code (leaves the model on the OLD library on +/// load failure). Not supported for whisper models. NOT thread-safe against +/// in-flight inference — serialize the caller. +int baseRT_reload_metallib(baseRT_model_t model, const char *metallib_path); /// Override the KV cache element width for the next baseRT_load_model call. /// bits = 0 → auto (per-model default; Q8_0 when head_dim%32==0) @@ -228,6 +244,97 @@ size_t baseRT_model_config_sizeof(void); /// Get total GPU memory used by model (bytes). size_t baseRT_model_memory(baseRT_model_t model); +/// The device memory budget the engine allocates within, in bytes, without +/// needing a loaded model. NOT installed RAM: on Metal this is the unified- +/// memory working set the OS recommends (~75% of physical RAM), on CUDA the +/// device's total memory. 0 when no supported device is present. +/// +/// This is the number that matters for sizing decisions — it is what the load +/// path compares the model's working set against before warning that the OS +/// will start paging weights. +size_t baseRT_device_memory_budget(void); + +/// Floor on a window derived by `baseRT_suggest_max_context`. Policy, not a +/// hardware limit: an explicit max_context is honoured below it. There is no +/// corresponding ceiling — a derived window is bounded by the device's memory +/// and by the model's trained window, and by nothing else. +enum { + BASERT_AUTO_CONTEXT_MIN = 4096, +}; + +/// Suggest a `max_context` for `model_path` that fits this device: the memory +/// budget above, less the weights and activation/scratch headroom, divided by +/// what one token of KV costs across `max_batch` concurrent decode lanes. The +/// result is a multiple of 1024, at least BASERT_AUTO_CONTEXT_MIN, and never +/// above the model's trained window — a device with the memory for a model's +/// whole window gets the whole window. +/// +/// Reads the bundle's metadata only — no tensor upload, no allocation — so it +/// is cheap enough to call before `baseRT_load_model`. A serving front end +/// calls it once at startup instead of shipping a fixed default that is too +/// small on a workstation and too large on a laptop. +/// +/// max_batch: concurrent sequences the KV pool must hold (0 or 1 = single). +/// Only counted for models the load will actually page, since a +/// contiguous cache holds one history however wide the batch is. +/// kv_bits: the value that will be passed to `baseRT_set_kv_bits` (0 = auto, +/// or 4 / 8 / 16 / 84), since KV precision changes the answer. +/// paged_kv: non-zero if the load will enable paged KV (`baseRT_set_paged_kv` +/// or BaseRTLoadOptions::paged_kv). It changes both the per-layer +/// shape and whether lanes multiply, and the two caches differ by +/// several times on the hybrid decoders. +/// Returns 0 if the bundle cannot be read, the device budget is unknown, or the +/// bundle declares no trained window; the caller keeps its own default then. +int baseRT_suggest_max_context(const char *model_path, int max_batch, int kv_bits, int paged_kv); + +/// As above for a set of models that will be resident AT THE SAME TIME, such +/// as a server's eagerly loaded set. They share one budget, so their weights +/// add up, their KV pools add up, and the window returned is the one that fits +/// all of them at once, capped at the SHORTEST trained window among them. +/// +/// This is not the same as calling the single-model form on each and taking +/// the minimum: that answers "what fits this model alone", and two models that +/// each fit alone can exceed the budget together. +int baseRT_suggest_max_context_multi(const char *const *model_paths, int n_models, int max_batch, int kv_bits, + int paged_kv); + +/// Does `window` fit `model_paths` co-resident on this device, under the same +/// budget the suggestion above derives from? 1 = yes, 0 = no, -1 = cannot tell +/// (a bundle would not open, no trained window, unknown budget). +/// +/// Distinct from comparing against `baseRT_suggest_max_context_multi`: that +/// clamps its answer into the policy floor, so a set that fits NO tokens at all +/// still reports the floor and would compare equal to it. Ask this instead +/// before admitting a model into a window that was chosen without it. +int baseRT_context_window_fits(const char *const *model_paths, int n_models, int max_batch, int kv_bits, int paged_kv, + int window); +/// KV-cache allocation strategy reported by BaseRTMemoryStats. +typedef enum BaseRTKVCacheLayout { + BASERT_KV_CACHE_NONE = 0, + BASERT_KV_CACHE_CONTIGUOUS = 1, + BASERT_KV_CACHE_PAGED = 2, +} BaseRTKVCacheLayout; + +/// Runtime-owned memory counters for an idle model handle. Capacity is the +/// memory reserved for the KV cache; used bytes are the logical occupied +/// portion (contiguous cache) or occupied physical blocks (paged cache). +/// Neither value is process RSS, device-global usage, nor a peak. +typedef struct BaseRTMemoryStats { + uint64_t runtime_allocated_bytes; + uint64_t kv_cache_capacity_bytes; + uint64_t kv_cache_used_bytes; + uint64_t kv_cache_blocks_total; + uint64_t kv_cache_blocks_used; + uint64_t kv_cache_tokens_used; + BaseRTKVCacheLayout kv_cache_layout; + uint32_t reserved; +} BaseRTMemoryStats; + +/// Read current runtime and KV-cache memory counters. This is a boundary +/// observation: call only while no inference operation is mutating `model`. +/// Returns false for a null model or output pointer. +bool baseRT_model_memory_stats(baseRT_model_t model, BaseRTMemoryStats *out_stats); + /// Get last error message (thread-local). The string is valid until /// the next API call from the same thread that fails or that explicitly /// resets the error state. Returns "" when there is no pending error. @@ -260,6 +367,25 @@ const char *baseRT_decode_token(baseRT_model_t model, uint32_t token_id); /// string lives in a thread-local buffer that is overwritten on each call. const char *baseRT_decode_token_static(baseRT_model_t model, uint32_t token_id); +/// A caller-owned incremental-decode stream: the same UTF-8 assembly and +/// channel-protocol normalization `baseRT_decode_token` applies (gpt-oss +/// Harmony and Muse framing become reasoning spans / ChatML tool-call +/// blocks), but with state private to this stream. One per continuously +/// batched lane; the model's own stream is untouched. +typedef struct baseRT_decode_stream *baseRT_decode_stream_t; +baseRT_decode_stream_t baseRT_decode_stream_create(baseRT_model_t model); +void baseRT_decode_stream_reset(baseRT_decode_stream_t stream); +/// Returns a pointer into the stream's own buffer, valid until the next +/// call on the same stream. +const char *baseRT_decode_stream_token(baseRT_decode_stream_t stream, uint32_t token_id); +void baseRT_decode_stream_free(baseRT_decode_stream_t stream); + +/// 1 when `token_id` ends the assistant turn: the bundle's eos_token_id OR +/// any of its extra end-of-turn ids (Llama 3's <|eot_id|>, gpt-oss's +/// <|call|> and <|endoftext|>, …). `baseRT_eos_token_id` reports only the +/// first; a lane that compares against that alone runs past the others. +int baseRT_is_eos_token(baseRT_model_t model, uint32_t token_id); + /// Length-preserving variant of `baseRT_decode_token_static` for callers /// that need the token's EXACT raw bytes. Byte-level BPE / byte-fallback /// tokens can decode to bytes containing 0x00, which the C-string variants @@ -810,6 +936,18 @@ uint32_t baseRT_prefill(baseRT_model_t model, const uint32_t *tokens, int n_toke /// on error). int baseRT_read_logits(baseRT_model_t model, float *out, int max_logits); +/// Sliding-window teacher-forced perplexity of `tokens[0..n_tokens)` on this +/// model. For each start position (advanced by `stride`, up to `max_positions`; +/// <=0 means unbounded), prefills the preceding `window` tokens with a fresh KV +/// cache and accumulates -log P(true next token); PPL = exp(mean NLL). Leaves +/// the model's KV cache reset. Writes exp(mean NLL) to `*out_ppl` and the number +/// of scored positions to `*out_positions` (either may be NULL). This is the +/// exact loop the `baseRT_ppl` tool runs, exposed so a resident server can gate +/// accuracy without a second model load. Returns BASERT_OK or an error code. +/// NOT thread-safe against concurrent inference on the same model. +int baseRT_perplexity(baseRT_model_t model, const uint32_t *tokens, int n_tokens, int window, int stride, + int max_positions, double *out_ppl, int *out_positions); + /// Multimodal prefill: run vision tower on image, then prefill tokens with /// image features spliced at positions where tokens[i] == config.image_token_id. /// The number of image placeholder tokens in the stream must equal the image's @@ -876,6 +1014,13 @@ int baseRT_load_state(baseRT_model_t model, const char *path); /// every forward pass that runs a GEMM with a tensor_name registered in /// the adapter has a post-GEMM low-rank delta applied (`y += B @ A @ x`). /// +/// While an adapter is active, single-sequence decode skips the baked +/// dispatch-table replay (the table carries no delta dispatches) and takes +/// the per-token immediate encode instead — correct output at reduced +/// decode throughput. Speculative decode is likewise disabled for the +/// duration. The batched multi-sequence API (`baseRT_batch_step*`) does +/// NOT apply adapters. +/// /// Calling `baseRT_lora_load` again replaces the active adapter (no /// stacking). Returns 0 on success, negative on failure (see /// `baseRT_get_error`). @@ -1013,6 +1158,13 @@ const char *baseRT_bos_token(baseRT_model_t model); /// BOS token id, for callers that need to prepend BOS to raw token /// sequences (e.g. perplexity windows on BOS-sensitive models). uint32_t baseRT_bos_id(baseRT_model_t model); + +/// 1 when the tokenizer prepends BOS to encoded text (the model was trained +/// with a leading BOS), 0 when it does not (byte-level BPE families such as +/// GPT-2 / o200k). baseRT_bos_id can still name a token in the 0 case; +/// callers that synthesize a sequence start (perplexity windows) must key +/// on this, not on bos_id being valid. +int baseRT_add_bos(baseRT_model_t model); const char *baseRT_eos_token(baseRT_model_t model); /// Primary end-of-sequence token id (the one the continuous-batching engine and diff --git a/bindings/swift/Sources/CBaseRT/include/types.h b/bindings/swift/Sources/CBaseRT/include/types.h index e607967..6be7daf 100644 --- a/bindings/swift/Sources/CBaseRT/include/types.h +++ b/bindings/swift/Sources/CBaseRT/include/types.h @@ -27,6 +27,22 @@ typedef enum { } BaseRTErrorCode; /// Model configuration extracted from weight file metadata. +/// +/// ABI NOTE (0.2.4): this struct's layout CHANGED. New family support landed +/// its fields beside the ones they belong with (`ssm_*` next to `gdn_*`, MLA +/// and gpt-oss fields beside the attention block) rather than appended, so +/// every offset after the first insertion moved, and the struct grew. That is +/// a deliberate readability choice under the pre-1.0 clause in baseRT.h ("the +/// above is intent, not contract"), not an oversight — but it means a caller +/// built against 0.2.3 MUST be recompiled against this header. Note that +/// appending would not have been sufficient on its own either: +/// `baseRT_get_config` returns this struct BY VALUE, so growth alone +/// overruns an older caller's return slot whatever the field order. +/// +/// Mismatches are detectable rather than silent: `baseRT_model_config_sizeof()` +/// reports the library's own `sizeof`, and the language bindings compare it +/// against their mirrored definition at load time. Check it if you bind to +/// this struct from outside the tree. typedef struct { // Decoder (or decoder-only LLM) parameters uint32_t dim; // embedding dimension @@ -106,6 +122,18 @@ typedef struct { uint32_t gdn_value_head_dim; // linear_value_head_dim uint32_t gdn_conv_kernel; // linear_conv_kernel_dim (short causal conv width) + // ── Nemotron-H hybrid Mamba-2 SSM (0 = not a Mamba-2 hybrid) ───── + // Per-layer schedule: linear_attn_layers bit set = Mamba-2 layer, + // n_kv_heads_per_layer[il] > 0 = attention layer, otherwise (MoE) FFN. + // The gdn_* fields above are additionally aliased at load so the shared + // GDNStateCache sizes the SSM state pool: nk=ssm_num_groups, + // khd=ssm_state_size, nv=ssm_num_heads, vhd=ssm_inner/ssm_heads. + uint32_t ssm_state_size; // per-head state width N (Nemotron 3 Nano: 128) + uint32_t ssm_conv_kernel; // depthwise causal conv taps (4) + uint32_t ssm_num_groups; // B/C groups (8) + uint32_t ssm_inner_size; // heads * head_dim (4096) + uint32_t ssm_num_heads; // SSM heads (64) + // Mixture-of-Experts (0 = dense model). // Gemma 4 26B-A4B: n_experts=128, n_experts_used=8, n_experts_shared=1 (via dense ffn.*), expert_gating=0 // (softmax), norm_topk_prob=0 Qwen3.6-35B-A3B (qwen35moe): n_experts=128, n_experts_used=8, n_experts_shared=0, @@ -117,6 +145,9 @@ typedef struct { uint8_t expert_gating; // 0 = softmax, 1 = sigmoid uint8_t norm_topk_prob; // 1 = renormalize top-k weights to sum to 1 (Qwen), 0 = leave as-is (Gemma) uint8_t _moe_pad[2]; // align to 4 bytes + // Routed-expert output scale (DeepSeek/Nemotron routed_scaling_factor; + // applied to the renormalized top-k weights). 0 = disabled (treat as 1). + float expert_weights_scale; // Vision tower (Gemma 4, PaliGemma, Llava-style multimodal). // All zero = no vision tower. @@ -234,6 +265,39 @@ typedef struct { uint32_t vision_pos_embed_w; // learned position grid width (e.g. 32) uint32_t vision_adapter_dim; // projector hidden width (`projector_hidden_size`, e.g. 4096) uint32_t video_token_id; // text-side placeholder token for video features (0 = none) + + // ── GLM 5.2 / glm-dsa: Multi-head Latent Attention + MoE ───────── + // DeepSeek-V3.2-style decoder. All zero = not an MLA model. The + // struct's uniform head_dim can't express MLA's compressed/asymmetric + // Q/K/V dims, so the MLA encoder (src/core/models/glm_dsa.cpp) reads + // these explicit fields instead. See arch_descriptor.h is_mla flag. + uint32_t q_lora_rank; // MLA query compression rank (attn_q_a output); 0 = not MLA + uint32_t kv_lora_rank; // MLA KV latent rank (attn_kv_a_mqa kv part); K cache row = kv_lora_rank + qk_rope_head_dim + uint32_t qk_nope_head_dim; // per-head non-positional Q/K dim (attends in latent space via k_b) + uint32_t qk_rope_head_dim; // per-head decoupled-RoPE Q/K dim (the only rotated part; NORM/GPT-J rope) + uint32_t v_head_dim; // per-head value dim after v_b up-projection + float routed_scaling_factor; // routed-expert output scale (DeepSeek routed_scaling_factor); 0 = none + uint32_t first_k_dense_replace; // first N layers use a dense SwiGLU FFN instead of MoE (GLM 5.2 = 3) + uint32_t nextn_predict_layers; // Multi-Token-Prediction head layers (dropped at convert; recorded only) + // DSA lightning-indexer geometry. Loaded but UNUSED by the dense MLA + // path (llama.cpp's glm-dsa never runs the indexer at any context + // length). Kept for a future true-sparse-attention path. 0 = none. + uint32_t indexer_head_count; + uint32_t indexer_key_length; + uint32_t indexer_top_k; + // ── gpt-oss (0 / empty = not applicable) ───────────────────────── + // YaRN correction-range betas (HF `rope_scaling.beta_fast` / `beta_slow`); + // used with rope_scaling_type == 4 (yarn). 0 = the YaRN defaults (32 / 1). + float rope_yarn_beta_fast; + float rope_yarn_beta_slow; + // Clamped SwiGLU: gate = min(gate, limit), up = clamp(up, -limit, limit), + // act = (up + 1) * gate * sigmoid(alpha * gate). limit 0 = plain SwiGLU. + float swiglu_limit; + float swiglu_alpha; + uint8_t attention_sinks; // 1 = learned per-head sink logits (`layers.N.attention.sinks`) + uint8_t attention_bias; // 1 = q/k/v/o projections carry biases + uint8_t rope_yarn_truncate; // 1 = floor/ceil the YaRN correction range (HF default); gpt-oss: 0 + uint8_t _gptoss_pad; // align to 4 bytes } BaseRTModelConfig; /// Transcription result statistics. diff --git a/include/baseRT/baseRT.h b/include/baseRT/baseRT.h index 0decfaa..1738b64 100644 --- a/include/baseRT/baseRT.h +++ b/include/baseRT/baseRT.h @@ -64,7 +64,7 @@ extern "C" { #define BASERT_VERSION_MAJOR 0 #define BASERT_VERSION_MINOR 2 -#define BASERT_VERSION_PATCH 3 +#define BASERT_VERSION_PATCH 4 /// Compile-time version, packed as `(MAJOR<<16) | (MINOR<<8) | PATCH`. /// Useful for `#if BASERT_VERSION >= 0x000200` feature checks. @@ -89,7 +89,11 @@ typedef void *baseRT_model_t; /// binary itself (single-file distributions ship the shared library with the /// kernels linked in, so NULL just works). Named generically so non-Metal /// backends (CUDA/ROCm, future) can reuse the same parameter. -/// max_context: maximum context window (0 = use model default, capped at 4096). +/// max_context: maximum context window. 0 = the model's trained window, capped +/// at the chip's max prefill chunk — a shape default that ignores how much +/// memory this device actually has. A serving front end should instead pass +/// `baseRT_suggest_max_context()`, which derives the window from the device +/// budget, or a number the operator chose. /// Returns NULL on failure. baseRT_model_t baseRT_load_model(const char *model_path, const char *kernel_library_path, int max_context); @@ -123,6 +127,18 @@ enum { /// operator-facing "continuous batching disabled because ..." message), call /// the gated entry point once and read baseRT_get_error(). uint32_t baseRT_capabilities(baseRT_model_t model); +/// Hot-swap the kernel library on an already-loaded model, WITHOUT reloading +/// weights. Weights are mmap-backed Metal buffers independent of the kernel +/// library, so only the compiled pipelines + baked decode dispatch tables are +/// rebuilt against the new metallib (sub-second). Built for the automated +/// kernel-tuning loop: edit a `.metal`, `make shaders`, reload — no ~400 GB +/// model reload per iteration. `metallib_path` NULL uses the default sidecar +/// search (`build/baseRT.metallib`). Picks up `.metal`-only changes; a +/// C++/param/dispatch change still needs a rebuilt binary + restart. +/// Returns BASERT_OK, or an error code (leaves the model on the OLD library on +/// load failure). Not supported for whisper models. NOT thread-safe against +/// in-flight inference — serialize the caller. +int baseRT_reload_metallib(baseRT_model_t model, const char *metallib_path); /// Override the KV cache element width for the next baseRT_load_model call. /// bits = 0 → auto (per-model default; Q8_0 when head_dim%32==0) @@ -228,6 +244,97 @@ size_t baseRT_model_config_sizeof(void); /// Get total GPU memory used by model (bytes). size_t baseRT_model_memory(baseRT_model_t model); +/// The device memory budget the engine allocates within, in bytes, without +/// needing a loaded model. NOT installed RAM: on Metal this is the unified- +/// memory working set the OS recommends (~75% of physical RAM), on CUDA the +/// device's total memory. 0 when no supported device is present. +/// +/// This is the number that matters for sizing decisions — it is what the load +/// path compares the model's working set against before warning that the OS +/// will start paging weights. +size_t baseRT_device_memory_budget(void); + +/// Floor on a window derived by `baseRT_suggest_max_context`. Policy, not a +/// hardware limit: an explicit max_context is honoured below it. There is no +/// corresponding ceiling — a derived window is bounded by the device's memory +/// and by the model's trained window, and by nothing else. +enum { + BASERT_AUTO_CONTEXT_MIN = 4096, +}; + +/// Suggest a `max_context` for `model_path` that fits this device: the memory +/// budget above, less the weights and activation/scratch headroom, divided by +/// what one token of KV costs across `max_batch` concurrent decode lanes. The +/// result is a multiple of 1024, at least BASERT_AUTO_CONTEXT_MIN, and never +/// above the model's trained window — a device with the memory for a model's +/// whole window gets the whole window. +/// +/// Reads the bundle's metadata only — no tensor upload, no allocation — so it +/// is cheap enough to call before `baseRT_load_model`. A serving front end +/// calls it once at startup instead of shipping a fixed default that is too +/// small on a workstation and too large on a laptop. +/// +/// max_batch: concurrent sequences the KV pool must hold (0 or 1 = single). +/// Only counted for models the load will actually page, since a +/// contiguous cache holds one history however wide the batch is. +/// kv_bits: the value that will be passed to `baseRT_set_kv_bits` (0 = auto, +/// or 4 / 8 / 16 / 84), since KV precision changes the answer. +/// paged_kv: non-zero if the load will enable paged KV (`baseRT_set_paged_kv` +/// or BaseRTLoadOptions::paged_kv). It changes both the per-layer +/// shape and whether lanes multiply, and the two caches differ by +/// several times on the hybrid decoders. +/// Returns 0 if the bundle cannot be read, the device budget is unknown, or the +/// bundle declares no trained window; the caller keeps its own default then. +int baseRT_suggest_max_context(const char *model_path, int max_batch, int kv_bits, int paged_kv); + +/// As above for a set of models that will be resident AT THE SAME TIME, such +/// as a server's eagerly loaded set. They share one budget, so their weights +/// add up, their KV pools add up, and the window returned is the one that fits +/// all of them at once, capped at the SHORTEST trained window among them. +/// +/// This is not the same as calling the single-model form on each and taking +/// the minimum: that answers "what fits this model alone", and two models that +/// each fit alone can exceed the budget together. +int baseRT_suggest_max_context_multi(const char *const *model_paths, int n_models, int max_batch, int kv_bits, + int paged_kv); + +/// Does `window` fit `model_paths` co-resident on this device, under the same +/// budget the suggestion above derives from? 1 = yes, 0 = no, -1 = cannot tell +/// (a bundle would not open, no trained window, unknown budget). +/// +/// Distinct from comparing against `baseRT_suggest_max_context_multi`: that +/// clamps its answer into the policy floor, so a set that fits NO tokens at all +/// still reports the floor and would compare equal to it. Ask this instead +/// before admitting a model into a window that was chosen without it. +int baseRT_context_window_fits(const char *const *model_paths, int n_models, int max_batch, int kv_bits, int paged_kv, + int window); +/// KV-cache allocation strategy reported by BaseRTMemoryStats. +typedef enum BaseRTKVCacheLayout { + BASERT_KV_CACHE_NONE = 0, + BASERT_KV_CACHE_CONTIGUOUS = 1, + BASERT_KV_CACHE_PAGED = 2, +} BaseRTKVCacheLayout; + +/// Runtime-owned memory counters for an idle model handle. Capacity is the +/// memory reserved for the KV cache; used bytes are the logical occupied +/// portion (contiguous cache) or occupied physical blocks (paged cache). +/// Neither value is process RSS, device-global usage, nor a peak. +typedef struct BaseRTMemoryStats { + uint64_t runtime_allocated_bytes; + uint64_t kv_cache_capacity_bytes; + uint64_t kv_cache_used_bytes; + uint64_t kv_cache_blocks_total; + uint64_t kv_cache_blocks_used; + uint64_t kv_cache_tokens_used; + BaseRTKVCacheLayout kv_cache_layout; + uint32_t reserved; +} BaseRTMemoryStats; + +/// Read current runtime and KV-cache memory counters. This is a boundary +/// observation: call only while no inference operation is mutating `model`. +/// Returns false for a null model or output pointer. +bool baseRT_model_memory_stats(baseRT_model_t model, BaseRTMemoryStats *out_stats); + /// Get last error message (thread-local). The string is valid until /// the next API call from the same thread that fails or that explicitly /// resets the error state. Returns "" when there is no pending error. @@ -260,6 +367,25 @@ const char *baseRT_decode_token(baseRT_model_t model, uint32_t token_id); /// string lives in a thread-local buffer that is overwritten on each call. const char *baseRT_decode_token_static(baseRT_model_t model, uint32_t token_id); +/// A caller-owned incremental-decode stream: the same UTF-8 assembly and +/// channel-protocol normalization `baseRT_decode_token` applies (gpt-oss +/// Harmony and Muse framing become reasoning spans / ChatML tool-call +/// blocks), but with state private to this stream. One per continuously +/// batched lane; the model's own stream is untouched. +typedef struct baseRT_decode_stream *baseRT_decode_stream_t; +baseRT_decode_stream_t baseRT_decode_stream_create(baseRT_model_t model); +void baseRT_decode_stream_reset(baseRT_decode_stream_t stream); +/// Returns a pointer into the stream's own buffer, valid until the next +/// call on the same stream. +const char *baseRT_decode_stream_token(baseRT_decode_stream_t stream, uint32_t token_id); +void baseRT_decode_stream_free(baseRT_decode_stream_t stream); + +/// 1 when `token_id` ends the assistant turn: the bundle's eos_token_id OR +/// any of its extra end-of-turn ids (Llama 3's <|eot_id|>, gpt-oss's +/// <|call|> and <|endoftext|>, …). `baseRT_eos_token_id` reports only the +/// first; a lane that compares against that alone runs past the others. +int baseRT_is_eos_token(baseRT_model_t model, uint32_t token_id); + /// Length-preserving variant of `baseRT_decode_token_static` for callers /// that need the token's EXACT raw bytes. Byte-level BPE / byte-fallback /// tokens can decode to bytes containing 0x00, which the C-string variants @@ -810,6 +936,18 @@ uint32_t baseRT_prefill(baseRT_model_t model, const uint32_t *tokens, int n_toke /// on error). int baseRT_read_logits(baseRT_model_t model, float *out, int max_logits); +/// Sliding-window teacher-forced perplexity of `tokens[0..n_tokens)` on this +/// model. For each start position (advanced by `stride`, up to `max_positions`; +/// <=0 means unbounded), prefills the preceding `window` tokens with a fresh KV +/// cache and accumulates -log P(true next token); PPL = exp(mean NLL). Leaves +/// the model's KV cache reset. Writes exp(mean NLL) to `*out_ppl` and the number +/// of scored positions to `*out_positions` (either may be NULL). This is the +/// exact loop the `baseRT_ppl` tool runs, exposed so a resident server can gate +/// accuracy without a second model load. Returns BASERT_OK or an error code. +/// NOT thread-safe against concurrent inference on the same model. +int baseRT_perplexity(baseRT_model_t model, const uint32_t *tokens, int n_tokens, int window, int stride, + int max_positions, double *out_ppl, int *out_positions); + /// Multimodal prefill: run vision tower on image, then prefill tokens with /// image features spliced at positions where tokens[i] == config.image_token_id. /// The number of image placeholder tokens in the stream must equal the image's @@ -876,6 +1014,13 @@ int baseRT_load_state(baseRT_model_t model, const char *path); /// every forward pass that runs a GEMM with a tensor_name registered in /// the adapter has a post-GEMM low-rank delta applied (`y += B @ A @ x`). /// +/// While an adapter is active, single-sequence decode skips the baked +/// dispatch-table replay (the table carries no delta dispatches) and takes +/// the per-token immediate encode instead — correct output at reduced +/// decode throughput. Speculative decode is likewise disabled for the +/// duration. The batched multi-sequence API (`baseRT_batch_step*`) does +/// NOT apply adapters. +/// /// Calling `baseRT_lora_load` again replaces the active adapter (no /// stacking). Returns 0 on success, negative on failure (see /// `baseRT_get_error`). @@ -1013,6 +1158,13 @@ const char *baseRT_bos_token(baseRT_model_t model); /// BOS token id, for callers that need to prepend BOS to raw token /// sequences (e.g. perplexity windows on BOS-sensitive models). uint32_t baseRT_bos_id(baseRT_model_t model); + +/// 1 when the tokenizer prepends BOS to encoded text (the model was trained +/// with a leading BOS), 0 when it does not (byte-level BPE families such as +/// GPT-2 / o200k). baseRT_bos_id can still name a token in the 0 case; +/// callers that synthesize a sequence start (perplexity windows) must key +/// on this, not on bos_id being valid. +int baseRT_add_bos(baseRT_model_t model); const char *baseRT_eos_token(baseRT_model_t model); /// Primary end-of-sequence token id (the one the continuous-batching engine and diff --git a/include/baseRT/types.h b/include/baseRT/types.h index e607967..6be7daf 100644 --- a/include/baseRT/types.h +++ b/include/baseRT/types.h @@ -27,6 +27,22 @@ typedef enum { } BaseRTErrorCode; /// Model configuration extracted from weight file metadata. +/// +/// ABI NOTE (0.2.4): this struct's layout CHANGED. New family support landed +/// its fields beside the ones they belong with (`ssm_*` next to `gdn_*`, MLA +/// and gpt-oss fields beside the attention block) rather than appended, so +/// every offset after the first insertion moved, and the struct grew. That is +/// a deliberate readability choice under the pre-1.0 clause in baseRT.h ("the +/// above is intent, not contract"), not an oversight — but it means a caller +/// built against 0.2.3 MUST be recompiled against this header. Note that +/// appending would not have been sufficient on its own either: +/// `baseRT_get_config` returns this struct BY VALUE, so growth alone +/// overruns an older caller's return slot whatever the field order. +/// +/// Mismatches are detectable rather than silent: `baseRT_model_config_sizeof()` +/// reports the library's own `sizeof`, and the language bindings compare it +/// against their mirrored definition at load time. Check it if you bind to +/// this struct from outside the tree. typedef struct { // Decoder (or decoder-only LLM) parameters uint32_t dim; // embedding dimension @@ -106,6 +122,18 @@ typedef struct { uint32_t gdn_value_head_dim; // linear_value_head_dim uint32_t gdn_conv_kernel; // linear_conv_kernel_dim (short causal conv width) + // ── Nemotron-H hybrid Mamba-2 SSM (0 = not a Mamba-2 hybrid) ───── + // Per-layer schedule: linear_attn_layers bit set = Mamba-2 layer, + // n_kv_heads_per_layer[il] > 0 = attention layer, otherwise (MoE) FFN. + // The gdn_* fields above are additionally aliased at load so the shared + // GDNStateCache sizes the SSM state pool: nk=ssm_num_groups, + // khd=ssm_state_size, nv=ssm_num_heads, vhd=ssm_inner/ssm_heads. + uint32_t ssm_state_size; // per-head state width N (Nemotron 3 Nano: 128) + uint32_t ssm_conv_kernel; // depthwise causal conv taps (4) + uint32_t ssm_num_groups; // B/C groups (8) + uint32_t ssm_inner_size; // heads * head_dim (4096) + uint32_t ssm_num_heads; // SSM heads (64) + // Mixture-of-Experts (0 = dense model). // Gemma 4 26B-A4B: n_experts=128, n_experts_used=8, n_experts_shared=1 (via dense ffn.*), expert_gating=0 // (softmax), norm_topk_prob=0 Qwen3.6-35B-A3B (qwen35moe): n_experts=128, n_experts_used=8, n_experts_shared=0, @@ -117,6 +145,9 @@ typedef struct { uint8_t expert_gating; // 0 = softmax, 1 = sigmoid uint8_t norm_topk_prob; // 1 = renormalize top-k weights to sum to 1 (Qwen), 0 = leave as-is (Gemma) uint8_t _moe_pad[2]; // align to 4 bytes + // Routed-expert output scale (DeepSeek/Nemotron routed_scaling_factor; + // applied to the renormalized top-k weights). 0 = disabled (treat as 1). + float expert_weights_scale; // Vision tower (Gemma 4, PaliGemma, Llava-style multimodal). // All zero = no vision tower. @@ -234,6 +265,39 @@ typedef struct { uint32_t vision_pos_embed_w; // learned position grid width (e.g. 32) uint32_t vision_adapter_dim; // projector hidden width (`projector_hidden_size`, e.g. 4096) uint32_t video_token_id; // text-side placeholder token for video features (0 = none) + + // ── GLM 5.2 / glm-dsa: Multi-head Latent Attention + MoE ───────── + // DeepSeek-V3.2-style decoder. All zero = not an MLA model. The + // struct's uniform head_dim can't express MLA's compressed/asymmetric + // Q/K/V dims, so the MLA encoder (src/core/models/glm_dsa.cpp) reads + // these explicit fields instead. See arch_descriptor.h is_mla flag. + uint32_t q_lora_rank; // MLA query compression rank (attn_q_a output); 0 = not MLA + uint32_t kv_lora_rank; // MLA KV latent rank (attn_kv_a_mqa kv part); K cache row = kv_lora_rank + qk_rope_head_dim + uint32_t qk_nope_head_dim; // per-head non-positional Q/K dim (attends in latent space via k_b) + uint32_t qk_rope_head_dim; // per-head decoupled-RoPE Q/K dim (the only rotated part; NORM/GPT-J rope) + uint32_t v_head_dim; // per-head value dim after v_b up-projection + float routed_scaling_factor; // routed-expert output scale (DeepSeek routed_scaling_factor); 0 = none + uint32_t first_k_dense_replace; // first N layers use a dense SwiGLU FFN instead of MoE (GLM 5.2 = 3) + uint32_t nextn_predict_layers; // Multi-Token-Prediction head layers (dropped at convert; recorded only) + // DSA lightning-indexer geometry. Loaded but UNUSED by the dense MLA + // path (llama.cpp's glm-dsa never runs the indexer at any context + // length). Kept for a future true-sparse-attention path. 0 = none. + uint32_t indexer_head_count; + uint32_t indexer_key_length; + uint32_t indexer_top_k; + // ── gpt-oss (0 / empty = not applicable) ───────────────────────── + // YaRN correction-range betas (HF `rope_scaling.beta_fast` / `beta_slow`); + // used with rope_scaling_type == 4 (yarn). 0 = the YaRN defaults (32 / 1). + float rope_yarn_beta_fast; + float rope_yarn_beta_slow; + // Clamped SwiGLU: gate = min(gate, limit), up = clamp(up, -limit, limit), + // act = (up + 1) * gate * sigmoid(alpha * gate). limit 0 = plain SwiGLU. + float swiglu_limit; + float swiglu_alpha; + uint8_t attention_sinks; // 1 = learned per-head sink logits (`layers.N.attention.sinks`) + uint8_t attention_bias; // 1 = q/k/v/o projections carry biases + uint8_t rope_yarn_truncate; // 1 = floor/ceil the YaRN correction range (HF default); gpt-oss: 0 + uint8_t _gptoss_pad; // align to 4 bytes } BaseRTModelConfig; /// Transcription result statistics.