hf: resolve compressed-tensors targets by checkpoint module class

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]);

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<ModuleClass>) -> 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<Format>, Option<&str>) {
+    fn declared(
+        &self,
+        module: &str,
+        class_of: impl FnOnce() -> Option<ModuleClass>,
+    ) -> (Vec<Format>, 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<Item = (&'a str, Dtype, &'a [usize])>,
+    class_of: impl Fn(&str) -> Option<ModuleClass>,
 ) -> 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);

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.
+

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()));