teeny_quant/cli/
quantize.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 quantize`: reads a `.safetensors` checkpoint and writes a quantized one.
18
19use std::collections::HashMap;
20use std::path::PathBuf;
21
22use anyhow::{Context, Result};
23use clap::Args;
24
25use crate::format::{self, PackedTensorInfo};
26use crate::quant::{Fp8Variant, Granularity, Scheme};
27use crate::read::read_f32;
28use crate::write::{OutputTensor, write_safetensors};
29
30/// `teeny-quant quantize` arguments.
31#[derive(Args, Debug)]
32pub struct QuantizeArgs {
33    /// Input `.safetensors` checkpoint.
34    #[arg(long)]
35    pub input: PathBuf,
36
37    /// Output `.safetensors` path.
38    #[arg(long)]
39    pub output: PathBuf,
40
41    /// Quantization scheme.
42    #[arg(long, value_enum)]
43    pub scheme: SchemeArg,
44
45    /// Use asymmetric (rather than symmetric) affine quantization. Ignored for `--scheme fp8`,
46    /// which is always amax-scaled/symmetric.
47    #[arg(long)]
48    pub asymmetric: bool,
49
50    /// Quantization granularity.
51    #[arg(long, value_enum, default_value = "channel")]
52    pub granularity: GranularityArg,
53
54    /// Axis for `--granularity channel`/`group` (default `0`, the output-channel axis for
55    /// `Linear`/`Conv` weights shaped `[out, in, ...]`). Group-wise quantization typically wants
56    /// the reduction axis instead, e.g. `--axis 1`.
57    #[arg(long, default_value_t = 0)]
58    pub axis: usize,
59
60    /// Elements per group; required for `--granularity group`.
61    #[arg(long)]
62    pub group_size: Option<usize>,
63
64    /// FP8 encoding; only used for `--scheme fp8`.
65    #[arg(long, value_enum, default_value = "e4m3")]
66    pub fp8_variant: Fp8VariantArg,
67}
68
69/// `--scheme` values.
70#[derive(clap::ValueEnum, Clone, Copy, Debug)]
71pub enum SchemeArg {
72    /// `Scheme::Int8`.
73    #[value(name = "int8")]
74    Int8,
75    /// `Scheme::Int4`.
76    #[value(name = "int4")]
77    Int4,
78    /// `Scheme::Fp8`.
79    #[value(name = "fp8")]
80    Fp8,
81}
82
83/// `--granularity` values.
84#[derive(clap::ValueEnum, Clone, Copy, Debug)]
85pub enum GranularityArg {
86    /// `Granularity::PerTensor`.
87    #[value(name = "tensor")]
88    Tensor,
89    /// `Granularity::PerChannel`.
90    #[value(name = "channel")]
91    Channel,
92    /// `Granularity::Group`.
93    #[value(name = "group")]
94    Group,
95}
96
97/// `--fp8-variant` values.
98#[derive(clap::ValueEnum, Clone, Copy, Debug)]
99pub enum Fp8VariantArg {
100    /// `Fp8Variant::E4M3`.
101    #[value(name = "e4m3")]
102    E4M3,
103    /// `Fp8Variant::E5M2`.
104    #[value(name = "e5m2")]
105    E5M2,
106}
107
108/// Runs `teeny-quant quantize`.
109pub fn run(args: QuantizeArgs) -> Result<()> {
110    let scheme = match args.scheme {
111        SchemeArg::Int8 => Scheme::Int8 {
112            symmetric: !args.asymmetric,
113        },
114        SchemeArg::Int4 => Scheme::Int4 {
115            symmetric: !args.asymmetric,
116        },
117        SchemeArg::Fp8 => Scheme::Fp8 {
118            variant: match args.fp8_variant {
119                Fp8VariantArg::E4M3 => Fp8Variant::E4M3,
120                Fp8VariantArg::E5M2 => Fp8Variant::E5M2,
121            },
122        },
123    };
124
125    let granularity = match args.granularity {
126        GranularityArg::Tensor => Granularity::PerTensor,
127        GranularityArg::Channel => Granularity::PerChannel { axis: args.axis },
128        GranularityArg::Group => {
129            let group_size = args
130                .group_size
131                .context("--group-size is required for --granularity group")?;
132            Granularity::Group {
133                axis: args.axis,
134                group_size,
135            }
136        }
137    };
138
139    let mapped = teeny_data::safetensors::SafeTensors::from_pretrained(&args.input)
140        .with_context(|| format!("failed to open '{}'", args.input.display()))?;
141    let tensors = mapped.tensors().with_context(|| {
142        format!(
143            "failed to read tensor headers from '{}'",
144            args.input.display()
145        )
146    })?;
147
148    let mut outputs: HashMap<String, OutputTensor> = HashMap::new();
149    let mut ignored = Vec::new();
150    let mut packed_int4 = HashMap::new();
151    let mut quantized_count = 0usize;
152
153    for name in tensors.names() {
154        let view = tensors
155            .tensor(name)
156            .with_context(|| format!("reading tensor '{name}'"))?;
157        let shape = view.shape().to_vec();
158
159        if !format::should_quantize(&shape, view.dtype()) {
160            ignored.push(name.to_string());
161            outputs.insert(
162                name.to_string(),
163                OutputTensor {
164                    dtype: view.dtype(),
165                    shape,
166                    data: view.data().to_vec(),
167                },
168            );
169            continue;
170        }
171
172        let data =
173            read_f32(&view, name).with_context(|| format!("reading tensor '{name}' as f32"))?;
174        let quantized = format::quantize_tensor(name, &data, &shape, scheme, granularity)
175            .with_context(|| format!("quantizing tensor '{name}'"))?;
176
177        if matches!(scheme, Scheme::Int4 { .. }) {
178            packed_int4.insert(
179                name.to_string(),
180                PackedTensorInfo {
181                    logical_shape: shape.clone(),
182                    elements: shape.iter().product(),
183                },
184            );
185        }
186
187        outputs.extend(quantized);
188        quantized_count += 1;
189    }
190
191    let ignored_count = ignored.len();
192    let config = format::build_config(scheme, granularity, ignored, packed_int4);
193    let metadata = format::config_to_metadata(&config)?;
194
195    write_safetensors(&args.output, &outputs, metadata)
196        .with_context(|| format!("failed to write '{}'", args.output.display()))?;
197
198    println!(
199        "Quantized {quantized_count} tensor(s) ({scheme}, {granularity}), left {ignored_count} tensor(s) unquantized -> {output}",
200        scheme = scheme.short_name(),
201        granularity = granularity.strategy_name(),
202        output = args.output.display(),
203    );
204
205    Ok(())
206}
teeny_quant/cli/quantize.rs

teeny_quant/cli/
quantize.rs
