Files
rustytorch/crates/production/rtx-inference/tests/quantization.rs
T
2026-03-04 00:08:42 +00:00

577 lines
20 KiB
Rust

//! Comprehensive tests for quantization functionality
//!
//! Tests cover:
//! - INT8, INT4, and FP8 quantization schemes
//! - Dynamic quantization with calibration
//! - Accuracy validation and loss measurement
//! - Performance optimization verification
use rtx_inference::DataType;
use rtx_inference::quantization::TestTensor;
/// Test quantization configuration
#[tokio::test]
async fn test_quantization_config_creation() {
use rtx_inference::quantization::{CalibrationMethod, QuantizationConfig, QuantizationScheme};
let config = QuantizationConfig::new()
.with_scheme(QuantizationScheme::INT8)
.with_calibration_method(CalibrationMethod::MinMax)
.with_calibration_samples(1000)
.with_accuracy_threshold(0.02)
.with_dynamic_range_optimization(true);
assert_eq!(config.scheme, QuantizationScheme::INT8);
assert_eq!(config.calibration_samples, 1000);
assert_eq!(config.accuracy_threshold, 0.02);
assert!(config.dynamic_range_optimization);
// Test validation
let invalid_config = QuantizationConfig::new().with_calibration_samples(0); // Should be > 0
assert!(invalid_config.validate().is_err());
}
/// Test INT8 quantization with symmetric scaling
#[tokio::test]
async fn test_int8_symmetric_quantization() {
use rtx_inference::quantization::{
QuantizationConfig, QuantizationScheme, Quantizer, ScalingMethod,
};
let config = QuantizationConfig::new()
.with_scheme(QuantizationScheme::INT8)
.with_scaling_method(ScalingMethod::Symmetric);
let quantizer = Quantizer::new(config);
// Test data with known range
let input = TestTensor::new(vec![-2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0, 5.0], vec![8]);
let quantized = quantizer.quantize(&input).await.unwrap();
assert_eq!(quantized.dtype, DataType::INT8);
// Test dequantization
let dequantized = quantizer.dequantize(&quantized).await.unwrap();
assert_eq!(dequantized.dtype, DataType::F32);
// Verify accuracy (should be close to original within quantization error)
let max_error = input
.data
.iter()
.zip(dequantized.data.iter())
.map(|(orig, deq)| (orig - deq).abs())
.fold(0.0f32, |a, b| a.max(b));
assert!(
max_error < 0.1,
"Quantization error too large: {}",
max_error
);
}
/// Test INT8 quantization with asymmetric scaling
#[tokio::test]
async fn test_int8_asymmetric_quantization() {
use rtx_inference::quantization::{
QuantizationConfig, QuantizationScheme, Quantizer, ScalingMethod,
};
let config = QuantizationConfig::new()
.with_scheme(QuantizationScheme::INT8)
.with_scaling_method(ScalingMethod::Asymmetric);
let quantizer = Quantizer::new(config);
// Test data with asymmetric range
let input = TestTensor::new(vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0], vec![8]);
let quantized = quantizer.quantize(&input).await.unwrap();
let dequantized = quantizer.dequantize(&quantized).await.unwrap();
// Asymmetric quantization should handle this range better
let max_error = input
.data
.iter()
.zip(dequantized.data.iter())
.map(|(orig, deq)| (orig - deq).abs())
.fold(0.0f32, |a, b| a.max(b));
assert!(
max_error < 0.05,
"Asymmetric quantization error: {}",
max_error
);
}
/// Test INT4 quantization with weight sharing
#[tokio::test]
async fn test_int4_quantization() {
use rtx_inference::quantization::{
GroupSize, QuantizationConfig, QuantizationScheme, Quantizer,
};
let config = QuantizationConfig::new()
.with_scheme(QuantizationScheme::INT4)
.with_group_size(GroupSize::G128); // Group quantization for better accuracy
let quantizer = Quantizer::new(config);
// Create larger tensor for group quantization
let input = TestTensor::from_random(vec![512, 256], -1.0, 1.0);
let quantized = quantizer.quantize(&input).await.unwrap();
assert_eq!(quantized.dtype, DataType::INT4);
// Verify storage efficiency (INT4 should use ~1/8 the space of F32)
let compression_ratio =
quantized.storage_size_bytes() as f32 / input.storage_size_bytes() as f32;
assert!(
compression_ratio < 0.15,
"INT4 compression ratio: {}",
compression_ratio
);
let dequantized = quantizer.dequantize(&quantized).await.unwrap();
// Measure quantization quality
let snr = calculate_snr(&input.data, &dequantized.data);
assert!(snr > 20.0, "Signal-to-noise ratio too low: {} dB", snr);
}
/// Test FP8 E4M3 quantization
#[tokio::test]
async fn test_fp8_e4m3_quantization() {
use rtx_inference::quantization::{QuantizationConfig, QuantizationScheme, Quantizer};
let config = QuantizationConfig::new().with_scheme(QuantizationScheme::FP8E4M3);
let quantizer = Quantizer::new(config);
// Test with moderate dynamic range (good for E4M3 format)
let input = TestTensor::from_random(vec![128, 128], -8.0, 8.0);
let quantized = quantizer.quantize(&input).await.unwrap();
assert_eq!(quantized.dtype, DataType::FP8E4M3);
let dequantized = quantizer.dequantize(&quantized).await.unwrap();
// FP8 should maintain reasonable accuracy for this range
let mse = calculate_mse(&input.data, &dequantized.data);
assert!(mse < 0.1, "FP8 E4M3 MSE too high: {}", mse);
}
/// Test FP8 E5M2 quantization
#[tokio::test]
async fn test_fp8_e5m2_quantization() {
use rtx_inference::quantization::{QuantizationConfig, QuantizationScheme, Quantizer};
let config = QuantizationConfig::new().with_scheme(QuantizationScheme::FP8E5M2);
let quantizer = Quantizer::new(config);
// Test with larger dynamic range (better for E5M2 format)
let input = TestTensor::from_random(vec![128, 128], -32.0, 32.0);
let quantized = quantizer.quantize(&input).await.unwrap();
assert_eq!(quantized.dtype, DataType::FP8E5M2);
let dequantized = quantizer.dequantize(&quantized).await.unwrap();
// E5M2 should handle larger ranges better than E4M3
let relative_error = calculate_relative_error(&input.data, &dequantized.data);
assert!(
relative_error < 0.05,
"FP8 E5M2 relative error: {}",
relative_error
);
}
/// Test dynamic quantization with calibration
#[tokio::test]
async fn test_dynamic_quantization_calibration() {
use rtx_inference::quantization::{CalibrationMethod, DynamicQuantizer, QuantizationConfig};
let config = QuantizationConfig::new()
.with_calibration_method(CalibrationMethod::KLDivergence)
.with_calibration_samples(500);
let mut quantizer = DynamicQuantizer::new(config);
// Collect calibration samples
let mut calibration_data = Vec::new();
for _ in 0..500 {
let sample = TestTensor::from_random(vec![64, 64], -10.0, 10.0);
calibration_data.push(sample);
}
// Perform calibration
quantizer.calibrate(&calibration_data).await.unwrap();
// Test quantization with calibrated parameters
let test_input = TestTensor::from_random(vec![64, 64], -8.0, 12.0);
let quantized = quantizer.quantize(&test_input).await.unwrap();
let dequantized = quantizer.dequantize(&quantized).await.unwrap();
// Calibration should improve accuracy
let calibrated_error = calculate_mse(&test_input.data, &dequantized.data);
// Compare with uncalibrated quantization
let uncalibrated_config = QuantizationConfig::new();
let uncalibrated_quantizer = rtx_inference::quantization::Quantizer::new(uncalibrated_config);
let uncalibrated_quantized = uncalibrated_quantizer.quantize(&test_input).await.unwrap();
let uncalibrated_dequantized = uncalibrated_quantizer
.dequantize(&uncalibrated_quantized)
.await
.unwrap();
let uncalibrated_error = calculate_mse(&test_input.data, &uncalibrated_dequantized.data);
assert!(
calibrated_error < uncalibrated_error,
"Calibrated error ({}) should be less than uncalibrated error ({})",
calibrated_error,
uncalibrated_error
);
}
/// Test accuracy validation with thresholds
#[tokio::test]
async fn test_accuracy_validation() {
use rtx_inference::quantization::{
QuantizationConfig, QuantizationValidator, ValidationMetrics,
};
let config = QuantizationConfig::new().with_accuracy_threshold(0.01);
let validator = QuantizationValidator::new(config);
// Create high-quality quantization (should pass)
let original = TestTensor::from_random(vec![100, 100], -1.0, 1.0);
let mut quantized_good = original.clone();
// Add small noise to simulate good quantization
for val in quantized_good.data.iter_mut() {
*val += (fastrand::f32() - 0.5) * 0.005; // Small noise
}
let metrics_good = validator
.validate(&original, &quantized_good)
.await
.unwrap();
assert!(
metrics_good.passes_threshold(),
"Good quantization should pass validation"
);
assert!(
metrics_good.mse < 0.01,
"MSE should be low: {}",
metrics_good.mse
);
// Create poor quantization (should fail)
let mut quantized_bad = original.clone();
// Add large noise to simulate poor quantization
for val in quantized_bad.data.iter_mut() {
*val += (fastrand::f32() - 0.5) * 0.5; // Large noise
}
let metrics_bad = validator.validate(&original, &quantized_bad).await.unwrap();
assert!(
!metrics_bad.passes_threshold(),
"Bad quantization should fail validation"
);
assert!(
metrics_bad.mse > 0.01,
"MSE should be high: {}",
metrics_bad.mse
);
}
/// Test per-channel vs per-tensor quantization
#[tokio::test]
async fn test_per_channel_vs_per_tensor() {
use rtx_inference::quantization::{QuantizationConfig, QuantizationGranularity, Quantizer};
let input = TestTensor::new(
vec![
// Channel 1: small values
0.1, 0.2, 0.15, 0.25, // Channel 2: large values
10.0, 15.0, 12.0, 18.0, // Channel 3: mixed values
-5.0, 2.0, -1.0, 3.0,
],
vec![3, 4], // 3 channels, 4 elements each
);
// Test per-tensor quantization
let per_tensor_config =
QuantizationConfig::new().with_granularity(QuantizationGranularity::PerTensor);
let per_tensor_quantizer = Quantizer::new(per_tensor_config);
let per_tensor_result = per_tensor_quantizer.quantize(&input).await.unwrap();
let per_tensor_dequantized = per_tensor_quantizer
.dequantize(&per_tensor_result)
.await
.unwrap();
// Test per-channel quantization
let per_channel_config =
QuantizationConfig::new().with_granularity(QuantizationGranularity::PerChannel);
let per_channel_quantizer = Quantizer::new(per_channel_config);
let per_channel_result = per_channel_quantizer.quantize(&input).await.unwrap();
let per_channel_dequantized = per_channel_quantizer
.dequantize(&per_channel_result)
.await
.unwrap();
// Per-channel should be more accurate for this case
let per_tensor_error = calculate_mse(&input.data, &per_tensor_dequantized.data);
let per_channel_error = calculate_mse(&input.data, &per_channel_dequantized.data);
assert!(
per_channel_error < per_tensor_error,
"Per-channel error ({}) should be less than per-tensor error ({})",
per_channel_error,
per_tensor_error
);
}
/// Test quantization with outlier handling
#[tokio::test]
async fn test_outlier_handling() {
use rtx_inference::quantization::{OutlierHandling, QuantizationConfig, Quantizer};
let mut data = vec![1.0; 1000]; // Mostly normal values
data[500] = 100.0; // Large outlier
data[501] = -100.0; // Large outlier
let input = TestTensor::new(data, vec![1000]);
// Test without outlier handling
let no_outlier_config = QuantizationConfig::new().with_outlier_handling(OutlierHandling::None);
let no_outlier_quantizer = Quantizer::new(no_outlier_config);
let no_outlier_result = no_outlier_quantizer.quantize(&input).await.unwrap();
let no_outlier_dequantized = no_outlier_quantizer
.dequantize(&no_outlier_result)
.await
.unwrap();
// Test with outlier clipping
let clipping_config = QuantizationConfig::new()
.with_outlier_handling(OutlierHandling::Clipping { percentile: 99.0 });
let clipping_quantizer = Quantizer::new(clipping_config);
let clipping_result = clipping_quantizer.quantize(&input).await.unwrap();
let clipping_dequantized = clipping_quantizer
.dequantize(&clipping_result)
.await
.unwrap();
// Outlier handling should improve quantization of normal values
let normal_indices: Vec<usize> = (0..1000).filter(|&i| i != 500 && i != 501).collect();
let no_outlier_error =
calculate_mse_subset(&input.data, &no_outlier_dequantized.data, &normal_indices);
let clipping_error =
calculate_mse_subset(&input.data, &clipping_dequantized.data, &normal_indices);
assert!(
clipping_error < no_outlier_error,
"Outlier handling should improve normal value accuracy: {} vs {}",
clipping_error,
no_outlier_error
);
}
/// Test quantization performance optimization
#[tokio::test]
async fn test_quantization_performance() {
use rtx_inference::quantization::{OptimizationLevel, QuantizationConfig, Quantizer};
use std::time::Instant;
let large_input = TestTensor::from_random(vec![1024, 1024], -1.0, 1.0);
// Test basic quantization
let basic_config = QuantizationConfig::new().with_optimization_level(OptimizationLevel::None);
let basic_quantizer = Quantizer::new(basic_config);
let start = Instant::now();
let _basic_result = basic_quantizer.quantize(&large_input).await.unwrap();
let basic_time = start.elapsed();
// Test optimized quantization
let optimized_config =
QuantizationConfig::new().with_optimization_level(OptimizationLevel::Aggressive);
let optimized_quantizer = Quantizer::new(optimized_config);
let start = Instant::now();
let _optimized_result = optimized_quantizer.quantize(&large_input).await.unwrap();
let optimized_time = start.elapsed();
// Optimized should be faster (though this test might be flaky in CI)
println!(
"Basic time: {:?}, Optimized time: {:?}",
basic_time, optimized_time
);
// At least verify both complete successfully
assert!(basic_time > std::time::Duration::ZERO);
assert!(optimized_time > std::time::Duration::ZERO);
}
/// Test mixed precision quantization
#[tokio::test]
async fn test_mixed_precision_quantization() {
use rtx_inference::quantization::{LayerConfig, MixedPrecisionQuantizer, QuantizationScheme};
let mut quantizer = MixedPrecisionQuantizer::new();
// Configure different layers with different precisions
quantizer.add_layer_config("embedding", LayerConfig::new(QuantizationScheme::INT8));
quantizer.add_layer_config("attention", LayerConfig::new(QuantizationScheme::FP8E4M3));
quantizer.add_layer_config("ffn", LayerConfig::new(QuantizationScheme::INT4));
quantizer.add_layer_config("output", LayerConfig::new(QuantizationScheme::FP8E5M2));
// Test each layer type
let test_cases = vec![
(
"embedding",
TestTensor::from_random(vec![512, 768], -0.1, 0.1),
),
(
"attention",
TestTensor::from_random(vec![64, 64], -2.0, 2.0),
),
("ffn", TestTensor::from_random(vec![768, 3072], -1.0, 1.0)),
(
"output",
TestTensor::from_random(vec![3072, 50257], -0.5, 0.5),
),
];
for (layer_name, input) in test_cases {
let quantized = quantizer.quantize_layer(layer_name, &input).await.unwrap();
let dequantized = quantizer
.dequantize_layer(layer_name, &quantized)
.await
.unwrap();
// Verify that appropriate precision was used
let expected_dtype = match layer_name {
"embedding" => DataType::INT8,
"attention" => DataType::FP8E4M3,
"ffn" => DataType::INT4,
"output" => DataType::FP8E5M2,
_ => panic!("Unexpected layer name"),
};
assert_eq!(quantized.dtype, expected_dtype);
// Verify reasonable accuracy
let error = calculate_mse(&input.data, &dequantized.data);
assert!(
error < 0.5,
"Layer {} error too high: {}",
layer_name,
error
);
}
}
/// Test error handling in quantization
#[tokio::test]
async fn test_quantization_error_handling() {
use rtx_inference::quantization::{QuantizationConfig, QuantizationScheme, Quantizer};
let quantizer = Quantizer::new(QuantizationConfig::new());
// Test empty tensor
let empty_tensor = TestTensor::new(vec![], vec![0]);
let result = quantizer.quantize(&empty_tensor).await;
assert!(result.is_err(), "Should fail on empty tensor");
// Test invalid shape
let invalid_tensor = TestTensor::new(vec![1.0, 2.0], vec![3]); // Data/shape mismatch
let result = quantizer.quantize(&invalid_tensor).await;
assert!(result.is_err(), "Should fail on shape mismatch");
// Test extreme values
let extreme_tensor = TestTensor::new(vec![f32::INFINITY, f32::NEG_INFINITY, f32::NAN], vec![3]);
let result = quantizer.quantize(&extreme_tensor).await;
assert!(result.is_err(), "Should fail on extreme values");
}
/// Test quantization with batch processing
#[tokio::test]
async fn test_batch_quantization() {
use rtx_inference::quantization::{BatchQuantizer, QuantizationConfig};
let config = QuantizationConfig::new();
let quantizer = BatchQuantizer::new(config);
// Create batch of tensors
let batch = vec![
TestTensor::from_random(vec![128, 256], -1.0, 1.0),
TestTensor::from_random(vec![128, 256], -2.0, 2.0),
TestTensor::from_random(vec![128, 256], -0.5, 0.5),
TestTensor::from_random(vec![128, 256], -3.0, 3.0),
];
let quantized_batch = quantizer.quantize_batch(&batch).await.unwrap();
assert_eq!(quantized_batch.len(), batch.len());
let dequantized_batch = quantizer.dequantize_batch(&quantized_batch).await.unwrap();
assert_eq!(dequantized_batch.len(), batch.len());
// Verify accuracy for each tensor in batch
for (i, (original, dequantized)) in batch.iter().zip(dequantized_batch.iter()).enumerate() {
let error = calculate_mse(&original.data, &dequantized.data);
assert!(error < 0.1, "Batch item {} error too high: {}", i, error);
}
}
// Helper functions
fn calculate_snr(original: &[f32], quantized: &[f32]) -> f32 {
let signal_power: f32 = original.iter().map(|x| x * x).sum();
let noise_power: f32 = original
.iter()
.zip(quantized.iter())
.map(|(o, q)| (o - q) * (o - q))
.sum();
if noise_power == 0.0 {
f32::INFINITY
} else {
10.0 * (signal_power / noise_power).log10()
}
}
fn calculate_mse(original: &[f32], quantized: &[f32]) -> f32 {
assert_eq!(original.len(), quantized.len());
let sum_squared_error: f32 = original
.iter()
.zip(quantized.iter())
.map(|(o, q)| (o - q) * (o - q))
.sum();
sum_squared_error / original.len() as f32
}
fn calculate_relative_error(original: &[f32], quantized: &[f32]) -> f32 {
assert_eq!(original.len(), quantized.len());
let sum_relative_error: f32 = original
.iter()
.zip(quantized.iter())
.map(|(o, q)| {
if o.abs() < 1e-8 {
if q.abs() < 1e-8 { 0.0 } else { 1.0 }
} else {
((o - q) / o).abs()
}
})
.sum();
sum_relative_error / original.len() as f32
}
fn calculate_mse_subset(original: &[f32], quantized: &[f32], indices: &[usize]) -> f32 {
let sum_squared_error: f32 = indices
.iter()
.map(|&i| (original[i] - quantized[i]) * (original[i] - quantized[i]))
.sum();
sum_squared_error / indices.len() as f32
}