1074 lines
33 KiB
Rust
1074 lines
33 KiB
Rust
//! Quantization Module for Model Compression and Acceleration
|
|
//!
|
|
//! This module implements various quantization schemes (INT8, INT4, FP8) with
|
|
//! calibration methods to maintain model accuracy while reducing memory usage
|
|
//! and improving inference performance.
|
|
//!
|
|
//! Key features:
|
|
//! - Multiple quantization schemes: INT8, INT4, FP8 E4M3, FP8 E5M2
|
|
//! - Dynamic quantization with calibration
|
|
//! - Per-tensor and per-channel granularity
|
|
//! - Outlier handling and accuracy validation
|
|
//! - Mixed precision support for different layers
|
|
|
|
use crate::{InferenceError, InferenceResult};
|
|
use std::collections::HashMap;
|
|
use std::sync::Arc;
|
|
use tokio::sync::Mutex;
|
|
use tracing::{debug, info};
|
|
|
|
/// Data types supported by quantization
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum DataType {
|
|
F32,
|
|
F16,
|
|
BF16,
|
|
INT8,
|
|
INT4,
|
|
FP8E4M3,
|
|
FP8E5M2,
|
|
}
|
|
|
|
/// Quantization schemes available
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum QuantizationScheme {
|
|
INT8,
|
|
INT4,
|
|
FP8E4M3,
|
|
FP8E5M2,
|
|
}
|
|
|
|
/// Scaling methods for quantization
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum ScalingMethod {
|
|
Symmetric,
|
|
Asymmetric,
|
|
}
|
|
|
|
/// Calibration methods for dynamic quantization
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum CalibrationMethod {
|
|
MinMax,
|
|
Percentile,
|
|
KLDivergence,
|
|
MSE,
|
|
}
|
|
|
|
/// Quantization granularity
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum QuantizationGranularity {
|
|
PerTensor,
|
|
PerChannel,
|
|
PerGroup,
|
|
}
|
|
|
|
/// Group size for grouped quantization
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum GroupSize {
|
|
G32,
|
|
G64,
|
|
G128,
|
|
G256,
|
|
}
|
|
|
|
impl GroupSize {
|
|
#[must_use]
|
|
pub fn size(&self) -> usize {
|
|
match self {
|
|
Self::G32 => 32,
|
|
Self::G64 => 64,
|
|
Self::G128 => 128,
|
|
Self::G256 => 256,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Outlier handling strategies
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub enum OutlierHandling {
|
|
None,
|
|
Clipping { percentile: f32 },
|
|
Separate { threshold: f32 },
|
|
}
|
|
|
|
/// Optimization levels for quantization performance
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub enum OptimizationLevel {
|
|
None,
|
|
Basic,
|
|
Aggressive,
|
|
}
|
|
|
|
/// Configuration for quantization operations
|
|
#[derive(Debug, Clone)]
|
|
pub struct QuantizationConfig {
|
|
pub scheme: QuantizationScheme,
|
|
pub scaling_method: ScalingMethod,
|
|
pub calibration_method: CalibrationMethod,
|
|
pub calibration_samples: usize,
|
|
pub accuracy_threshold: f32,
|
|
pub granularity: QuantizationGranularity,
|
|
pub group_size: Option<GroupSize>,
|
|
pub dynamic_range_optimization: bool,
|
|
pub outlier_handling: OutlierHandling,
|
|
pub optimization_level: OptimizationLevel,
|
|
}
|
|
|
|
impl QuantizationConfig {
|
|
/// Create a new configuration with defaults
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self {
|
|
scheme: QuantizationScheme::INT8,
|
|
scaling_method: ScalingMethod::Symmetric,
|
|
calibration_method: CalibrationMethod::MinMax,
|
|
calibration_samples: 512,
|
|
accuracy_threshold: 0.01,
|
|
granularity: QuantizationGranularity::PerTensor,
|
|
group_size: None,
|
|
dynamic_range_optimization: false,
|
|
outlier_handling: OutlierHandling::None,
|
|
optimization_level: OptimizationLevel::Basic,
|
|
}
|
|
}
|
|
|
|
/// Set quantization scheme
|
|
#[must_use]
|
|
pub fn with_scheme(mut self, scheme: QuantizationScheme) -> Self {
|
|
self.scheme = scheme;
|
|
self
|
|
}
|
|
|
|
/// Set scaling method
|
|
#[must_use]
|
|
pub fn with_scaling_method(mut self, method: ScalingMethod) -> Self {
|
|
self.scaling_method = method;
|
|
self
|
|
}
|
|
|
|
/// Set calibration method
|
|
#[must_use]
|
|
pub fn with_calibration_method(mut self, method: CalibrationMethod) -> Self {
|
|
self.calibration_method = method;
|
|
self
|
|
}
|
|
|
|
/// Set number of calibration samples
|
|
#[must_use]
|
|
pub fn with_calibration_samples(mut self, samples: usize) -> Self {
|
|
self.calibration_samples = samples;
|
|
self
|
|
}
|
|
|
|
/// Set accuracy threshold for validation
|
|
#[must_use]
|
|
pub fn with_accuracy_threshold(mut self, threshold: f32) -> Self {
|
|
self.accuracy_threshold = threshold;
|
|
self
|
|
}
|
|
|
|
/// Set quantization granularity
|
|
#[must_use]
|
|
pub fn with_granularity(mut self, granularity: QuantizationGranularity) -> Self {
|
|
self.granularity = granularity;
|
|
self
|
|
}
|
|
|
|
/// Set group size for grouped quantization
|
|
#[must_use]
|
|
pub fn with_group_size(mut self, size: GroupSize) -> Self {
|
|
self.group_size = Some(size);
|
|
self
|
|
}
|
|
|
|
/// Enable dynamic range optimization
|
|
#[must_use]
|
|
pub fn with_dynamic_range_optimization(mut self, enabled: bool) -> Self {
|
|
self.dynamic_range_optimization = enabled;
|
|
self
|
|
}
|
|
|
|
/// Set outlier handling method
|
|
#[must_use]
|
|
pub fn with_outlier_handling(mut self, handling: OutlierHandling) -> Self {
|
|
self.outlier_handling = handling;
|
|
self
|
|
}
|
|
|
|
/// Set optimization level
|
|
#[must_use]
|
|
pub fn with_optimization_level(mut self, level: OptimizationLevel) -> Self {
|
|
self.optimization_level = level;
|
|
self
|
|
}
|
|
|
|
/// Validate configuration parameters
|
|
pub fn validate(&self) -> InferenceResult<()> {
|
|
if self.calibration_samples == 0 {
|
|
return Err(InferenceError::invalid_request(
|
|
"calibration_samples must be greater than 0",
|
|
));
|
|
}
|
|
|
|
if !(0.0..=1.0).contains(&self.accuracy_threshold) {
|
|
return Err(InferenceError::invalid_request(
|
|
"accuracy_threshold must be between 0.0 and 1.0",
|
|
));
|
|
}
|
|
|
|
if matches!(self.granularity, QuantizationGranularity::PerGroup)
|
|
&& self.group_size.is_none()
|
|
{
|
|
return Err(InferenceError::invalid_request(
|
|
"group_size must be specified for per-group quantization",
|
|
));
|
|
}
|
|
|
|
if let OutlierHandling::Clipping { percentile } = self.outlier_handling
|
|
&& !(0.0..=100.0).contains(&percentile)
|
|
{
|
|
return Err(InferenceError::invalid_request(
|
|
"outlier clipping percentile must be between 0.0 and 100.0",
|
|
));
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Default for QuantizationConfig {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
/// Test tensor representation for quantization operations
|
|
#[derive(Debug, Clone)]
|
|
pub struct TestTensor {
|
|
pub data: Vec<f32>,
|
|
pub shape: Vec<usize>,
|
|
pub dtype: DataType,
|
|
}
|
|
|
|
impl TestTensor {
|
|
#[must_use]
|
|
pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
|
|
Self {
|
|
data,
|
|
shape,
|
|
dtype: DataType::F32,
|
|
}
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn with_dtype(mut self, dtype: DataType) -> Self {
|
|
self.dtype = dtype;
|
|
self
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn num_elements(&self) -> usize {
|
|
self.shape.iter().product()
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn from_random(shape: Vec<usize>, min: f32, max: f32) -> Self {
|
|
let num_elements = shape.iter().product();
|
|
let mut data = Vec::with_capacity(num_elements);
|
|
for _ in 0..num_elements {
|
|
data.push(min + fastrand::f32() * (max - min));
|
|
}
|
|
Self::new(data, shape)
|
|
}
|
|
|
|
#[must_use]
|
|
pub fn storage_size_bytes(&self) -> usize {
|
|
match self.dtype {
|
|
DataType::F32 => self.data.len() * 4,
|
|
DataType::F16 => self.data.len() * 2,
|
|
DataType::BF16 => self.data.len() * 2,
|
|
DataType::INT8 => self.data.len(),
|
|
DataType::INT4 => self.data.len().div_ceil(2), // 4 bits per element
|
|
DataType::FP8E4M3 => self.data.len(),
|
|
DataType::FP8E5M2 => self.data.len(),
|
|
}
|
|
}
|
|
|
|
fn validate(&self) -> InferenceResult<()> {
|
|
if self.data.is_empty() {
|
|
return Err(InferenceError::invalid_request("Tensor data is empty"));
|
|
}
|
|
|
|
let expected_elements = self.shape.iter().product::<usize>();
|
|
if self.data.len() != expected_elements {
|
|
return Err(InferenceError::invalid_request(format!(
|
|
"Data length ({}) doesn't match shape ({:?})",
|
|
self.data.len(),
|
|
self.shape
|
|
)));
|
|
}
|
|
|
|
// Check for invalid values
|
|
for (i, &value) in self.data.iter().enumerate() {
|
|
if !value.is_finite() {
|
|
return Err(InferenceError::invalid_request(format!(
|
|
"Non-finite value {value} at index {i}"
|
|
)));
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// Quantization parameters computed during calibration
|
|
#[derive(Debug, Clone)]
|
|
pub struct QuantizationParams {
|
|
pub scale: f32,
|
|
pub zero_point: i32,
|
|
pub min_value: f32,
|
|
pub max_value: f32,
|
|
pub scheme: QuantizationScheme,
|
|
}
|
|
|
|
impl QuantizationParams {
|
|
#[must_use]
|
|
pub fn new(scheme: QuantizationScheme) -> Self {
|
|
Self {
|
|
scale: 1.0,
|
|
zero_point: 0,
|
|
min_value: 0.0,
|
|
max_value: 0.0,
|
|
scheme,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Main quantization implementation
|
|
pub struct Quantizer {
|
|
config: QuantizationConfig,
|
|
params: Option<QuantizationParams>,
|
|
}
|
|
|
|
impl Quantizer {
|
|
/// Create a new quantizer
|
|
#[must_use]
|
|
pub fn new(config: QuantizationConfig) -> Self {
|
|
Self {
|
|
config,
|
|
params: None,
|
|
}
|
|
}
|
|
|
|
/// Quantize a tensor
|
|
pub async fn quantize(&self, tensor: &TestTensor) -> InferenceResult<TestTensor> {
|
|
tensor.validate()?;
|
|
|
|
debug!(
|
|
"Quantizing tensor with shape {:?} using {:?}",
|
|
tensor.shape, self.config.scheme
|
|
);
|
|
|
|
// Apply outlier handling if configured
|
|
let processed_data = self.handle_outliers(&tensor.data).await?;
|
|
|
|
// Compute quantization parameters
|
|
let params = self.compute_quantization_params(&processed_data).await?;
|
|
|
|
// Perform quantization based on scheme
|
|
let quantized_data = match self.config.scheme {
|
|
QuantizationScheme::INT8 => self.quantize_int8(&processed_data, ¶ms).await?,
|
|
QuantizationScheme::INT4 => self.quantize_int4(&processed_data, ¶ms).await?,
|
|
QuantizationScheme::FP8E4M3 => self.quantize_fp8_e4m3(&processed_data, ¶ms).await?,
|
|
QuantizationScheme::FP8E5M2 => self.quantize_fp8_e5m2(&processed_data, ¶ms).await?,
|
|
};
|
|
|
|
let target_dtype = match self.config.scheme {
|
|
QuantizationScheme::INT8 => DataType::INT8,
|
|
QuantizationScheme::INT4 => DataType::INT4,
|
|
QuantizationScheme::FP8E4M3 => DataType::FP8E4M3,
|
|
QuantizationScheme::FP8E5M2 => DataType::FP8E5M2,
|
|
};
|
|
|
|
Ok(TestTensor {
|
|
data: quantized_data,
|
|
shape: tensor.shape.clone(),
|
|
dtype: target_dtype,
|
|
})
|
|
}
|
|
|
|
/// Dequantize a tensor
|
|
pub async fn dequantize(&self, tensor: &TestTensor) -> InferenceResult<TestTensor> {
|
|
tensor.validate()?;
|
|
|
|
debug!(
|
|
"Dequantizing tensor with shape {:?} from {:?}",
|
|
tensor.shape, tensor.dtype
|
|
);
|
|
|
|
// Use stored parameters or compute them
|
|
let params = if let Some(ref params) = self.params {
|
|
params.clone()
|
|
} else {
|
|
// For dequantization, we need to infer parameters from the quantized data
|
|
self.infer_quantization_params(tensor).await?
|
|
};
|
|
|
|
// Perform dequantization based on original scheme
|
|
let dequantized_data = match tensor.dtype {
|
|
DataType::INT8 => self.dequantize_int8(&tensor.data, ¶ms).await?,
|
|
DataType::INT4 => self.dequantize_int4(&tensor.data, ¶ms).await?,
|
|
DataType::FP8E4M3 => self.dequantize_fp8_e4m3(&tensor.data, ¶ms).await?,
|
|
DataType::FP8E5M2 => self.dequantize_fp8_e5m2(&tensor.data, ¶ms).await?,
|
|
_ => {
|
|
return Err(InferenceError::invalid_request(format!(
|
|
"Cannot dequantize tensor with dtype {:?}",
|
|
tensor.dtype
|
|
)));
|
|
}
|
|
};
|
|
|
|
Ok(TestTensor {
|
|
data: dequantized_data,
|
|
shape: tensor.shape.clone(),
|
|
dtype: DataType::F32,
|
|
})
|
|
}
|
|
|
|
// Implementation methods
|
|
|
|
async fn handle_outliers(&self, data: &[f32]) -> InferenceResult<Vec<f32>> {
|
|
match &self.config.outlier_handling {
|
|
OutlierHandling::None => Ok(data.to_vec()),
|
|
OutlierHandling::Clipping { percentile } => {
|
|
let mut sorted_data = data.to_vec();
|
|
sorted_data.sort_by(f32::total_cmp);
|
|
|
|
let lower_idx = ((100.0 - percentile) / 200.0 * sorted_data.len() as f32) as usize;
|
|
let upper_idx = (percentile / 100.0 * sorted_data.len() as f32) as usize;
|
|
|
|
let min_val = sorted_data[lower_idx];
|
|
let max_val = sorted_data[upper_idx.min(sorted_data.len() - 1)];
|
|
|
|
Ok(data.iter().map(|&x| x.clamp(min_val, max_val)).collect())
|
|
}
|
|
OutlierHandling::Separate { threshold: _ } => {
|
|
// For now, implement as clipping at 3 standard deviations
|
|
let mean = data.iter().sum::<f32>() / data.len() as f32;
|
|
let variance =
|
|
data.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / data.len() as f32;
|
|
let std_dev = variance.sqrt();
|
|
|
|
let min_val = mean - 3.0 * std_dev;
|
|
let max_val = mean + 3.0 * std_dev;
|
|
|
|
Ok(data.iter().map(|&x| x.clamp(min_val, max_val)).collect())
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn compute_quantization_params(
|
|
&self,
|
|
data: &[f32],
|
|
) -> InferenceResult<QuantizationParams> {
|
|
let mut params = QuantizationParams::new(self.config.scheme);
|
|
|
|
// Find min and max values
|
|
let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
|
|
let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
|
|
|
params.min_value = min_val;
|
|
params.max_value = max_val;
|
|
|
|
// Compute scale and zero point based on scaling method and scheme
|
|
match self.config.scaling_method {
|
|
ScalingMethod::Symmetric => {
|
|
let abs_max = max_val.abs().max(min_val.abs());
|
|
params.scale = match self.config.scheme {
|
|
QuantizationScheme::INT8 => abs_max / 127.0,
|
|
QuantizationScheme::INT4 => abs_max / 7.0,
|
|
QuantizationScheme::FP8E4M3 | QuantizationScheme::FP8E5M2 => abs_max / 240.0, // Approximate
|
|
};
|
|
params.zero_point = 0;
|
|
}
|
|
ScalingMethod::Asymmetric => {
|
|
let (qmin, qmax) = match self.config.scheme {
|
|
QuantizationScheme::INT8 => (-128i32, 127i32),
|
|
QuantizationScheme::INT4 => (-8i32, 7i32),
|
|
QuantizationScheme::FP8E4M3 | QuantizationScheme::FP8E5M2 => (-240i32, 239i32), // Approximate
|
|
};
|
|
|
|
params.scale = (max_val - min_val) / (qmax - qmin) as f32;
|
|
params.zero_point = qmin - (min_val / params.scale).round() as i32;
|
|
}
|
|
}
|
|
|
|
// Ensure scale is not zero
|
|
if params.scale == 0.0 {
|
|
params.scale = 1e-8;
|
|
}
|
|
|
|
Ok(params)
|
|
}
|
|
|
|
async fn infer_quantization_params(
|
|
&self,
|
|
tensor: &TestTensor,
|
|
) -> InferenceResult<QuantizationParams> {
|
|
// For dequantization, we need to infer reasonable parameters
|
|
// This is a simplified approach - in practice, these would be stored with the quantized tensor
|
|
let mut params = QuantizationParams::new(match tensor.dtype {
|
|
DataType::INT8 => QuantizationScheme::INT8,
|
|
DataType::INT4 => QuantizationScheme::INT4,
|
|
DataType::FP8E4M3 => QuantizationScheme::FP8E4M3,
|
|
DataType::FP8E5M2 => QuantizationScheme::FP8E5M2,
|
|
_ => {
|
|
return Err(InferenceError::internal_error(
|
|
"quantization",
|
|
"Invalid dtype for inference",
|
|
));
|
|
}
|
|
});
|
|
|
|
// Use default scale factors for inference
|
|
params.scale = match tensor.dtype {
|
|
DataType::INT8 => 1.0 / 127.0,
|
|
DataType::INT4 => 1.0 / 7.0,
|
|
DataType::FP8E4M3 | DataType::FP8E5M2 => 1.0 / 240.0,
|
|
_ => 1.0,
|
|
};
|
|
params.zero_point = 0;
|
|
|
|
Ok(params)
|
|
}
|
|
|
|
async fn quantize_int8(
|
|
&self,
|
|
data: &[f32],
|
|
params: &QuantizationParams,
|
|
) -> InferenceResult<Vec<f32>> {
|
|
let quantized: Vec<f32> = data
|
|
.iter()
|
|
.map(|&x| {
|
|
let q_val = (x / params.scale).round() + params.zero_point as f32;
|
|
q_val.clamp(-128.0, 127.0)
|
|
})
|
|
.collect();
|
|
Ok(quantized)
|
|
}
|
|
|
|
async fn dequantize_int8(
|
|
&self,
|
|
data: &[f32],
|
|
params: &QuantizationParams,
|
|
) -> InferenceResult<Vec<f32>> {
|
|
let dequantized: Vec<f32> = data
|
|
.iter()
|
|
.map(|&q| (q - params.zero_point as f32) * params.scale)
|
|
.collect();
|
|
Ok(dequantized)
|
|
}
|
|
|
|
async fn quantize_int4(
|
|
&self,
|
|
data: &[f32],
|
|
params: &QuantizationParams,
|
|
) -> InferenceResult<Vec<f32>> {
|
|
let quantized: Vec<f32> = data
|
|
.iter()
|
|
.map(|&x| {
|
|
let q_val = (x / params.scale).round() + params.zero_point as f32;
|
|
q_val.clamp(-8.0, 7.0)
|
|
})
|
|
.collect();
|
|
Ok(quantized)
|
|
}
|
|
|
|
async fn dequantize_int4(
|
|
&self,
|
|
data: &[f32],
|
|
params: &QuantizationParams,
|
|
) -> InferenceResult<Vec<f32>> {
|
|
let dequantized: Vec<f32> = data
|
|
.iter()
|
|
.map(|&q| (q - params.zero_point as f32) * params.scale)
|
|
.collect();
|
|
Ok(dequantized)
|
|
}
|
|
|
|
async fn quantize_fp8_e4m3(
|
|
&self,
|
|
data: &[f32],
|
|
params: &QuantizationParams,
|
|
) -> InferenceResult<Vec<f32>> {
|
|
// Simplified FP8 E4M3 quantization (4-bit exponent, 3-bit mantissa)
|
|
let quantized: Vec<f32> = data
|
|
.iter()
|
|
.map(|&x| {
|
|
let scaled = x / params.scale;
|
|
// Simulate FP8 precision by rounding to representable values
|
|
let rounded = (scaled * 8.0).round() / 8.0; // Coarse approximation
|
|
rounded.clamp(-240.0, 239.0)
|
|
})
|
|
.collect();
|
|
Ok(quantized)
|
|
}
|
|
|
|
async fn dequantize_fp8_e4m3(
|
|
&self,
|
|
data: &[f32],
|
|
params: &QuantizationParams,
|
|
) -> InferenceResult<Vec<f32>> {
|
|
let dequantized: Vec<f32> = data.iter().map(|&q| q * params.scale).collect();
|
|
Ok(dequantized)
|
|
}
|
|
|
|
async fn quantize_fp8_e5m2(
|
|
&self,
|
|
data: &[f32],
|
|
params: &QuantizationParams,
|
|
) -> InferenceResult<Vec<f32>> {
|
|
// Simplified FP8 E5M2 quantization (5-bit exponent, 2-bit mantissa)
|
|
let quantized: Vec<f32> = data
|
|
.iter()
|
|
.map(|&x| {
|
|
let scaled = x / params.scale;
|
|
// Simulate FP8 precision with different rounding for E5M2
|
|
let rounded = (scaled * 4.0).round() / 4.0; // Coarse approximation
|
|
rounded.clamp(-57344.0, 57344.0) // Approximate E5M2 range
|
|
})
|
|
.collect();
|
|
Ok(quantized)
|
|
}
|
|
|
|
async fn dequantize_fp8_e5m2(
|
|
&self,
|
|
data: &[f32],
|
|
params: &QuantizationParams,
|
|
) -> InferenceResult<Vec<f32>> {
|
|
let dequantized: Vec<f32> = data.iter().map(|&q| q * params.scale).collect();
|
|
Ok(dequantized)
|
|
}
|
|
}
|
|
|
|
/// Dynamic quantizer with calibration support
|
|
pub struct DynamicQuantizer {
|
|
config: QuantizationConfig,
|
|
calibration_data: Arc<Mutex<Vec<TestTensor>>>,
|
|
is_calibrated: Arc<Mutex<bool>>,
|
|
params: Arc<Mutex<Option<QuantizationParams>>>,
|
|
}
|
|
|
|
impl DynamicQuantizer {
|
|
/// Create a new dynamic quantizer
|
|
#[must_use]
|
|
pub fn new(config: QuantizationConfig) -> Self {
|
|
Self {
|
|
config,
|
|
calibration_data: Arc::new(Mutex::new(Vec::new())),
|
|
is_calibrated: Arc::new(Mutex::new(false)),
|
|
params: Arc::new(Mutex::new(None)),
|
|
}
|
|
}
|
|
|
|
/// Perform calibration with provided data
|
|
pub async fn calibrate(&mut self, data: &[TestTensor]) -> InferenceResult<()> {
|
|
info!("Starting calibration with {} samples", data.len());
|
|
|
|
// Collect calibration statistics
|
|
let mut all_values = Vec::new();
|
|
for tensor in data {
|
|
tensor.validate()?;
|
|
all_values.extend_from_slice(&tensor.data);
|
|
}
|
|
|
|
if all_values.is_empty() {
|
|
return Err(InferenceError::invalid_request(
|
|
"No calibration data provided",
|
|
));
|
|
}
|
|
|
|
// Compute optimal quantization parameters using specified method
|
|
let params = match self.config.calibration_method {
|
|
CalibrationMethod::MinMax => self.calibrate_minmax(&all_values).await?,
|
|
CalibrationMethod::Percentile => self.calibrate_percentile(&all_values).await?,
|
|
CalibrationMethod::KLDivergence => self.calibrate_kl_divergence(&all_values).await?,
|
|
CalibrationMethod::MSE => self.calibrate_mse(&all_values).await?,
|
|
};
|
|
|
|
*self.params.lock().await = Some(params);
|
|
*self.is_calibrated.lock().await = true;
|
|
|
|
info!("Calibration completed successfully");
|
|
Ok(())
|
|
}
|
|
|
|
/// Quantize using calibrated parameters
|
|
pub async fn quantize(&self, tensor: &TestTensor) -> InferenceResult<TestTensor> {
|
|
let is_calibrated = *self.is_calibrated.lock().await;
|
|
if !is_calibrated {
|
|
return Err(InferenceError::invalid_request(
|
|
"Quantizer must be calibrated before use",
|
|
));
|
|
}
|
|
|
|
// Create a temporary quantizer with calibrated parameters
|
|
let quantizer = Quantizer {
|
|
config: self.config.clone(),
|
|
params: self.params.lock().await.clone(),
|
|
};
|
|
|
|
quantizer.quantize(tensor).await
|
|
}
|
|
|
|
/// Dequantize using calibrated parameters
|
|
pub async fn dequantize(&self, tensor: &TestTensor) -> InferenceResult<TestTensor> {
|
|
let is_calibrated = *self.is_calibrated.lock().await;
|
|
if !is_calibrated {
|
|
return Err(InferenceError::invalid_request(
|
|
"Quantizer must be calibrated before use",
|
|
));
|
|
}
|
|
|
|
let quantizer = Quantizer {
|
|
config: self.config.clone(),
|
|
params: self.params.lock().await.clone(),
|
|
};
|
|
|
|
quantizer.dequantize(tensor).await
|
|
}
|
|
|
|
// Calibration method implementations
|
|
|
|
async fn calibrate_minmax(&self, data: &[f32]) -> InferenceResult<QuantizationParams> {
|
|
let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
|
|
let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
|
|
|
|
let mut params = QuantizationParams::new(self.config.scheme);
|
|
params.min_value = min_val;
|
|
params.max_value = max_val;
|
|
|
|
let abs_max = max_val.abs().max(min_val.abs());
|
|
params.scale = match self.config.scheme {
|
|
QuantizationScheme::INT8 => abs_max / 127.0,
|
|
QuantizationScheme::INT4 => abs_max / 7.0,
|
|
QuantizationScheme::FP8E4M3 | QuantizationScheme::FP8E5M2 => abs_max / 240.0,
|
|
};
|
|
params.zero_point = 0;
|
|
|
|
if params.scale == 0.0 {
|
|
params.scale = 1e-8;
|
|
}
|
|
|
|
Ok(params)
|
|
}
|
|
|
|
async fn calibrate_percentile(&self, data: &[f32]) -> InferenceResult<QuantizationParams> {
|
|
let mut sorted_data = data.to_vec();
|
|
sorted_data.sort_by(f32::total_cmp);
|
|
|
|
// Use 99.9th percentile to handle outliers
|
|
let percentile = 99.9;
|
|
let lower_idx = ((100.0 - percentile) / 200.0 * sorted_data.len() as f32) as usize;
|
|
let upper_idx = (percentile / 100.0 * sorted_data.len() as f32) as usize;
|
|
|
|
let min_val = sorted_data[lower_idx];
|
|
let max_val = sorted_data[upper_idx.min(sorted_data.len() - 1)];
|
|
|
|
let mut params = QuantizationParams::new(self.config.scheme);
|
|
params.min_value = min_val;
|
|
params.max_value = max_val;
|
|
|
|
let abs_max = max_val.abs().max(min_val.abs());
|
|
params.scale = match self.config.scheme {
|
|
QuantizationScheme::INT8 => abs_max / 127.0,
|
|
QuantizationScheme::INT4 => abs_max / 7.0,
|
|
QuantizationScheme::FP8E4M3 | QuantizationScheme::FP8E5M2 => abs_max / 240.0,
|
|
};
|
|
params.zero_point = 0;
|
|
|
|
if params.scale == 0.0 {
|
|
params.scale = 1e-8;
|
|
}
|
|
|
|
Ok(params)
|
|
}
|
|
|
|
async fn calibrate_kl_divergence(&self, data: &[f32]) -> InferenceResult<QuantizationParams> {
|
|
// Simplified KL divergence calibration
|
|
// In practice, this would involve computing histogram distributions
|
|
// and finding the threshold that minimizes KL divergence
|
|
|
|
// For now, use a percentile-based approach as an approximation
|
|
self.calibrate_percentile(data).await
|
|
}
|
|
|
|
async fn calibrate_mse(&self, data: &[f32]) -> InferenceResult<QuantizationParams> {
|
|
// Simplified MSE-based calibration
|
|
// This would involve trying different scales and finding the one with minimum MSE
|
|
|
|
let mut best_params = self.calibrate_minmax(data).await?;
|
|
let mut best_mse = f32::INFINITY;
|
|
|
|
// Try different scale factors
|
|
let base_params = best_params.clone();
|
|
for factor in [0.8, 0.9, 1.0, 1.1, 1.2] {
|
|
let mut test_params = base_params.clone();
|
|
test_params.scale *= factor;
|
|
|
|
// Compute MSE for this scale
|
|
let mut mse = 0.0;
|
|
for &value in data {
|
|
let quantized = (value / test_params.scale).round() * test_params.scale;
|
|
mse += (value - quantized).powi(2);
|
|
}
|
|
mse /= data.len() as f32;
|
|
|
|
if mse < best_mse {
|
|
best_mse = mse;
|
|
best_params = test_params;
|
|
}
|
|
}
|
|
|
|
Ok(best_params)
|
|
}
|
|
}
|
|
|
|
/// Validation metrics for quantization quality
|
|
#[derive(Debug, Clone)]
|
|
pub struct ValidationMetrics {
|
|
pub mse: f32,
|
|
pub snr: f32,
|
|
pub relative_error: f32,
|
|
pub max_error: f32,
|
|
pub threshold: f32,
|
|
}
|
|
|
|
impl ValidationMetrics {
|
|
#[must_use]
|
|
pub fn passes_threshold(&self) -> bool {
|
|
self.mse < self.threshold
|
|
}
|
|
}
|
|
|
|
/// Validator for quantization accuracy
|
|
pub struct QuantizationValidator {
|
|
config: QuantizationConfig,
|
|
}
|
|
|
|
impl QuantizationValidator {
|
|
#[must_use]
|
|
pub fn new(config: QuantizationConfig) -> Self {
|
|
Self { config }
|
|
}
|
|
|
|
pub async fn validate(
|
|
&self,
|
|
original: &TestTensor,
|
|
quantized: &TestTensor,
|
|
) -> InferenceResult<ValidationMetrics> {
|
|
if original.data.len() != quantized.data.len() {
|
|
return Err(InferenceError::invalid_request("Tensor size mismatch"));
|
|
}
|
|
|
|
let mse = self.calculate_mse(&original.data, &quantized.data);
|
|
let snr = self.calculate_snr(&original.data, &quantized.data);
|
|
let relative_error = self.calculate_relative_error(&original.data, &quantized.data);
|
|
let max_error = self.calculate_max_error(&original.data, &quantized.data);
|
|
|
|
Ok(ValidationMetrics {
|
|
mse,
|
|
snr,
|
|
relative_error,
|
|
max_error,
|
|
threshold: self.config.accuracy_threshold,
|
|
})
|
|
}
|
|
|
|
fn calculate_mse(&self, original: &[f32], quantized: &[f32]) -> f32 {
|
|
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_snr(&self, 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_relative_error(&self, original: &[f32], quantized: &[f32]) -> f32 {
|
|
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_max_error(&self, original: &[f32], quantized: &[f32]) -> f32 {
|
|
original
|
|
.iter()
|
|
.zip(quantized.iter())
|
|
.map(|(o, q)| (o - q).abs())
|
|
.fold(0.0f32, f32::max)
|
|
}
|
|
}
|
|
|
|
/// Configuration for individual layers in mixed precision
|
|
#[derive(Debug, Clone)]
|
|
pub struct LayerConfig {
|
|
pub scheme: QuantizationScheme,
|
|
pub config: QuantizationConfig,
|
|
}
|
|
|
|
impl LayerConfig {
|
|
#[must_use]
|
|
pub fn new(scheme: QuantizationScheme) -> Self {
|
|
Self {
|
|
scheme,
|
|
config: QuantizationConfig::new().with_scheme(scheme),
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Mixed precision quantizer for different model layers
|
|
pub struct MixedPrecisionQuantizer {
|
|
layer_configs: HashMap<String, LayerConfig>,
|
|
quantizers: HashMap<String, Quantizer>,
|
|
}
|
|
|
|
impl Default for MixedPrecisionQuantizer {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl MixedPrecisionQuantizer {
|
|
#[must_use]
|
|
pub fn new() -> Self {
|
|
Self {
|
|
layer_configs: HashMap::new(),
|
|
quantizers: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
pub fn add_layer_config(&mut self, layer_name: &str, config: LayerConfig) {
|
|
let quantizer = Quantizer::new(config.config.clone());
|
|
self.layer_configs.insert(layer_name.to_string(), config);
|
|
self.quantizers.insert(layer_name.to_string(), quantizer);
|
|
}
|
|
|
|
pub async fn quantize_layer(
|
|
&self,
|
|
layer_name: &str,
|
|
tensor: &TestTensor,
|
|
) -> InferenceResult<TestTensor> {
|
|
let quantizer = self.quantizers.get(layer_name).ok_or_else(|| {
|
|
InferenceError::invalid_request(format!("No configuration for layer: {layer_name}"))
|
|
})?;
|
|
|
|
quantizer.quantize(tensor).await
|
|
}
|
|
|
|
pub async fn dequantize_layer(
|
|
&self,
|
|
layer_name: &str,
|
|
tensor: &TestTensor,
|
|
) -> InferenceResult<TestTensor> {
|
|
let quantizer = self.quantizers.get(layer_name).ok_or_else(|| {
|
|
InferenceError::invalid_request(format!("No configuration for layer: {layer_name}"))
|
|
})?;
|
|
|
|
quantizer.dequantize(tensor).await
|
|
}
|
|
}
|
|
|
|
/// Batch quantizer for processing multiple tensors
|
|
pub struct BatchQuantizer {
|
|
quantizer: Quantizer,
|
|
}
|
|
|
|
impl BatchQuantizer {
|
|
#[must_use]
|
|
pub fn new(config: QuantizationConfig) -> Self {
|
|
Self {
|
|
quantizer: Quantizer::new(config),
|
|
}
|
|
}
|
|
|
|
pub async fn quantize_batch(&self, tensors: &[TestTensor]) -> InferenceResult<Vec<TestTensor>> {
|
|
let mut results = Vec::with_capacity(tensors.len());
|
|
for tensor in tensors {
|
|
results.push(self.quantizer.quantize(tensor).await?);
|
|
}
|
|
Ok(results)
|
|
}
|
|
|
|
pub async fn dequantize_batch(
|
|
&self,
|
|
tensors: &[TestTensor],
|
|
) -> InferenceResult<Vec<TestTensor>> {
|
|
let mut results = Vec::with_capacity(tensors.len());
|
|
for tensor in tensors {
|
|
results.push(self.quantizer.dequantize(tensor).await?);
|
|
}
|
|
Ok(results)
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_config_validation() {
|
|
let valid_config = QuantizationConfig::new();
|
|
assert!(valid_config.validate().is_ok());
|
|
|
|
let invalid_config = QuantizationConfig::new().with_calibration_samples(0);
|
|
assert!(invalid_config.validate().is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_tensor_creation() {
|
|
let tensor = TestTensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![2, 2]);
|
|
assert_eq!(tensor.data.len(), 4);
|
|
assert_eq!(tensor.shape, vec![2, 2]);
|
|
assert_eq!(tensor.num_elements(), 4);
|
|
assert_eq!(tensor.dtype, DataType::F32);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_basic_quantization() {
|
|
let config = QuantizationConfig::new();
|
|
let quantizer = Quantizer::new(config);
|
|
|
|
let tensor = TestTensor::new(vec![1.0, 2.0, 3.0, 4.0], vec![4]);
|
|
let quantized = quantizer.quantize(&tensor).await.unwrap();
|
|
let dequantized = quantizer.dequantize(&quantized).await.unwrap();
|
|
|
|
assert_eq!(quantized.dtype, DataType::INT8);
|
|
assert_eq!(dequantized.dtype, DataType::F32);
|
|
assert_eq!(dequantized.data.len(), tensor.data.len());
|
|
}
|
|
}
|