foxhunt/ml/tests/test_tft_varmap_quantization.rs

//! Integration test for TFT VarMap bulk quantization
//!
//! Tests the full workflow:
//! 1. Create FP32 TFT model with trained weights
//! 2. Quantize all 3,288 tensors to INT8
//! 3. Save quantized weights to SafeTensors
//! 4. Load quantized weights from disk
//! 5. Verify numerical accuracy (within 1e-2)

use candle_core::{DType, Device, Tensor, Var};
use candle_nn::{VarBuilder, VarMap};
use ml::memory_optimization::quantization::{QuantizationConfig, QuantizationType, Quantizer};
use ml::tft::varmap_quantization::{load_quantized_weights, quantize_varmap, save_quantized_weights};
use std::sync::Arc;

/// Test basic VarMap quantization with small model
#[test]
fn test_quantize_small_varmap() {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());

    // Create a small TFT-like model structure
    {
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);

        // Attention weights (Q, K, V, O)
        let _ = vb.get((256, 256), "attention.q_proj.weight").unwrap();
        let _ = vb.get((256,), "attention.q_proj.bias").unwrap();
        let _ = vb.get((256, 256), "attention.k_proj.weight").unwrap();
        let _ = vb.get((256,), "attention.k_proj.bias").unwrap();
        let _ = vb.get((256, 256), "attention.v_proj.weight").unwrap();
        let _ = vb.get((256,), "attention.v_proj.bias").unwrap();
        let _ = vb.get((256, 256), "attention.o_proj.weight").unwrap();
        let _ = vb.get((256,), "attention.o_proj.bias").unwrap();

        // LSTM weights (input gate)
        let _ = vb.get((256, 256), "lstm.layer0.w_ii").unwrap();
        let _ = vb.get((256, 256), "lstm.layer0.w_hi").unwrap();
    }

    // Quantize VarMap
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device);

    let result = quantize_varmap(varmap, &mut quantizer);
    assert!(result.is_ok(), "Quantization failed: {:?}", result.err());

    let weights = result.unwrap();
    assert_eq!(weights.len(), 10, "Expected 10 quantized tensors");

    // Verify all tensors were quantized
    assert!(weights.contains_key("attention.q_proj.weight"));
    assert!(weights.contains_key("attention.q_proj.bias"));
    assert!(weights.contains_key("lstm.layer0.w_ii"));

    // Verify quantization type
    for (name, qweight) in weights.iter() {
        assert_eq!(
            qweight.quant_type,
            QuantizationType::Int8,
            "Tensor {} has wrong quantization type",
            name
        );
        assert!(
            qweight.scale > 0.0,
            "Tensor {} has invalid scale",
            name
        );
    }
}

/// Test save and load round-trip
#[test]
fn test_save_load_round_trip() {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());

    // Create test tensors with known values
    {
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
        let _ = vb.get((10, 20), "layer1.weight").unwrap();
        let _ = vb.get((10,), "layer1.bias").unwrap();
        let _ = vb.get((20, 30), "layer2.weight").unwrap();
    }

    // Quantize
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());
    let original_weights = quantize_varmap(varmap, &mut quantizer).unwrap();

    // Save to temporary file
    let temp_dir = std::env::temp_dir();
    let temp_path = temp_dir.join("test_tft_quantized_round_trip");
    let temp_path_str = temp_path.to_str().unwrap();

    let save_result = save_quantized_weights(&original_weights, temp_path_str);
    assert!(save_result.is_ok(), "Save failed: {:?}", save_result.err());

    // Load back
    let load_result = load_quantized_weights(temp_path_str, &device);
    assert!(load_result.is_ok(), "Load failed: {:?}", load_result.err());

    let loaded_weights = load_result.unwrap();

    // Verify same number of tensors
    assert_eq!(
        loaded_weights.len(),
        original_weights.len(),
        "Loaded tensor count mismatch"
    );

    // Verify scale and zero_point preserved
    for (name, original) in original_weights.iter() {
        let loaded = loaded_weights
            .get(name)
            .expect(&format!("Missing tensor '{}' after load", name));

        assert!(
            (original.scale - loaded.scale).abs() < 1e-6,
            "Scale mismatch for '{}': original={}, loaded={}",
            name,
            original.scale,
            loaded.scale
        );

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

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

    // Cleanup
    let _ = std::fs::remove_file(format!("{}.safetensors", temp_path_str));
}

/// Test quantization with real-valued tensors (verify numerical accuracy)
#[test]
fn test_quantization_accuracy() {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());

    // Create tensor with known values for accuracy testing
    {
        let tensor_data = vec![
            -2.0f32, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5,
        ];
        let tensor = Tensor::from_vec(tensor_data.clone(), (10,), &device).unwrap();
        let var = Var::from_tensor(&tensor).unwrap();

        varmap
            .data()
            .lock()
            .unwrap()
            .insert("test_tensor".to_string(), var);
    }

    // Quantize
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());
    let weights = quantize_varmap(varmap, &mut quantizer).unwrap();

    // Dequantize and verify accuracy
    let quantized = weights.get("test_tensor").unwrap();
    let dequantized = quantizer.dequantize_tensor(quantized).unwrap();

    // Get original values
    let original_data = vec![-2.0f32, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 2.5];
    let dequantized_data = dequantized.to_vec1::<f32>().unwrap();

    // Verify accuracy within 1e-2 (INT8 quantization error)
    for (i, (&original, &reconstructed)) in original_data
        .iter()
        .zip(dequantized_data.iter())
        .enumerate()
    {
        let error = (original - reconstructed).abs();
        assert!(
            error < 0.05, // 5% error tolerance for INT8
            "Quantization error too large at index {}: original={}, reconstructed={}, error={}",
            i,
            original,
            reconstructed,
            error
        );
    }
}

/// Test handling of invalid tensors (NaN, Inf, empty)
#[test]
fn test_invalid_tensor_handling() {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());

    // Add valid and invalid tensors
    {
        // Valid tensor
        let valid = Tensor::zeros((10, 20), DType::F32, &device).unwrap();
        varmap
            .data()
            .lock()
            .unwrap()
            .insert("valid".to_string(), Var::from_tensor(&valid).unwrap());

        // Empty tensor (0 elements) - should be skipped
        let empty = Tensor::zeros((0,), DType::F32, &device).unwrap();
        varmap
            .data()
            .lock()
            .unwrap()
            .insert("empty".to_string(), Var::from_tensor(&empty).unwrap());

        // Tensor with NaN - should be skipped
        let nan_data = vec![1.0f32, 2.0, f32::NAN, 4.0];
        let nan_tensor = Tensor::from_vec(nan_data, (4,), &device).unwrap();
        varmap.data().lock().unwrap().insert(
            "nan_tensor".to_string(),
            Var::from_tensor(&nan_tensor).unwrap(),
        );

        // Tensor with Inf - should be skipped
        let inf_data = vec![1.0f32, 2.0, f32::INFINITY, 4.0];
        let inf_tensor = Tensor::from_vec(inf_data, (4,), &device).unwrap();
        varmap.data().lock().unwrap().insert(
            "inf_tensor".to_string(),
            Var::from_tensor(&inf_tensor).unwrap(),
        );
    }

    // Quantize (should skip invalid tensors)
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device);
    let weights = quantize_varmap(varmap, &mut quantizer).unwrap();

    // Only the valid tensor should be quantized
    assert_eq!(
        weights.len(),
        1,
        "Should only quantize 1 valid tensor, got {}",
        weights.len()
    );
    assert!(
        weights.contains_key("valid"),
        "Valid tensor should be quantized"
    );
    assert!(
        !weights.contains_key("empty"),
        "Empty tensor should be skipped"
    );
    assert!(
        !weights.contains_key("nan_tensor"),
        "NaN tensor should be skipped"
    );
    assert!(
        !weights.contains_key("inf_tensor"),
        "Inf tensor should be skipped"
    );
}

/// Test performance: quantize 100 tensors and verify <30s target
#[test]
fn test_quantization_performance() {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());

    // Create 100 tensors of varying sizes
    {
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
        for i in 0..100 {
            let size = 100 + i * 10; // Varying sizes: 100, 110, 120, ..., 1090
            let _ = vb
                .get((size, size), &format!("layer_{}.weight", i))
                .unwrap();
            let _ = vb.get((size,), &format!("layer_{}.bias", i)).unwrap();
        }
    }

    // Measure quantization time
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device);

    let start = std::time::Instant::now();
    let weights = quantize_varmap(varmap, &mut quantizer).unwrap();
    let elapsed = start.elapsed();

    // Verify all tensors quantized
    assert_eq!(weights.len(), 200, "Should quantize 200 tensors (100 weights + 100 biases)");

    // Verify performance (should be well under 30s for 200 tensors)
    let elapsed_secs = elapsed.as_secs_f32();
    println!(
        "Quantized {} tensors in {:.2}s ({:.0} tensors/sec)",
        weights.len(),
        elapsed_secs,
        weights.len() as f32 / elapsed_secs
    );

    // Even on slow CPU, 200 tensors should take <10s (conservative target)
    assert!(
        elapsed_secs < 10.0,
        "Quantization too slow: {:.2}s for 200 tensors (expected <10s)",
        elapsed_secs
    );
}

/// Test file size reduction (INT8 should be ~25% of FP32)
#[test]
fn test_file_size_reduction() {
    let device = Device::Cpu;
    let varmap = Arc::new(VarMap::new());

    // Create moderately sized tensors
    {
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
        let _ = vb.get((256, 256), "large_weight").unwrap();
        let _ = vb.get((256,), "large_bias").unwrap();
    }

    // Calculate expected FP32 size
    let fp32_size = (256 * 256 + 256) * 4; // 4 bytes per F32
    let expected_int8_size = (256 * 256 + 256) * 1; // 1 byte per INT8

    // Quantize and save
    let config = QuantizationConfig {
        quant_type: QuantizationType::Int8,
        symmetric: true,
        per_channel: false,
        calibration_samples: None,
    };
    let mut quantizer = Quantizer::new(config, device.clone());
    let weights = quantize_varmap(varmap, &mut quantizer).unwrap();

    let temp_dir = std::env::temp_dir();
    let temp_path = temp_dir.join("test_tft_size_reduction");
    let temp_path_str = temp_path.to_str().unwrap();

    save_quantized_weights(&weights, temp_path_str).unwrap();

    // Check file size
    let safetensors_path = format!("{}.safetensors", temp_path_str);
    let metadata = std::fs::metadata(&safetensors_path).unwrap();
    let actual_size = metadata.len() as usize;

    println!(
        "FP32 size: {} bytes, INT8 size: {} bytes, Expected INT8: {} bytes",
        fp32_size, actual_size, expected_int8_size
    );

    // SafeTensors adds metadata overhead, but size should still be much smaller than FP32
    // Allow up to 50% overhead for metadata (scale/zero_point scalars + SafeTensors headers)
    let max_acceptable_size = (expected_int8_size as f64 * 1.5) as usize;
    assert!(
        actual_size < fp32_size / 2,
        "INT8 file size ({} bytes) should be <50% of FP32 size ({} bytes)",
        actual_size,
        fp32_size
    );

    // Cleanup
    let _ = std::fs::remove_file(safetensors_path);
}