quantize_rs/
config.rs

1//! YAML and TOML configuration file support.
2//!
3//! A configuration file can specify global quantization settings
4//! (`bits`, `per_channel`), per-model overrides, and batch processing
5//! parameters.
6
7use crate::errors::{QuantizeError, Result};
8use serde::{Deserialize, Serialize};
9use std::path::Path;
10
11/// Top-level quantization configuration.
12///
13/// Can be loaded from a YAML or TOML file with [`Config::from_file`].
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct Config {
16    /// Default bit width (4 or 8). Defaults to 8.
17    #[serde(default = "default_bits")]
18    pub bits: u8,
19
20    /// Default per-channel setting. Defaults to `false`.
21    #[serde(default)]
22    pub per_channel: bool,
23
24    /// Layer names to exclude from quantization globally.
25    #[serde(default)]
26    pub excluded_layers: Vec<String>,
27
28    /// Minimum number of elements a tensor must have to be quantized.
29    /// Tensors smaller than this are kept in FP32. Defaults to 0 (no minimum).
30    #[serde(default)]
31    pub min_elements: usize,
32
33    /// Store INT4 weights as native ONNX `DataType::Int4` (opset 21).
34    /// Only affects models quantized with `bits=4` or per-layer INT4 overrides.
35    /// Defaults to `false` (widen INT4 to INT8, compatible with opset 10+).
36    #[serde(default)]
37    pub native_int4: bool,
38
39    /// Use symmetric quantization (zero_point == 0).  Required by most
40    /// ONNX Runtime / TensorRT INT8 matmul kernels for per-channel weight
41    /// quantization.  Defaults to `false` (asymmetric).
42    #[serde(default)]
43    pub symmetric: bool,
44
45    /// Per-model configuration overrides.
46    #[serde(default)]
47    pub models: Vec<ModelConfig>,
48
49    /// Batch processing configuration.
50    #[serde(default)]
51    pub batch: Option<BatchConfig>,
52}
53
54/// Per-model quantization overrides.
55#[derive(Debug, Clone, Serialize, Deserialize)]
56pub struct ModelConfig {
57    /// Path to the input ONNX model.
58    pub input: String,
59
60    /// Path for the quantized output model.
61    pub output: String,
62
63    /// Override bit width for this model.
64    #[serde(default)]
65    pub bits: Option<u8>,
66
67    /// Override per-channel setting for this model.
68    #[serde(default)]
69    pub per_channel: Option<bool>,
70
71    /// Skip this model if the output file already exists.
72    #[serde(default)]
73    pub skip_existing: bool,
74
75    /// Layer names to exclude from quantization for this model.
76    /// Merged with (but does not replace) the global `excluded_layers`.
77    #[serde(default)]
78    pub excluded_layers: Vec<String>,
79
80    /// Per-layer bit-width overrides for this model.
81    /// Key = initializer name, value = 4 or 8.
82    #[serde(default)]
83    pub layer_bits: std::collections::HashMap<String, u8>,
84
85    /// Override the global `min_elements` threshold for this model.
86    #[serde(default)]
87    pub min_elements: Option<usize>,
88
89    /// Override the global `native_int4` flag for this model.
90    #[serde(default)]
91    pub native_int4: Option<bool>,
92
93    /// Override the global `symmetric` flag for this model.
94    #[serde(default)]
95    pub symmetric: Option<bool>,
96}
97
98/// Batch processing configuration for quantizing multiple models.
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct BatchConfig {
101    /// Glob pattern or directory for input models.
102    pub input_dir: String,
103
104    /// Output directory for quantized models.
105    pub output_dir: String,
106
107    /// Skip models whose output already exists.
108    #[serde(default)]
109    pub skip_existing: bool,
110
111    /// Continue processing remaining models after a failure.
112    #[serde(default)]
113    pub continue_on_error: bool,
114
115    /// Number of models to quantize in parallel.  Defaults to 1 (serial).
116    #[serde(default = "default_jobs")]
117    pub jobs: usize,
118}
119
120fn default_jobs() -> usize {
121    1
122}
123
124fn default_bits() -> u8 {
125    8
126}
127
128impl Config {
129    /// Load a config from a YAML or TOML file (auto-detected by extension).
130    ///
131    /// # Errors
132    ///
133    /// Returns [`QuantizeError::Config`] on I/O, parse, or unsupported format errors.
134    pub fn from_file(path: impl AsRef<Path>) -> Result<Self> {
135        let path = path.as_ref();
136        let extension =
137            path.extension()
138                .and_then(|s| s.to_str())
139                .ok_or_else(|| QuantizeError::Config {
140                    reason: "Config file has no extension".into(),
141                })?;
142
143        let content = std::fs::read_to_string(path).map_err(|e| QuantizeError::Config {
144            reason: format!("Failed to read config file '{}': {e}", path.display()),
145        })?;
146
147        match extension {
148            "yaml" | "yml" => Self::from_yaml(&content),
149            "toml" => Self::from_toml(&content),
150            _ => Err(QuantizeError::Config {
151                reason: format!("Unsupported config format: {}", extension),
152            }),
153        }
154    }
155
156    /// Parse configuration from a YAML string.
157    pub fn from_yaml(content: &str) -> Result<Self> {
158        serde_yaml::from_str(content).map_err(|e| QuantizeError::Config {
159            reason: format!("Failed to parse YAML config: {e}"),
160        })
161    }
162
163    /// Parse configuration from a TOML string.
164    pub fn from_toml(content: &str) -> Result<Self> {
165        toml::from_str(content).map_err(|e| QuantizeError::Config {
166            reason: format!("Failed to parse TOML config: {e}"),
167        })
168    }
169
170    /// Validate the configuration (bits values, non-empty paths).
171    ///
172    /// # Errors
173    ///
174    /// Returns [`QuantizeError::Config`] if any field is invalid.
175    pub fn validate(&self) -> Result<()> {
176        if self.bits != 4 && self.bits != 8 {
177            return Err(QuantizeError::Config {
178                reason: format!("Invalid bits value: {}. Must be 4 or 8", self.bits),
179            });
180        }
181
182        for (idx, model) in self.models.iter().enumerate() {
183            if model.input.is_empty() {
184                return Err(QuantizeError::Config {
185                    reason: format!("Model {}: input path is empty", idx),
186                });
187            }
188            if model.output.is_empty() {
189                return Err(QuantizeError::Config {
190                    reason: format!("Model {}: output path is empty", idx),
191                });
192            }
193            if let Some(bits) = model.bits {
194                if bits != 4 && bits != 8 {
195                    return Err(QuantizeError::Config {
196                        reason: format!("Model {}: invalid bits value: {}", idx, bits),
197                    });
198                }
199            }
200            for (layer, &bits) in &model.layer_bits {
201                if layer.is_empty() {
202                    return Err(QuantizeError::Config {
203                        reason: format!("Model {}: layer_bits contains an empty layer name", idx),
204                    });
205                }
206                if bits != 4 && bits != 8 {
207                    return Err(QuantizeError::Config {
208                        reason: format!(
209                            "Model {}: invalid bits {} for layer '{}'",
210                            idx, bits, layer
211                        ),
212                    });
213                }
214            }
215        }
216
217        if let Some(batch) = &self.batch {
218            if batch.input_dir.is_empty() {
219                return Err(QuantizeError::Config {
220                    reason: "Batch input_dir is empty".into(),
221                });
222            }
223            if batch.output_dir.is_empty() {
224                return Err(QuantizeError::Config {
225                    reason: "Batch output_dir is empty".into(),
226                });
227            }
228        }
229
230        Ok(())
231    }
232
233    /// Effective bit width for a model (model override or global default).
234    pub fn get_bits(&self, model: &ModelConfig) -> u8 {
235        model.bits.unwrap_or(self.bits)
236    }
237
238    /// Effective per-channel setting for a model (model override or global default).
239    pub fn get_per_channel(&self, model: &ModelConfig) -> bool {
240        model.per_channel.unwrap_or(self.per_channel)
241    }
242
243    /// Effective excluded-layers list: global list merged with model-level list.
244    pub fn get_excluded_layers(&self, model: &ModelConfig) -> Vec<String> {
245        let mut layers = self.excluded_layers.clone();
246        for l in &model.excluded_layers {
247            if !layers.contains(l) {
248                layers.push(l.clone());
249            }
250        }
251        layers
252    }
253
254    /// Effective min-elements threshold for a model.
255    pub fn get_min_elements(&self, model: &ModelConfig) -> usize {
256        model.min_elements.unwrap_or(self.min_elements)
257    }
258
259    /// Effective native-INT4 flag for a model (override or global default).
260    pub fn get_native_int4(&self, model: &ModelConfig) -> bool {
261        model.native_int4.unwrap_or(self.native_int4)
262    }
263
264    /// Effective symmetric flag for a model (override or global default).
265    pub fn get_symmetric(&self, model: &ModelConfig) -> bool {
266        model.symmetric.unwrap_or(self.symmetric)
267    }
268
269    /// Effective per-layer bit-width overrides for a model.
270    ///
271    /// Layer names are model-specific so there is no global map to merge;
272    /// this simply returns the model's own `layer_bits` map.
273    pub fn get_layer_bits(&self, model: &ModelConfig) -> std::collections::HashMap<String, u8> {
274        model.layer_bits.clone()
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    #[test]
283    fn test_yaml_config() {
284        let yaml = r#"
285bits: 8
286per_channel: true
287
288models:
289  - input: model1.onnx
290    output: model1_int8.onnx
291  
292  - input: model2.onnx
293    output: model2_int8.onnx
294    per_channel: false
295
296batch:
297  input_dir: "models/*.onnx"
298  output_dir: quantized/
299  skip_existing: true
300"#;
301
302        let config = Config::from_yaml(yaml).unwrap();
303        assert_eq!(config.bits, 8);
304        assert!(config.per_channel);
305        assert_eq!(config.models.len(), 2);
306        assert!(config.batch.is_some());
307    }
308
309    #[test]
310    fn test_empty_layer_bits_key_rejected() {
311        let yaml = r#"
312bits: 8
313models:
314  - input: model.onnx
315    output: out.onnx
316    layer_bits:
317      "": 4
318"#;
319        let config = Config::from_yaml(yaml).unwrap();
320        let err = config.validate().unwrap_err();
321        assert!(matches!(err, crate::errors::QuantizeError::Config { .. }));
322        assert!(err.to_string().contains("empty layer name"));
323    }
324
325    #[test]
326    fn test_toml_config() {
327        let toml = r#"
328bits = 8
329per_channel = true
330
331[[models]]
332input = "model1.onnx"
333output = "model1_int8.onnx"
334
335[[models]]
336input = "model2.onnx"
337output = "model2_int8.onnx"
338per_channel = false
339
340[batch]
341input_dir = "models/*.onnx"
342output_dir = "quantized/"
343skip_existing = true
344"#;
345
346        let config = Config::from_toml(toml).unwrap();
347        assert_eq!(config.bits, 8);
348        assert!(config.per_channel);
349        assert_eq!(config.models.len(), 2);
350        assert!(config.batch.is_some());
351    }
352}
quantize_rs/config.rs

quantize_rs/
config.rs
