foxhunt/ml/examples/quantize_tft_varmap.rs

//! Example: Quantize TFT VarMap to INT8
//!
//! Demonstrates the full workflow:
//! 1. Load FP32 TFT model weights from safetensors
//! 2. Quantize all 3,288 tensors to INT8
//! 3. Save quantized weights
//! 4. Load and verify quantized weights
//!
//! Usage:
//!   cargo run --example quantize_tft_varmap --release --features cuda -- \
//!     --input ml/trained_models/tft_225_epoch_0.safetensors \
//!     --output ml/trained_models/tft_225_epoch_0_int8

use candle_core::Device;
use candle_nn::VarMap;
use clap::Parser;
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::tft::varmap_quantization::{
    load_quantized_weights, quantize_varmap, save_quantized_weights,
};
use ml::MLError;
use std::sync::Arc;

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Path to FP32 TFT model (safetensors)
    #[arg(short, long)]
    input: String,

    /// Path to save quantized INT8 model
    #[arg(short, long)]
    output: String,

    /// Use CPU instead of CUDA
    #[arg(long)]
    cpu: bool,
}

fn main() -> Result<(), MLError> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_target(false)
        .with_thread_ids(true)
        .with_level(true)
        .init();

    let args = Args::parse();

    // Select device
    let device = if args.cpu {
        Device::Cpu
    } else {
        Device::cuda_if_available(0).unwrap_or(Device::Cpu)
    };

    println!("Using device: {:?}", device);
    println!("Input: {}", args.input);
    println!("Output: {}", args.output);
    println!();

    // Step 1: Load FP32 model weights into VarMap
    println!("Step 1: Loading FP32 model weights from {}", args.input);
    let varmap = Arc::new(VarMap::new());

    // Load tensors from safetensors
    let tensors = candle_core::safetensors::load(&args.input, &device)
        .map_err(|e| MLError::CheckpointError(format!("Failed to load FP32 model: {}", e)))?;

    // Populate VarMap with loaded tensors
    {
        use candle_core::Var;
        let mut vars_data = varmap
            .data()
            .lock()
            .map_err(|e| MLError::LockError(format!("Failed to lock VarMap: {}", e)))?;

        for (name, tensor) in tensors.iter() {
            let var = Var::from_tensor(tensor)?;
            vars_data.insert(name.clone(), var);
        }
    }

    println!("✓ Loaded {} tensors from FP32 model", tensors.len());
    println!();

    // Step 2: Quantize all tensors to INT8
    println!("Step 2: Quantizing VarMap to INT8");
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());

    let quantized_weights = quantize_varmap(varmap.clone(), &mut quantizer)?;
    println!("✓ Quantized {} tensors", quantized_weights.len());
    println!();

    // Step 3: Save quantized weights
    println!("Step 3: Saving quantized weights to {}", args.output);
    save_quantized_weights(&quantized_weights, &args.output)?;
    println!();

    // Step 4: Verify by loading back
    println!("Step 4: Verifying quantized weights (load and compare)");
    let loaded_weights = load_quantized_weights(&args.output, &device)?;

    // Verify same number of tensors
    assert_eq!(
        loaded_weights.len(),
        quantized_weights.len(),
        "Tensor count mismatch after load"
    );

    // Verify scale and zero_point preserved
    let mut total_error = 0.0;
    let mut max_error = 0.0;

    for (name, original) in quantized_weights.iter() {
        let loaded = loaded_weights
            .get(name)
            .expect(&format!("Missing tensor '{}' after load", name));

        // Check scale
        let scale_error = (original.scale - loaded.scale).abs();
        total_error += scale_error;
        max_error = max_error.max(scale_error);

        // Check zero_point
        assert_eq!(
            original.zero_point, loaded.zero_point,
            "Zero point mismatch for '{}'",
            name
        );

        // Check shape
        assert_eq!(
            original.data.dims(),
            loaded.data.dims(),
            "Shape mismatch for '{}'",
            name
        );
    }

    let avg_error = total_error / quantized_weights.len() as f32;

    println!("✓ Verification passed:");
    println!("  - Tensor count: {} tensors", loaded_weights.len());
    println!("  - Avg scale error: {:.2e}", avg_error);
    println!("  - Max scale error: {:.2e}", max_error);
    println!();

    // Calculate memory savings
    let fp32_size_mb = tensors.len() * 4 / 1024 / 1024; // Rough estimate
    let metadata = std::fs::metadata(format!("{}.safetensors", args.output))
        .map_err(|e| MLError::CheckpointError(format!("Failed to stat output file: {}", e)))?;
    let int8_size_mb = metadata.len() as f32 / 1024.0 / 1024.0;

    println!("Memory Savings:");
    println!("  - FP32 model: ~{} MB (estimated)", fp32_size_mb);
    println!("  - INT8 model: {:.2} MB (actual)", int8_size_mb);
    println!(
        "  - Reduction: ~{:.1}%",
        (1.0 - int8_size_mb / fp32_size_mb as f32) * 100.0
    );
    println!();

    println!("✓ VarMap quantization complete!");
    println!("  Output: {}.safetensors", args.output);

    Ok(())
}