metrale_config/parsers/
quantization.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Parser for the `quantization_config` block (compressed-tensors, or a ModelOpt
4//! `hf_quant_config.json` placed in that slot) into `QuantizationConfig`.
5//!
6//! Owner: config (quantization).
7//! Invariants: none beyond the types.
8
9#![allow(unused_imports)]
10
11use anyhow::{Context, Result};
12use serde_json::Value;
13
14use super::super::{ModelConfig, QuantizationConfig};
15
16/// 2026-09-28: `Ok(None)` when there is no block or it declares nothing; an error when the
17/// block's precision declaration is malformed (`DeclaredPrecisionPlan`).
18pub fn parse_quantization_config(raw: &serde_json::Value) -> Result<Option<QuantizationConfig>> {
19    let Some(qc_raw) = raw.get("quantization_config") else {
20        return Ok(None);
21    };
22    // 2026-09-26: A ModelOpt `hf_quant_config.json`, which `merge_sidecar_quant_config` places
23    // in this slot, nests the fields under `"quantization"` and has no `quant_method`; the
24    // scheme is named by `producer.name == "modelopt"` (the Nemotron-3 Nano and Super
25    // checkpoints in the local HF cache):
26    //
27    //   { "producer": { "name": "modelopt", ... },
28    //     "quantization": { "quant_algo": "NVFP4",
29    //                       "exclude_modules": ["lm_head", ...] } }
30    //
31    // Read as is, it would give `None`. `normalize_modelopt_sidecar` turns it into the flat
32    // shape; a flat block passes through unchanged.
33    let canonical = normalize_modelopt_sidecar(qc_raw);
34    let qc = &canonical;
35
36    let quant_method = qc
37        .get("quant_method")
38        .and_then(serde_json::Value::as_str)
39        .unwrap_or("")
40        .to_string();
41    let quant_algo = qc
42        .get("quant_algo")
43        .and_then(serde_json::Value::as_str)
44        .or_else(|| {
45            // 2026-09-26: Without `quant_algo`, the label comes from
46            // `config_groups.group_0.weights`: a 4-bit float is NVFP4, an 8-bit float FP8.
47            let group = qc.get("config_groups")?.get("group_0")?;
48            let weights = group.get("weights")?;
49            let bits = weights.get("num_bits")?.as_u64()?;
50            let ty = weights.get("type")?.as_str()?;
51            match (bits, ty) {
52                (4, "float") => Some("NVFP4"),
53                (8, "float") => Some("FP8"),
54                _ => None,
55            }
56        })
57        .unwrap_or("")
58        .to_string();
59    let format = qc
60        .get("format")
61        .and_then(serde_json::Value::as_str)
62        .unwrap_or("")
63        .to_string();
64
65    // 2026-09-26: The ignore list is `ignore` followed by the `exclude_modules` entries not
66    // already in it.
67    let mut ignore_modules: Vec<String> = Vec::new();
68    if let Some(arr) = qc.get("ignore").and_then(serde_json::Value::as_array) {
69        for v in arr {
70            if let Some(s) = v.as_str() {
71                ignore_modules.push(s.to_string());
72            }
73        }
74    }
75    if let Some(arr) = qc
76        .get("exclude_modules")
77        .and_then(serde_json::Value::as_array)
78    {
79        for v in arr {
80            if let Some(s) = v.as_str()
81                && !ignore_modules.contains(&s.to_string())
82            {
83                ignore_modules.push(s.to_string());
84            }
85        }
86    }
87
88    // 2026-09-26: No method, algorithm or ignore list means no quantization config.
89    if quant_method.is_empty() && quant_algo.is_empty() && ignore_modules.is_empty() {
90        return Ok(None);
91    }
92    let precision = crate::precision_plan::DeclaredPrecisionPlan::from_quantization_config(qc_raw)
93        .context("quantization_config")?;
94
95    Ok(Some(QuantizationConfig {
96        quant_method,
97        quant_algo,
98        format,
99        ignore_modules,
100        precision,
101    }))
102}
103
104/// 2026-09-26: Flatten a ModelOpt `hf_quant_config.json` payload into the shape
105/// [`parse_quantization_config`] reads. Both steps leave a flat block unchanged:
106///   1. A nested `"quantization"` object replaces the payload.
107///   2. Without a non-empty `quant_method`, `producer.name == "modelopt"` (any case) adds
108///      `quant_method = "modelopt"`, the value `detect_quant_format` and
109///      `detect_nvfp4_variant` match on.
110fn normalize_modelopt_sidecar(qc_raw: &serde_json::Value) -> serde_json::Value {
111    let mut canonical = match qc_raw.get("quantization") {
112        Some(serde_json::Value::Object(inner)) => serde_json::Value::Object(inner.clone()),
113        _ => qc_raw.clone(),
114    };
115
116    let Some(obj) = canonical.as_object_mut() else {
117        return canonical;
118    };
119
120    let has_method = obj
121        .get("quant_method")
122        .and_then(serde_json::Value::as_str)
123        .is_some_and(|s| !s.is_empty());
124    if !has_method {
125        let producer_modelopt = qc_raw
126            .get("producer")
127            .and_then(|p| p.get("name"))
128            .and_then(serde_json::Value::as_str)
129            .is_some_and(|n| n.eq_ignore_ascii_case("modelopt"));
130        if producer_modelopt {
131            obj.insert(
132                "quant_method".to_string(),
133                serde_json::Value::String("modelopt".to_string()),
134            );
135        }
136    }
137
138    canonical
139}
140
141#[cfg(test)]
142mod tests {
143    use super::{normalize_modelopt_sidecar, parse_quantization_config};
144
145    /// 2026-09-26: The `hf_quant_config.json` shape of the local
146    /// `nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4` checkpoint (exclusion list cut to three),
147    /// placed in the `quantization_config` slot as `merge_sidecar_quant_config` does.
148    #[test]
149    fn modelopt_sidecar_nested_schema_preserves_exclusions() {
150        let raw = serde_json::json!({
151            "quantization_config": {
152                "producer": { "name": "modelopt", "version": "0.29.0" },
153                "quantization": {
154                    "quant_algo": "NVFP4",
155                    "kv_cache_quant_algo": "FP8",
156                    "group_size": 16,
157                    "exclude_modules": [
158                        "lm_head",
159                        "backbone.layers.4.mixer.in_proj",
160                        "backbone.layers.0.mixer.conv1d"
161                    ]
162                }
163            }
164        });
165        let qc = parse_quantization_config(&raw)
166            .expect("parses")
167            .expect("ModelOpt nested sidecar must yield a QuantizationConfig");
168        assert_eq!(qc.quant_method, "modelopt");
169        assert_eq!(qc.quant_algo, "NVFP4");
170        assert_eq!(qc.ignore_modules.len(), 3);
171        assert!(qc.ignore_modules.iter().any(|m| m == "lm_head"));
172        assert!(
173            qc.ignore_modules
174                .iter()
175                .any(|m| m == "backbone.layers.4.mixer.in_proj")
176        );
177        assert!(
178            qc.ignore_modules
179                .iter()
180                .any(|m| m == "backbone.layers.0.mixer.conv1d")
181        );
182    }
183
184    /// 2026-09-26: A flat compressed-tensors block is unchanged by the ModelOpt normalization.
185    #[test]
186    fn flat_compressed_tensors_block_is_not_normalized() {
187        let raw = serde_json::json!({
188            "quantization_config": {
189                "quant_method": "compressed-tensors",
190                "format": "nvfp4-pack-quantized",
191                "ignore": ["lm_head"]
192            }
193        });
194        let flat = &raw["quantization_config"];
195        assert_eq!(normalize_modelopt_sidecar(flat), flat.clone());
196
197        let qc = parse_quantization_config(&raw)
198            .expect("parses")
199            .expect("flat block must still parse");
200        assert_eq!(qc.quant_method, "compressed-tensors");
201        assert_eq!(qc.format, "nvfp4-pack-quantized");
202        assert_eq!(qc.ignore_modules, vec!["lm_head".to_string()]);
203    }
204
205    /// 2026-09-26: The ModelOpt mixed-precision shape of the local Nemotron-3 Super 120B
206    /// checkpoint (nested, no `exclude_modules`) resolves `quant_method = "modelopt"`.
207    #[test]
208    fn modelopt_mixed_precision_sidecar_parses_without_exclusions() {
209        let raw = serde_json::json!({
210            "quantization_config": {
211                "producer": { "name": "modelopt", "version": "0.43.0" },
212                "quantization": {
213                    "quant_algo": "MIXED_PRECISION",
214                    "kv_cache_quant_algo": "FP8",
215                    "quantized_layers": {
216                        "backbone.layers.0.mixer.in_proj": { "quant_algo": "FP8" }
217                    }
218                }
219            }
220        });
221        let qc = parse_quantization_config(&raw)
222            .expect("parses")
223            .expect("mixed-precision sidecar must yield a QuantizationConfig");
224        assert_eq!(qc.quant_method, "modelopt");
225        assert_eq!(qc.quant_algo, "MIXED_PRECISION");
226        assert!(qc.ignore_modules.is_empty());
227    }
228}
metrale_config/parsers/quantization.rs

metrale_config/parsers/
quantization.rs
