Skip to main content

teeny_quant/cli/
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//! `teeny-quant validate`: per-tensor quantization error metrics (see [`crate::validate`]).
18
19use std::path::PathBuf;
20
21use anyhow::{Context, Result};
22use clap::Args;
23
24use crate::validate::validate_checkpoint;
25
26/// `teeny-quant validate` arguments.
27#[derive(Args, Debug)]
28pub struct ValidateArgs {
29    /// The original (unquantized) `.safetensors` checkpoint.
30    #[arg(long)]
31    pub original: PathBuf,
32
33    /// The quantized `.safetensors` checkpoint (produced by `teeny-quant quantize`).
34    #[arg(long)]
35    pub quantized: PathBuf,
36
37    /// Exit non-zero if any tensor's max absolute error exceeds this.
38    #[arg(long)]
39    pub max_abs_error_threshold: Option<f32>,
40}
41
42/// Runs `teeny-quant validate`.
43pub fn run(args: ValidateArgs) -> Result<()> {
44    let reports = validate_checkpoint(&args.original, &args.quantized).with_context(|| {
45        format!(
46            "failed to validate '{}' against '{}'",
47            args.quantized.display(),
48            args.original.display()
49        )
50    })?;
51
52    println!(
53        "{:<55} {:>14} {:>14} {:>10}",
54        "tensor", "max_abs_err", "mean_abs_err", "sqnr_db"
55    );
56    let mut flagged = 0usize;
57    for r in &reports {
58        println!(
59            "{:<55} {:>14.6} {:>14.6} {:>10.2}",
60            r.name, r.max_abs_error, r.mean_abs_error, r.sqnr_db
61        );
62        if args
63            .max_abs_error_threshold
64            .is_some_and(|t| r.max_abs_error > t)
65        {
66            flagged += 1;
67        }
68    }
69    println!("\n{} tensor(s) checked", reports.len());
70
71    if let Some(threshold) = args.max_abs_error_threshold
72        && flagged > 0
73    {
74        anyhow::bail!("{flagged} tensor(s) exceeded max-abs-error threshold {threshold}");
75    }
76
77    Ok(())
78}