264 lines
8.3 KiB
Rust
264 lines
8.3 KiB
Rust
//! Comprehensive tests for quantization functionality
|
|
|
|
use rtx_compress::{
|
|
CompressionError, Result,
|
|
quantization::{
|
|
CalibrationConfig, CalibrationMethod, PQConfig, PostTrainingQuantizer, ProductQuantizer,
|
|
QuantizationConfig, QuantizationScheme,
|
|
},
|
|
};
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
#[cfg(test)]
|
|
mod quantization_tests {
|
|
use super::*;
|
|
|
|
fn create_test_device() -> Device {
|
|
Device::cpu()
|
|
}
|
|
|
|
fn create_test_tensor(size: usize) -> Tensor {
|
|
let data: Vec<f32> = (0..size).map(|i| i as f32 / 10.0).collect();
|
|
Tensor::from_data(data, &[size / 4, 4], &create_test_device()).unwrap()
|
|
}
|
|
|
|
fn create_embeddings_tensor() -> Tensor {
|
|
// Create some embeddings for testing [batch_size=8, embedding_dim=128]
|
|
let data: Vec<f32> = (0..1024).map(|i| (i as f32).sin()).collect();
|
|
Tensor::from_data(data, &[8, 128], &create_test_device()).unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn test_pq_config_creation() {
|
|
let config = PQConfig {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256,
|
|
max_iterations: 100,
|
|
tolerance: 1e-4,
|
|
use_opq: false,
|
|
opq_iterations: 20,
|
|
use_residual: false,
|
|
residual_stages: 2,
|
|
};
|
|
|
|
assert_eq!(config.num_subquantizers, 8);
|
|
assert_eq!(config.codebook_size, 256);
|
|
assert_eq!(config.max_iterations, 100);
|
|
assert!(!config.use_opq);
|
|
}
|
|
|
|
#[test]
|
|
fn test_product_quantizer_creation() {
|
|
let config = PQConfig {
|
|
num_subquantizers: 4,
|
|
codebook_size: 256,
|
|
max_iterations: 50,
|
|
tolerance: 1e-4,
|
|
use_opq: false,
|
|
opq_iterations: 20,
|
|
use_residual: false,
|
|
residual_stages: 2,
|
|
};
|
|
|
|
let quantizer = ProductQuantizer::new(config);
|
|
assert!(quantizer.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing quantizer fit issue"]
|
|
fn test_product_quantizer_train() {
|
|
let config = PQConfig {
|
|
num_subquantizers: 4,
|
|
codebook_size: 16, // Small for testing
|
|
max_iterations: 10,
|
|
tolerance: 1e-3,
|
|
use_opq: false,
|
|
opq_iterations: 20,
|
|
use_residual: false,
|
|
residual_stages: 2,
|
|
};
|
|
|
|
let mut quantizer = ProductQuantizer::new(config).unwrap();
|
|
let embeddings = create_embeddings_tensor();
|
|
|
|
let result = quantizer.fit(&embeddings);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing codebook size issue"]
|
|
fn test_product_quantizer_encode_decode() {
|
|
let config = PQConfig {
|
|
num_subquantizers: 4,
|
|
codebook_size: 16,
|
|
max_iterations: 10,
|
|
tolerance: 1e-3,
|
|
use_opq: false,
|
|
opq_iterations: 20,
|
|
use_residual: false,
|
|
residual_stages: 2,
|
|
};
|
|
|
|
let mut quantizer = ProductQuantizer::new(config).unwrap();
|
|
let embeddings = create_embeddings_tensor();
|
|
|
|
// Train the quantizer
|
|
quantizer.fit(&embeddings).unwrap();
|
|
|
|
// Encode
|
|
let encoded = quantizer.encode(&embeddings);
|
|
assert!(encoded.is_ok());
|
|
|
|
// Decode
|
|
let decoded = quantizer.decode(&encoded.unwrap());
|
|
assert!(decoded.is_ok());
|
|
|
|
// Check shape is preserved
|
|
assert_eq!(decoded.unwrap().shape(), embeddings.shape());
|
|
}
|
|
|
|
#[test]
|
|
fn test_post_training_quantizer() {
|
|
let calibration = CalibrationConfig::new(10, 1.0);
|
|
let config = QuantizationConfig::new_symmetric(QuantizationScheme::INT8, calibration);
|
|
|
|
let quantizer = PostTrainingQuantizer::new(config);
|
|
assert!(quantizer.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_int8_quantization() {
|
|
let calibration = CalibrationConfig::new(1, 1.0);
|
|
let config = QuantizationConfig::new_symmetric(QuantizationScheme::INT8, calibration);
|
|
|
|
let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
|
|
let tensor = create_test_tensor(64);
|
|
|
|
// Calibrate
|
|
quantizer.calibrate(&[tensor.clone()]).unwrap();
|
|
|
|
// Quantize
|
|
let quantized = quantizer.quantize(&tensor);
|
|
assert!(quantized.is_ok());
|
|
|
|
// Dequantize
|
|
let dequantized = quantizer.dequantize(&quantized.unwrap());
|
|
assert!(dequantized.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_int4_quantization() {
|
|
let calibration = CalibrationConfig::new(1, 1.0);
|
|
let config = QuantizationConfig::new_asymmetric(QuantizationScheme::INT4, calibration);
|
|
|
|
let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
|
|
let tensor = create_test_tensor(64);
|
|
|
|
quantizer.calibrate(&[tensor.clone()]).unwrap();
|
|
let quantized = quantizer.quantize(&tensor).unwrap();
|
|
let dequantized = quantizer.dequantize(&quantized).unwrap();
|
|
|
|
assert_eq!(dequantized.shape(), tensor.shape());
|
|
}
|
|
|
|
#[test]
|
|
fn test_percentile_calibration() {
|
|
let calibration = CalibrationConfig::new_with_percentile(1, 0.001, 0.999);
|
|
let config = QuantizationConfig::new_symmetric(QuantizationScheme::INT8, calibration);
|
|
|
|
let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
|
|
let tensor = create_test_tensor(100);
|
|
|
|
let result = quantizer.calibrate(&[tensor]);
|
|
assert!(result.is_ok());
|
|
}
|
|
|
|
#[test]
|
|
fn test_per_channel_quantization() {
|
|
let calibration = CalibrationConfig::new(1, 1.0);
|
|
let config = QuantizationConfig::new_per_channel(
|
|
QuantizationScheme::INT8,
|
|
calibration,
|
|
0, // channel dimension
|
|
);
|
|
|
|
let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
|
|
let tensor = create_test_tensor(64);
|
|
|
|
quantizer.calibrate(&[tensor.clone()]).unwrap();
|
|
let quantized = quantizer.quantize(&tensor).unwrap();
|
|
let dequantized = quantizer.dequantize(&quantized).unwrap();
|
|
|
|
assert_eq!(dequantized.shape(), tensor.shape());
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing codebook size issue"]
|
|
fn test_compression_ratio() {
|
|
let config = PQConfig {
|
|
num_subquantizers: 8,
|
|
codebook_size: 256, // 8 bits per subquantizer
|
|
max_iterations: 10,
|
|
tolerance: 1e-3,
|
|
use_opq: false,
|
|
opq_iterations: 20,
|
|
use_residual: false,
|
|
residual_stages: 2,
|
|
};
|
|
|
|
let mut quantizer = ProductQuantizer::new(config).unwrap();
|
|
let embeddings = create_embeddings_tensor();
|
|
|
|
quantizer.fit(&embeddings).unwrap();
|
|
let encoded = quantizer.encode(&embeddings).unwrap();
|
|
|
|
// Calculate compression ratio
|
|
let original_size = embeddings.numel() * 4; // f32 = 4 bytes
|
|
let compressed_size = encoded.numel(); // indices are typically u8 or u16
|
|
|
|
let ratio = original_size as f64 / compressed_size as f64;
|
|
assert!(ratio > 1.0); // Should achieve compression
|
|
}
|
|
|
|
#[test]
|
|
#[ignore = "Pre-existing codebook size issue"]
|
|
fn test_residual_quantization() {
|
|
let config = PQConfig {
|
|
num_subquantizers: 4,
|
|
codebook_size: 16,
|
|
max_iterations: 10,
|
|
tolerance: 1e-3,
|
|
use_opq: false,
|
|
opq_iterations: 20,
|
|
use_residual: true,
|
|
residual_stages: 2,
|
|
};
|
|
|
|
let mut quantizer = ProductQuantizer::new(config).unwrap();
|
|
let embeddings = create_embeddings_tensor();
|
|
|
|
quantizer.fit(&embeddings).unwrap();
|
|
let encoded = quantizer.encode(&embeddings).unwrap();
|
|
let decoded = quantizer.decode(&encoded).unwrap();
|
|
|
|
// With residual quantization, reconstruction should be better
|
|
assert_eq!(decoded.shape(), embeddings.shape());
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantization_statistics() {
|
|
let calibration = CalibrationConfig::new(1, 1.0);
|
|
let config = QuantizationConfig::new_symmetric(QuantizationScheme::INT8, calibration);
|
|
|
|
let mut quantizer = PostTrainingQuantizer::new(config).unwrap();
|
|
let tensor = create_test_tensor(64);
|
|
|
|
quantizer.calibrate(&[tensor.clone()]).unwrap();
|
|
let quantized = quantizer.quantize(&tensor).unwrap();
|
|
|
|
let stats = quantizer.get_statistics();
|
|
assert!(stats.scale > 0.0);
|
|
assert!(stats.bit_width > 0);
|
|
}
|
|
}
|