diff --git a/crates/hf/src/modelopt.rs b/crates/hf/src/modelopt.rs index 69cf0cbb5..57a67a329 100644 --- a/crates/hf/src/modelopt.rs +++ b/crates/hf/src/modelopt.rs @@ -991,8 +991,11 @@ mod tests { ) .unwrap(); let raws: Vec<_> = st.names().map(|n| (n, st.raw(n).unwrap())).collect(); - let plan = - crate::quant_config::plan(&config, raws.iter().map(|(n, r)| (*n, r.dtype, r.shape))); + let plan = crate::quant_config::plan( + &config, + raws.iter().map(|(n, r)| (*n, r.dtype, r.shape)), + |_| None, + ); let experts = ["e.0.weight", "e.1.weight"].map(|n| (n, &plan.placed[n])); let got = normalize(|n| st.raw(n), &experts, Use::Matrix(Plain::Q8_0)).unwrap(); assert_eq!(got.scale, [2.0, 3.0]); diff --git a/crates/hf/src/quant_config.rs b/crates/hf/src/quant_config.rs index b451d2058..ecc4ad7f5 100644 --- a/crates/hf/src/quant_config.rs +++ b/crates/hf/src/quant_config.rs @@ -15,6 +15,7 @@ //! `quant_method`, algorithm or scheme is refused when the config is read. `input_scale` (an //! activation scale) is reported, not used. +use std::cell::LazyCell; use std::collections::{BTreeMap, BTreeSet}; use std::fmt; use std::path::Path; @@ -241,10 +242,17 @@ enum Other { }, } -/// A compressed-tensors target: every linear, a module by name, or `re:` a regex from the start. +/// The module classes a model's checkpoint map can identify for class-name targets. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ModuleClass { + Linear, + Embedding, +} + +/// A compressed-tensors target: a module class, a module by name, or `re:` a regex from the start. #[derive(Debug)] enum Target { - Linear, + Class(ModuleClass), Name(String), Re(Regex), } @@ -256,14 +264,22 @@ impl Target { Regex::new(&format!("^(?:{re})")) .map_err(|e| bad(format!("compressed-tensors pattern {t}: {e}")))?, ), - None if t == "Linear" => Target::Linear, + None if t == "Linear" => Target::Class(ModuleClass::Linear), + None if t == "Embedding" => Target::Class(ModuleClass::Embedding), + None if !t.contains('.') + && t.as_bytes().first().is_some_and(u8::is_ascii_uppercase) => + { + return Err(bad(format!( + "compressed-tensors class {t} is not one this reader knows" + ))); + } None => Target::Name(t.to_string()), }) } - fn matches(&self, module: &str) -> bool { + fn matches(&self, module: &str, class: impl FnOnce() -> Option) -> bool { match self { - Target::Linear => true, + Target::Class(want) => class() == Some(*want), Target::Name(n) => n == module, Target::Re(re) => re.is_match(module), } @@ -313,7 +329,11 @@ impl Config { /// The formats `module` is declared, and the module whose scale a tensor-scale format shares; /// empty when the config leaves it unquantized. - fn declared(&self, module: &str) -> (Vec, Option<&str>) { + fn declared( + &self, + module: &str, + class_of: impl FnOnce() -> Option, + ) -> (Vec, Option<&str>) { if let Some(m) = &self.modelopt { let layer = m .layers @@ -350,12 +370,14 @@ impl Config { } } Some(Other::CompressedTensors { groups, ignore }) => { - if ignore.iter().any(|t| t.matches(module)) { + let class = LazyCell::new(class_of); + let matches = |t: &Target| t.matches(module, || *class); + if ignore.iter().any(matches) { vec![] } else { groups .iter() - .filter(|(ts, _)| ts.iter().any(|t| t.matches(module))) + .filter(|(ts, _)| ts.iter().any(matches)) .map(|(_, f)| *f) .collect() } @@ -627,10 +649,12 @@ const COMPANIONS: [&str; 7] = [ "weight_zero_point", ]; -/// Every tensor of `tensors` (name, dtype, shape) placed by `config`. +/// Every tensor of `tensors` (name, dtype, shape) placed by `config`. `class_of` supplies a module's +/// class from the model's map, or `None` for another or unknown role; only class targets ask for it. pub fn plan<'a>( config: &Config, tensors: impl IntoIterator, + class_of: impl Fn(&str) -> Option, ) -> Plan { let table: BTreeMap<&str, (Dtype, &[usize])> = tensors.into_iter().map(|(n, d, s)| (n, (d, s))).collect(); @@ -646,7 +670,7 @@ pub fn plan<'a>( if plain(dtype).is_some() { continue; } - let (formats, owner) = config.declared(module); + let (formats, owner) = config.declared(module, || class_of(module)); let found = formats .iter() .find_map(|&f| layout(f, name, module, owner, dtype, shape, &table).map(|c| (f, c))); @@ -1096,7 +1120,11 @@ mod tests { for (tag, probes) in cases { let c = config(tag).unwrap_or_else(|e| panic!("{tag}: {e}")); for (module, want) in *probes { - assert_eq!(c.declared(module).0, *want, "{tag}: {module}"); + assert_eq!( + c.declared(module, || Some(ModuleClass::Linear)).0, + *want, + "{tag}: {module}" + ); } } let e = config("nvidia__NVIDIA-Nemotron-3.5-Bolt-74B-A9B-24GB").unwrap_err(); @@ -1137,6 +1165,17 @@ mod tests { "weights": { "num_bits": 4, "type": "int", "strategy": "group", "group_size": group, "symmetric": true } } } } }) }; + for (field, class) in [("targets", "LayerNorm"), ("ignore", "Conv2d")] { + let mut config = int4(128.into()); + if field == "targets" { + config["quantization_config"]["config_groups"]["g"]["targets"] = + serde_json::json!([class]); + } else { + config["quantization_config"]["ignore"] = serde_json::json!([class]); + } + let e = Config::parse(&config, None).unwrap_err(); + assert!(e.to_string().contains(class), "{e}"); + } let fp8 = |key: &str, v: Value| { let mut q = serde_json::json!({ "quant_method": "fp8", "weight_block_size": [128, 128] }); @@ -1231,7 +1270,7 @@ mod tests { ("ple.ids", I64, &[3]), ("e.2.weight", U8, &[0, 9_223_372_036_854_775_840]), ]; - let p = plan(&c, t.iter().copied()); + let p = plan(&c, t.iter().copied(), |_| None); let got = |n: &str| p.placed.get(n).map(|pl| (pl.format, pl.companions.clone())); assert_eq!( got("e.0.weight").unwrap(), @@ -1277,7 +1316,7 @@ mod tests { ("l.mlp.gate.weight_packed", I32, &[8, 64]), ("y.weight_packed", I32, &[0, usize::MAX / 4]), ]; - let p = plan(&ct, t.iter().copied()); + let p = plan(&ct, t.iter().copied(), |_| Some(ModuleClass::Linear)); assert_eq!( p.placed["x.weight_packed"].format, Format::Int4Pack { group: 128 } @@ -1296,6 +1335,69 @@ mod tests { } } + #[test] + fn compressed_tensor_classes_do_not_quantize_each_others_weights() { + use Dtype::*; + let tensors: &[(&str, Dtype, &[usize])] = &[ + ("model.embed_tokens.weight_packed", I32, &[8, 16]), + ("model.embed_tokens.weight_scale", Bf16, &[8, 1]), + ("model.embed_tokens.weight_shape", I64, &[2]), + ( + "model.layers.0.self_attn.q_proj.weight_packed", + I32, + &[8, 16], + ), + ( + "model.layers.0.self_attn.q_proj.weight_scale", + Bf16, + &[8, 1], + ), + ("model.layers.0.self_attn.q_proj.weight_shape", I64, &[2]), + ]; + for (target, kept, refused) in [ + ( + "Embedding", + "model.embed_tokens", + "model.layers.0.self_attn.q_proj", + ), + ( + "Linear", + "model.layers.0.self_attn.q_proj", + "model.embed_tokens", + ), + ] { + let config = Config::parse( + &serde_json::json!({"quantization_config": { + "quant_method": "compressed-tensors", + "config_groups": {"g": { + "format": "pack-quantized", "targets": [target], + "weights": {"num_bits": 4, "type": "int", "strategy": "group", + "group_size": 128, "symmetric": true} + }} + }}), + None, + ) + .unwrap(); + let placed = plan(&config, tensors.iter().copied(), |module| match module { + "model.embed_tokens" => Some(ModuleClass::Embedding), + "model.layers.0.self_attn.q_proj" => Some(ModuleClass::Linear), + _ => None, + }); + assert_eq!( + placed + .placed + .get(&format!("{kept}.weight_packed")) + .map(|p| p.format), + Some(Format::Int4Pack { group: 128 }), + "{target}: {placed:?}" + ); + assert!( + placed.refused[&format!("{refused}.weight_packed")].contains("unquantized"), + "{target}: {placed:?}" + ); + } + } + /// Real headers: every tensor of nvidia/GLM-5.3-Flash-NVFP4's layers 0 (dense) and 3 (288 /// experts) and its unlayered tensors under its config, and a deepseek-ai/DeepSeek-V4.1-Flash /// shard under its own, each tensor placed and none refused. @@ -1339,6 +1441,7 @@ mod tests { entries .iter() .map(|(n, d, s)| (n.as_str(), *d, s.as_slice())), + |_| None, ); let mut count = BTreeMap::new(); for pl in p.placed.values() { @@ -1518,6 +1621,14 @@ mod tests { plan( &config(tag).unwrap(), entries.iter().map(|(n, d, s)| (*n, *d, s.as_slice())), + |module| { + matches!( + module, + "model.language_model.layers.20.mlp.experts.148.down_proj" + | "model.language_model.layers.20.mlp.experts.149.down_proj" + ) + .then_some(ModuleClass::Linear) + }, ) }; let m = Use::Matrix(Plain::Q8_0); diff --git a/crates/hf/tests/fixtures/quant-configs/README.md b/crates/hf/tests/fixtures/quant-configs/README.md index 63e33157f..073a15951 100644 --- a/crates/hf/tests/fixtures/quant-configs/README.md +++ b/crates/hf/tests/fixtures/quant-configs/README.md @@ -19,3 +19,11 @@ Bolt's two files list 27,921 modules (15 MB); they are cut to the first module of each algorithm (`quantized_layers`) and the first group (`config_groups`), the rest unchanged. `../checkpoint-headers/glm-5.3-flash-nvfp4-layers-0-3.tensors.json.gz` is the dtype and shape of every tensor of `nvidia/GLM-5.3-Flash-NVFP4` layers 0 and 3 and every unlayered `model.language_model.` tensor and `lm_head.weight`, gathered from all 33 shard headers. + +**Class-name targets.** `quant_config::plan` takes an explicit module-class resolver; the checkpoint-module +classifier in `model_hf::arch::module_class` supplies the projection and embedding roles read by the GLM and +Qwen3.8 maps. `Linear` and `Embedding` select only their own classes, in both `targets` and `ignore`. +Norms, convolutions and unknown roles do not match a class merely because their stored weights look like matrices. +Module-name and `re:` targets still match their names; an unsupported bare capitalized class name is refused by +name while reading the config. BF16 and F32 tensors remain placed as stored, independently of quantization targets. + diff --git a/crates/model-hf/src/arch/glm5next.rs b/crates/model-hf/src/arch/glm5next.rs index f71938d08..56d3d5b17 100644 --- a/crates/model-hf/src/arch/glm5next.rs +++ b/crates/model-hf/src/arch/glm5next.rs @@ -25,6 +25,8 @@ use model::placement::tp::{TpTensor, card_blocks}; use safetensors::{Dtype, Folder}; use serde_json::Value; +use super::module_class; + /// Why the checkpoint was refused. #[derive(Debug, thiserror::Error)] pub enum Error { @@ -357,6 +359,7 @@ pub fn tp_tensors( let placed = quant_config::plan( config, table.iter().map(|(n, d, s)| (n.as_str(), *d, s.as_slice())), + module_class, ); if let Some((name, why)) = placed.refused.iter().next() { return Err(refuse(name, why.clone())); diff --git a/crates/model-hf/src/arch/mod.rs b/crates/model-hf/src/arch/mod.rs index a59d9fe09..0ef229d38 100644 --- a/crates/model-hf/src/arch/mod.rs +++ b/crates/model-hf/src/arch/mod.rs @@ -2,3 +2,162 @@ pub mod glm5next; pub mod qwen35moe; + +/// The classes of weight modules the checkpoint maps read. Projection, embedding, norm and +/// convolution are distinct even when they share a placement role or a two-dimensional shape. +pub fn module_class(module: &str) -> Option { + use hf::quant_config::ModuleClass::{Embedding, Linear}; + let short = module + .strip_prefix("model.language_model.") + .or_else(|| module.strip_prefix("model.")) + .unwrap_or(module); + match short { + "embed_tokens" => return Some(Embedding), + "lm_head" => return Some(Linear), + "hyper_connection_mixer.input_mix_weight_down" + | "hyper_connection_mixer.input_mix_weight_up" => return Some(Linear), + _ => {} + } + let (layer, rest) = short.strip_prefix("layers.")?.split_once('.')?; + layer.parse::().ok()?; + if let Some(expert) = rest.strip_prefix("mlp.experts.") { + let (id, projection) = expert.split_once('.')?; + id.parse::().ok()?; + return matches!(projection, "gate_proj" | "up_proj" | "down_proj").then_some(Linear); + } + if let Some(part) = rest.strip_prefix("ple.ple_embedding.ngram_embedding.shard_") { + part.parse::().ok()?; + return Some(Embedding); + } + if let Some(mixer) = rest + .strip_prefix("attn_hyper_connection.") + .or_else(|| rest.strip_prefix("mlp_hyper_connection.")) + { + return matches!(mixer, "input_mix_weight_down" | "input_mix_weight_up").then_some(Linear); + } + matches!( + rest, + "eh_proj" + | "self_attn.q_proj" + | "self_attn.k_proj" + | "self_attn.v_proj" + | "self_attn.o_proj" + | "self_attn.f_a_proj" + | "self_attn.f_b_proj" + | "self_attn.g_a_proj" + | "self_attn.g_b_proj" + | "self_attn.b_proj" + | "self_attn.q_a_proj" + | "self_attn.q_b_proj" + | "self_attn.kv_a_proj_with_mqa" + | "self_attn.kv_b_proj" + | "self_attn.indexer.wq_b" + | "self_attn.indexer.wk" + | "self_attn.indexer.weights_proj" + | "self_attn.indexer.index_qk_proj" + | "linear_attn.in_proj_qkv" + | "linear_attn.in_proj_z" + | "linear_attn.in_proj_a" + | "linear_attn.in_proj_b" + | "linear_attn.out_proj" + | "mlp.gate" + | "mlp.gate_proj" + | "mlp.up_proj" + | "mlp.down_proj" + | "mlp.shared_experts.gate_proj" + | "mlp.shared_experts.up_proj" + | "mlp.shared_experts.down_proj" + | "mlp.shared_expert.gate_proj" + | "mlp.shared_expert.up_proj" + | "mlp.shared_expert.down_proj" + | "mlp.shared_expert_gate" + | "ple.key_proj" + | "ple.value_proj" + ) + .then_some(Linear) +} + +#[cfg(test)] +mod tests { + use hf::quant_config::{Config, ModuleClass, plan}; + use safetensors::Dtype; + + #[test] + fn class_targets_select_projection_and_embedding_roles_not_shapes() { + use ModuleClass::{Embedding, Linear}; + let modules = [ + ("model.embed_tokens", Some(Embedding)), + ("model.language_model.embed_tokens", Some(Embedding)), + ("lm_head", Some(Linear)), + ("model.layers.0.self_attn.q_proj", Some(Linear)), + ( + "model.language_model.layers.0.linear_attn.in_proj_a", + Some(Linear), + ), + ("model.layers.3.mlp.experts.0.down_proj", Some(Linear)), + ( + "model.language_model.layers.0.ple.ple_embedding.ngram_embedding.shard_0", + Some(Embedding), + ), + ("model.layers.0.input_layernorm", None), + ("model.language_model.layers.0.linear_attn.conv1d", None), + ( + "model.language_model.layers.0.attn_hyper_connection.block_inject_weight", + None, + ), + ("model.layers.0.self_attn.unknown_proj", None), + ]; + let tensors: Vec<_> = modules + .iter() + .flat_map(|(module, _)| { + [ + (format!("{module}.weight_packed"), Dtype::I32, vec![8, 16]), + (format!("{module}.weight_scale"), Dtype::Bf16, vec![8, 1]), + (format!("{module}.weight_shape"), Dtype::I64, vec![2]), + ] + }) + .collect(); + for (ignore, ignored) in [ + ("", None), + ("Linear", Some(Linear)), + ("Embedding", Some(Embedding)), + ] { + let config = Config::parse( + &serde_json::json!({"quantization_config": { + "quant_method": "compressed-tensors", + "config_groups": {"g": { + "format": "pack-quantized", "targets": ["Linear", "Embedding"], + "weights": {"num_bits": 4, "type": "int", "strategy": "group", + "group_size": 128, "symmetric": true} + }}, + "ignore": if ignore.is_empty() { vec![] } else { vec![ignore] } + }}), + None, + ) + .unwrap(); + let placed = plan( + &config, + tensors + .iter() + .map(|(name, dtype, shape)| (name.as_str(), *dtype, shape.as_slice())), + super::module_class, + ); + for (module, class) in modules { + let name = format!("{module}.weight_packed"); + if class.is_some() && class != ignored { + assert_eq!( + placed.placed.get(&name).map(|p| p.format), + Some(hf::quant_config::Format::Int4Pack { group: 128 }), + "{ignore}/{name}: {placed:?}" + ); + } else { + let why = &placed.refused[&name]; + assert!( + why.contains(module) && why.contains("unquantized"), + "{ignore}/{why}" + ); + } + } + } + } +} diff --git a/crates/model-hf/src/arch/qwen35moe.rs b/crates/model-hf/src/arch/qwen35moe.rs index ee3968392..3aa7c224e 100644 --- a/crates/model-hf/src/arch/qwen35moe.rs +++ b/crates/model-hf/src/arch/qwen35moe.rs @@ -1346,6 +1346,7 @@ mod tests { entries .iter() .map(|(n, d, s)| (n.as_str(), *d, s.as_slice())), + crate::arch::module_class, ); let logical: Vec<(&str, Vec)> = entries .iter() @@ -2120,6 +2121,7 @@ mod tests { entries .iter() .map(|(n, d, s)| (n.as_str(), *d, s.as_slice())), + crate::arch::module_class, ); let got: Vec<_> = m.parts(&p, |n| st.raw(n)).collect(); if let Some(want) = want { diff --git a/crates/model-hf/src/bin/prepare.rs b/crates/model-hf/src/bin/prepare.rs index 9d6b3e7ed..2aa069cc8 100644 --- a/crates/model-hf/src/bin/prepare.rs +++ b/crates/model-hf/src/bin/prepare.rs @@ -251,6 +251,7 @@ fn plan_tensors( entries .iter() .map(|(n, d, s)| (n.as_str(), *d, s.as_slice())), + model_hf::arch::module_class, ); let logical: Vec<(&str, Vec)> = entries .iter() diff --git a/crates/model-hf/tests/common/glm_hf.rs b/crates/model-hf/tests/common/glm_hf.rs index ff49d4363..d55f5d30a 100644 --- a/crates/model-hf/tests/common/glm_hf.rs +++ b/crates/model-hf/tests/common/glm_hf.rs @@ -5,7 +5,7 @@ use hf::quant_config::{Config, plan}; use safetensors::Dtype; use serde_json::Value; -use super::{Geometry, Layout}; +use super::{Geometry, Layout, module_class}; pub(super) const GLM: Geometry = Geometry { layers: 45, @@ -33,6 +33,7 @@ pub(super) fn logical( let p = plan( config, table.iter().map(|(n, d, s)| (n.as_str(), *d, s.as_slice())), + module_class, ); assert!(p.refused.is_empty(), "{:?}", p.refused); table diff --git a/crates/model-hf/tests/glm5next_map.rs b/crates/model-hf/tests/glm5next_map.rs index 4a03205c9..3bfaf4e4e 100644 --- a/crates/model-hf/tests/glm5next_map.rs +++ b/crates/model-hf/tests/glm5next_map.rs @@ -6,6 +6,7 @@ use std::collections::{BTreeMap, BTreeSet}; use hf::quant_config::Config; use model_hf::arch::glm5next::{Geometry, Layout, map}; +use model_hf::arch::module_class; use safetensors::Dtype; #[path = "common/glm_hf.rs"]