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}