teeny_quant/
validate.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17//! Tensor-level (not full-model) quantization error metrics: for each quantized tensor,
18//! dequantize and compare against the original `f32` weights. No forward pass is needed -- this
19//! is a pure diff between the original and quantized-then-reconstructed safetensors files.
20
21use std::path::Path;
22
23use crate::error::{Error, Result};
24use crate::format;
25use crate::quant::Scheme;
26use crate::quant::compute_groups;
27use crate::quant::fp8::Fp8Variant;
28use crate::quant::granularity::Granularity;
29use crate::read::read_f32;
30
31/// Per-tensor quantization error metrics.
32#[derive(Debug, Clone, serde::Serialize)]
33pub struct TensorReport {
34    /// The (original, unsuffixed) tensor name.
35    pub name: String,
36    /// Largest absolute error between original and dequantized-reconstructed elements.
37    pub max_abs_error: f32,
38    /// Mean absolute error across all elements.
39    pub mean_abs_error: f32,
40    /// Signal-to-quantization-noise ratio in dB (higher is better; `+inf` for an exact match).
41    pub sqnr_db: f32,
42}
43
44/// Compares `original` against `reconstructed` element-wise, computing [`TensorReport`] metrics.
45/// Both slices must be the same length (the caller is responsible for shape bookkeeping).
46pub fn compare_tensors(
47    name: &str,
48    original: &[f32],
49    reconstructed: &[f32],
50) -> Result<TensorReport> {
51    if original.len() != reconstructed.len() {
52        return Err(Error::ShapeMismatch {
53            name: name.to_string(),
54            a: vec![original.len()],
55            b: vec![reconstructed.len()],
56        });
57    }
58    if original.is_empty() {
59        return Ok(TensorReport {
60            name: name.to_string(),
61            max_abs_error: 0.0,
62            mean_abs_error: 0.0,
63            sqnr_db: f32::INFINITY,
64        });
65    }
66
67    let mut max_abs_error = 0f32;
68    let mut sum_abs_error = 0f64;
69    let mut signal_energy = 0f64;
70    let mut noise_energy = 0f64;
71    for (&o, &r) in original.iter().zip(reconstructed.iter()) {
72        let err = (o - r).abs();
73        max_abs_error = max_abs_error.max(err);
74        sum_abs_error += err as f64;
75        signal_energy += (o as f64) * (o as f64);
76        noise_energy += (err as f64) * (err as f64);
77    }
78
79    let mean_abs_error = (sum_abs_error / original.len() as f64) as f32;
80    let sqnr_db = if noise_energy == 0.0 {
81        f32::INFINITY
82    } else {
83        (10.0 * (signal_energy / noise_energy).log10()) as f32
84    };
85
86    Ok(TensorReport {
87        name: name.to_string(),
88        max_abs_error,
89        mean_abs_error,
90        sqnr_db,
91    })
92}
93
94fn granularity_from_config(weights: &format::WeightsConfig) -> Result<Granularity> {
95    match (weights.strategy.as_str(), weights.axis, weights.group_size) {
96        ("tensor", _, _) => Ok(Granularity::PerTensor),
97        ("channel", Some(axis), _) => Ok(Granularity::PerChannel { axis }),
98        ("group", Some(axis), Some(group_size)) => Ok(Granularity::Group { axis, group_size }),
99        _ => Err(Error::TensorNotFound(format!(
100            "quantization_config has strategy '{}' but is missing axis/group_size",
101            weights.strategy
102        ))),
103    }
104}
105
106fn scheme_from_config(weights: &format::WeightsConfig) -> Option<Scheme> {
107    match (weights.type_.as_str(), weights.num_bits) {
108        ("int", 8) => Some(Scheme::Int8 {
109            symmetric: weights.symmetric,
110        }),
111        ("int", 4) => Some(Scheme::Int4 {
112            symmetric: weights.symmetric,
113        }),
114        ("float", 8) => {
115            // The two FP8 encodings are not bit-compatible, so decoding with the wrong one
116            // silently produces garbage rather than a compile/runtime error -- the variant used
117            // must come from the file's own quantization_config, not be guessed.
118            let variant = match weights.fp8_variant.as_deref() {
119                Some("e4m3") => Fp8Variant::E4M3,
120                Some("e5m2") => Fp8Variant::E5M2,
121                _ => return None,
122            };
123            Some(Scheme::Fp8 { variant })
124        }
125        _ => None,
126    }
127}
128
129/// Compares every quantized tensor in `quantized_path` (as recorded by its embedded
130/// `quantization_config`) against the corresponding tensor in `original_path`, returning one
131/// [`TensorReport`] per quantized tensor, in the checkpoint's tensor order.
132pub fn validate_checkpoint(
133    original_path: &Path,
134    quantized_path: &Path,
135) -> Result<Vec<TensorReport>> {
136    let original_mapped = teeny_data::safetensors::SafeTensors::from_pretrained(original_path)?;
137    let original_tensors = original_mapped.tensors()?;
138
139    let quantized_mapped = teeny_data::safetensors::SafeTensors::from_pretrained(quantized_path)?;
140    let quantized_tensors = quantized_mapped.tensors()?;
141
142    let metadata = crate::read::read_metadata(quantized_path)?;
143    let config = format::config_from_metadata(&metadata)?.ok_or_else(|| {
144        Error::TensorNotFound(format!(
145            "'{}' has no quantization_config metadata -- was it produced by `teeny-quant quantize`?",
146            quantized_path.display()
147        ))
148    })?;
149
150    let weights = &config
151        .config_groups
152        .get("group_0")
153        .ok_or_else(|| {
154            Error::TensorNotFound("quantization_config.config_groups.group_0".to_string())
155        })?
156        .weights;
157    let scheme = scheme_from_config(weights).ok_or_else(|| {
158        Error::TensorNotFound(format!(
159            "unrecognized scheme in quantization_config: type={} num_bits={}",
160            weights.type_, weights.num_bits
161        ))
162    })?;
163
164    let mut reports = Vec::new();
165    for name in quantized_tensors.names() {
166        if config.ignore.iter().any(|i| i == name) {
167            continue;
168        }
169        if name.ends_with("_scale") || name.ends_with("_zero_point") {
170            continue;
171        }
172
173        let original_view = original_tensors
174            .tensor(name)
175            .map_err(|_| Error::TensorNotFound(name.to_string()))?;
176        let original_data = read_f32(&original_view, name)?;
177        let shape = original_view.shape().to_vec();
178
179        let packed = config.teenygrad_packed_int4.get(name);
180        let logical_shape = packed.map(|p| p.logical_shape.clone()).unwrap_or(shape);
181        let granularity = granularity_from_config(weights)?;
182        let (groups, _) = compute_groups(&logical_shape, granularity);
183
184        let quantized_view = quantized_tensors
185            .tensor(name)
186            .map_err(|_| Error::TensorNotFound(name.to_string()))?;
187        let scale_name = format::scale_tensor_name(name);
188        let scale_view = quantized_tensors
189            .tensor(&scale_name)
190            .map_err(|_| Error::TensorNotFound(scale_name.clone()))?;
191        let scales = read_f32(&scale_view, &scale_name)?;
192
193        let zero_points = if weights.symmetric {
194            None
195        } else {
196            let zp_name = format::zero_point_tensor_name(name);
197            let zp_view = quantized_tensors
198                .tensor(&zp_name)
199                .map_err(|_| Error::TensorNotFound(zp_name.clone()))?;
200            Some(
201                zp_view
202                    .data()
203                    .chunks_exact(4)
204                    .map(|b| i32::from_le_bytes(b.try_into().expect("chunks_exact(4)")))
205                    .collect::<Vec<_>>(),
206            )
207        };
208
209        let reconstructed = format::dequantize_tensor(
210            scheme,
211            &quantized_view,
212            &logical_shape,
213            &scales,
214            zero_points.as_deref(),
215            &groups,
216        );
217
218        reports.push(compare_tensors(name, &original_data, &reconstructed)?);
219    }
220
221    Ok(reports)
222}
teeny_quant/validate.rs

teeny_quant/
validate.rs
