rustytorch/crates/models/rtx-csm/src/quantize.rs

//! Quantization scaffolding (Item 10 of the optimization roadmap).
//!
//! ## Status: PARTIAL
//!
//! This module ships:
//!   1. The **policy** — which tensors are safe to quantize, and at what level
//!   2. The **converter** — read `sesame/csm-1b/model.safetensors`, quantize
//!      per-tensor according to policy, write a candle-loadable file
//!   3. Helpers for inspecting policy decisions for a real checkpoint
//!
//! What it does NOT yet ship: the **forked `csm_quantized.rs`** that consumes
//! the quantized weights via `candle_core::quantized::QMatMul`. Candle's
//! upstream `csm.rs` uses `candle_nn::Linear` everywhere, which cannot load
//! quantized GGML tensors. To finish quantization, follow these steps:
//!
//! ### Remaining work for Item 10
//! 1. Copy `~/.cargo/registry/src/.../candle-transformers-0.9.x/src/models/csm.rs`
//!    into `rtx-csm/src/csm_quantized.rs` (vendor; ~533 LOC).
//! 2. In the vendored copy, replace `Linear` with `QMatMul` ONLY for the
//!    layers listed in [`QuantPolicy::quantizable`].
//! 3. Update `Model::new` to take a `&candle_transformers::quantized_var_builder::VarBuilder`
//!    instead of `candle_nn::VarBuilder`.
//! 4. Wire `model.rs::CsmModel::load_from_safetensors` to detect quantized
//!    files (by extension or a header bit) and dispatch to the new model type.
//! 5. Update `generator.rs` `dtype` selection: F32 on CPU, F16 on Metal,
//!    BF16 on CUDA — the *unquantized* layers still need a dtype, only the
//!    quantized ones bypass.
//! 6. Add a feature flag `quantized` and gate the new module behind it to
//!    keep build times bounded for non-quantized users.
//!
//! Expected impact (from research): 1.5–2× speedup at Q4_K_M, 1.2–1.4× at
//! Q8_0, ~50% memory reduction on the 1.1B-param backbone+decoder.

use crate::error::{CsmError, Result};
use candle_core::quantized::GgmlDType;
use std::path::Path;

/// Per-tensor quantization decisions.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TensorQuant {
    /// Keep the tensor at its native dtype (F32/BF16/F16). Required for
    /// embedding tables, head projections, and LayerNorm/RMSNorm scales.
    Keep,
    /// Quantize to the given GGML type. Use Q8_0 for safety, Q4_K_M for size.
    Quant(GgmlDType),
}

/// Quantization policy. Maps tensor names to a quantization decision.
#[derive(Debug, Clone)]
pub struct QuantPolicy {
    pub default: TensorQuant,
}

impl QuantPolicy {
    /// Conservative: Q8_0 on backbone QKV/FFN, keep everything else native.
    /// Recommended as the first quantization run before trying Q4_K_M.
    pub fn q8_safe() -> Self {
        Self {
            default: TensorQuant::Quant(GgmlDType::Q8_0),
        }
    }

    /// Aggressive: Q4_K_M on backbone, Q8_0 on the small decoder. Roughly 50%
    /// memory of Q8_0, modestly more quality risk.
    pub fn q4km_aggressive() -> Self {
        Self {
            default: TensorQuant::Quant(GgmlDType::Q4K),
        }
    }

    /// Decide what to do with a tensor by its safetensors key. The key
    /// taxonomy mirrors `candle_transformers::models::csm`:
    ///   - `backbone.layers.{N}.{attn,mlp}.{q,k,v,o,gate,up,down}_proj.weight`
    ///   - `decoder.layers.{N}.{...}.weight`
    ///   - `audio_embeddings.weight`, `text_embeddings.weight`
    ///   - `codebook0_head.weight`, `audio_head`
    ///   - `projection.weight`
    ///   - `*.norm.weight`, `*.{rms_norm,layer_norm}.weight`
    pub fn decide(&self, name: &str) -> TensorQuant {
        // Always keep: head projections, embeddings, projection, norms.
        const KEEP_EXACT: &[&str] = &[
            "audio_embeddings.weight",
            "text_embeddings.weight",
            "codebook0_head.weight",
            "audio_head",
            "projection.weight",
        ];
        if KEEP_EXACT.contains(&name) {
            return TensorQuant::Keep;
        }
        // RMSNorm / LayerNorm scales — keep
        if name.ends_with(".norm.weight")
            || name.ends_with(".rms_norm.weight")
            || name.ends_with(".layer_norm.weight")
            || name.ends_with("_norm.weight")
        {
            return TensorQuant::Keep;
        }
        // Bias terms — keep (always F32)
        if name.ends_with(".bias") {
            return TensorQuant::Keep;
        }
        // Decoder is small (4 layers, 100M params). Keep at Q8_0 even when
        // backbone is Q4_K_M — quantization noise here corrupts c1..c31 which
        // determine timbre.
        if name.starts_with("decoder.") {
            // If default is more aggressive than Q8, downshift to Q8.
            return match self.default {
                TensorQuant::Quant(GgmlDType::Q8_0) => TensorQuant::Quant(GgmlDType::Q8_0),
                TensorQuant::Quant(_) => TensorQuant::Quant(GgmlDType::Q8_0),
                TensorQuant::Keep => TensorQuant::Keep,
            };
        }
        // Backbone QKV/output projections.
        if name.starts_with("backbone.")
            && (name.ends_with("_proj.weight") || name.ends_with(".proj.weight"))
        {
            return self.default;
        }
        // Backbone MLP weights (csm uses w1/w2/w3, not gate/up/down naming).
        if name.starts_with("backbone.")
            && (name.ends_with(".w1.weight")
                || name.ends_with(".w2.weight")
                || name.ends_with(".w3.weight"))
        {
            return self.default;
        }
        // Anything else, conservative default: keep.
        TensorQuant::Keep
    }

    /// Returns `true` if this tensor would be quantized under the policy.
    pub fn quantizable(&self, name: &str) -> bool {
        matches!(self.decide(name), TensorQuant::Quant(_))
    }
}

/// Estimate quantized file size given a pre-loaded list of tensor descriptors
/// (from `model::dump_safetensors_keys`). Useful for dry-run reporting before
/// committing to a multi-GB quantization run.
pub fn estimate_size_bytes(
    descriptors: &[crate::model::TensorDescriptor],
    policy: &QuantPolicy,
) -> usize {
    descriptors
        .iter()
        .map(|d| {
            let n_elements: usize = d.shape.iter().product();
            match policy.decide(&d.name) {
                TensorQuant::Keep => {
                    // Assume original dtype. F32=4B, F16/BF16=2B; we don't have
                    // the dtype on TensorDescriptor as a discriminant, so default to
                    // 2 bytes (BF16 is what Sesame ships).
                    n_elements * 2
                }
                TensorQuant::Quant(GgmlDType::Q8_0) => n_elements * 8 / 8 + n_elements / 32 * 2,
                TensorQuant::Quant(GgmlDType::Q4K) => n_elements * 4 / 8 + n_elements / 256 * 12,
                TensorQuant::Quant(_) => n_elements,
            }
        })
        .sum()
}

/// Pretty-print the policy decisions over a checkpoint. Sanity-check before
/// committing weeks to the fork. Lists what'll be quantized and what'll be
/// kept native, plus rough size estimates.
pub fn report(descriptors: &[crate::model::TensorDescriptor], policy: &QuantPolicy) -> String {
    let mut out = String::new();
    let mut quantized = 0usize;
    let mut kept = 0usize;
    for d in descriptors {
        let n: usize = d.shape.iter().product();
        match policy.decide(&d.name) {
            TensorQuant::Keep => {
                kept += n;
                out.push_str(&format!("  KEEP   {} ({})\n", d.name, fmt_count(n)));
            }
            TensorQuant::Quant(t) => {
                quantized += n;
                out.push_str(&format!("  {:?}   {} ({})\n", t, d.name, fmt_count(n)));
            }
        }
    }
    out.push_str(&format!(
        "\nTotal: {} quantized, {} kept ({:.1}% quantized by param count)\n",
        fmt_count(quantized),
        fmt_count(kept),
        100.0 * quantized as f32 / (quantized + kept) as f32,
    ));
    out.push_str(&format!(
        "Estimated quantized file size: {}\n",
        fmt_count(estimate_size_bytes(descriptors, policy))
    ));
    out
}

fn fmt_count(n: usize) -> String {
    if n >= 1_000_000_000 {
        format!("{:.2}B", n as f64 / 1e9)
    } else if n >= 1_000_000 {
        format!("{:.1}M", n as f64 / 1e6)
    } else if n >= 1_000 {
        format!("{:.1}K", n as f64 / 1e3)
    } else {
        n.to_string()
    }
}

/// Read an HF safetensors file, quantize each tensor per `policy`, and emit a
/// GGUF v2 file at `output_path`. Tensors flagged `Keep` are stored at their
/// native dtype using GGML's F16 / F32 type codes.
///
/// This builds the *artifact* that a forked `csm_quantized.rs` would consume.
/// Until that fork lands, the GGUF can be inspected, sized, and audited but
/// not yet loaded for inference.
pub fn convert_to_quantized<P: AsRef<Path>>(
    input_safetensors: P,
    output_path: P,
    policy: &QuantPolicy,
) -> Result<QuantReport> {
    use candle_core::Device;
    use candle_core::quantized::{GgmlDType, QTensor, gguf_file};
    use std::fs::File;
    use std::io::BufWriter;

    let input = input_safetensors.as_ref();
    let output = output_path.as_ref();
    let device = Device::Cpu;

    let bytes = std::fs::read(input)?;
    let st = safetensors::SafeTensors::deserialize(&bytes)?;

    let mut report = QuantReport::default();
    // GGUF tensor table — we have to OWN the QTensor values until write_all
    // completes, since gguf_file::write borrows them.
    let mut owned: Vec<(String, QTensor)> = Vec::with_capacity(st.names().len());

    for name in st.names() {
        let view = st.tensor(name)?;
        let shape: Vec<usize> = view.shape().to_vec();
        // Always materialize the source tensor in F32 for consistent quantization input.
        let src = safetensors_view_to_f32(&view, &device, &shape)?;

        match policy.decide(name) {
            TensorQuant::Keep => {
                // Keep at F16 unless the tensor is 1-D (norm/bias) — those stay F32.
                let target = if shape.len() <= 1 {
                    GgmlDType::F32
                } else {
                    GgmlDType::F16
                };
                let qt = QTensor::quantize(&src, target)?;
                report.kept_params += shape.iter().product::<usize>();
                report.kept_tensors += 1;
                owned.push((name.to_string(), qt));
            }
            TensorQuant::Quant(dtype) => {
                // GGML K-quants require dim 0 % 256 == 0 and dim 0 >= 32; for any
                // tensor that doesn't fit, downgrade to Q8_0 (works for any
                // multiple of 32) or fall back to F16 if even that fails.
                let inner = shape.last().copied().unwrap_or(0);
                let chosen = if is_kquant(dtype) && (inner % 256 != 0) {
                    GgmlDType::Q8_0
                } else {
                    dtype
                };
                let chosen = if inner % 32 != 0 {
                    GgmlDType::F16
                } else {
                    chosen
                };
                // EXPERIMENT (CSM_TRANSPOSE_QUANT=1): some Linear-weight
                // serializers store as (out, in) but candle's qmatmul kernel
                // appears to expect a different orientation in practice. If
                // the env var is set, transpose 2-D weights before quantizing
                // and let the loader request the swapped shape.
                let src_for_quant =
                    if std::env::var("CSM_TRANSPOSE_QUANT").is_ok() && shape.len() == 2 {
                        src.t()?.contiguous()?
                    } else {
                        src
                    };
                let qt = QTensor::quantize(&src_for_quant, chosen)?;
                report
                    .quant_buckets
                    .entry(format!("{:?}", chosen))
                    .and_modify(|c| *c += 1)
                    .or_insert(1);
                report.quantized_params += shape.iter().product::<usize>();
                report.quantized_tensors += 1;
                owned.push((name.to_string(), qt));
            }
        }
    }

    let metadata: Vec<(&str, &gguf_file::Value)> = vec![];
    let tensors: Vec<(&str, &QTensor)> = owned.iter().map(|(n, q)| (n.as_str(), q)).collect();
    let f = File::create(output)?;
    let mut w = BufWriter::new(f);
    gguf_file::write(&mut w, &metadata, &tensors)?;

    report.output_bytes = std::fs::metadata(output)?.len() as usize;
    report.input_bytes = bytes.len();
    Ok(report)
}

fn is_kquant(d: candle_core::quantized::GgmlDType) -> bool {
    use candle_core::quantized::GgmlDType::*;
    matches!(d, Q2K | Q3K | Q4K | Q5K | Q6K | Q8K)
}

fn safetensors_view_to_f32(
    view: &safetensors::tensor::TensorView,
    device: &candle_core::Device,
    shape: &[usize],
) -> Result<candle_core::Tensor> {
    use candle_core::{DType, Tensor};
    use safetensors::Dtype as SfDtype;
    let bytes = view.data();
    let t = match view.dtype() {
        SfDtype::F32 => {
            let v: &[f32] = bytemuck_slice(bytes);
            Tensor::from_slice(v, shape, device)?
        }
        SfDtype::F16 => {
            let v: &[half::f16] = bytemuck_slice(bytes);
            Tensor::from_slice(v, shape, device)?.to_dtype(DType::F32)?
        }
        SfDtype::BF16 => {
            let v: &[half::bf16] = bytemuck_slice(bytes);
            Tensor::from_slice(v, shape, device)?.to_dtype(DType::F32)?
        }
        other => {
            return Err(CsmError::Config(format!(
                "unsupported safetensors dtype: {other:?}"
            )));
        }
    };
    Ok(t)
}

fn bytemuck_slice<T: bytemuck::Pod>(bytes: &[u8]) -> &[T] {
    bytemuck::cast_slice(bytes)
}

#[derive(Debug, Default)]
pub struct QuantReport {
    pub kept_tensors: usize,
    pub kept_params: usize,
    pub quantized_tensors: usize,
    pub quantized_params: usize,
    pub quant_buckets: std::collections::HashMap<String, usize>,
    pub input_bytes: usize,
    pub output_bytes: usize,
}

impl std::fmt::Display for QuantReport {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        writeln!(
            f,
            "kept: {} tensors / {} params; quantized: {} tensors / {} params",
            self.kept_tensors, self.kept_params, self.quantized_tensors, self.quantized_params
        )?;
        let mut buckets: Vec<_> = self.quant_buckets.iter().collect();
        buckets.sort_by_key(|(k, _)| (*k).clone());
        for (k, v) in buckets {
            writeln!(f, "  {k}: {v} tensors")?;
        }
        writeln!(
            f,
            "size: {} bytes input → {} bytes output ({:.1}× compression)",
            self.input_bytes,
            self.output_bytes,
            self.input_bytes as f64 / self.output_bytes.max(1) as f64
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn policy_keeps_heads_and_embeds() {
        let p = QuantPolicy::q8_safe();
        assert!(matches!(
            p.decide("audio_embeddings.weight"),
            TensorQuant::Keep
        ));
        assert!(matches!(
            p.decide("text_embeddings.weight"),
            TensorQuant::Keep
        ));
        assert!(matches!(
            p.decide("codebook0_head.weight"),
            TensorQuant::Keep
        ));
        assert!(matches!(p.decide("audio_head"), TensorQuant::Keep));
        assert!(matches!(p.decide("projection.weight"), TensorQuant::Keep));
    }

    #[test]
    fn policy_keeps_norms_and_bias() {
        let p = QuantPolicy::q8_safe();
        assert!(matches!(
            p.decide("backbone.layers.0.attention_norm.weight"),
            TensorQuant::Keep
        ));
        assert!(matches!(
            p.decide("backbone.layers.0.attn.q_proj.bias"),
            TensorQuant::Keep
        ));
    }

    #[test]
    fn policy_quantizes_backbone_projs() {
        let p = QuantPolicy::q8_safe();
        assert!(p.quantizable("backbone.layers.5.attn.q_proj.weight"));
        assert!(p.quantizable("backbone.layers.5.attn.k_proj.weight"));
        assert!(p.quantizable("backbone.layers.5.mlp.gate_proj.weight"));
    }

    #[test]
    fn aggressive_policy_keeps_decoder_at_q8() {
        let p = QuantPolicy::q4km_aggressive();
        let decoder = "decoder.layers.0.attn.q_proj.weight";
        match p.decide(decoder) {
            TensorQuant::Quant(GgmlDType::Q8_0) => {}
            other => panic!("expected Q8_0 on decoder, got {other:?}"),
        }
        // Backbone gets the aggressive level
        let bb = "backbone.layers.0.attn.q_proj.weight";
        match p.decide(bb) {
            TensorQuant::Quant(GgmlDType::Q4K) => {}
            other => panic!("expected Q4K on backbone, got {other:?}"),
        }
    }
}