232 lines
7.4 KiB
Rust
232 lines
7.4 KiB
Rust
//! Quantization module for model compression.
|
|
|
|
use forge_shared::{ModelInfo, QuantMethod, QuantPrecision, QuantizationConfig};
|
|
|
|
/// Quantizer for model weight quantization.
|
|
#[derive(Debug)]
|
|
pub struct Quantizer {
|
|
/// RNG state for calibration.
|
|
rng_state: u64,
|
|
}
|
|
|
|
impl Default for Quantizer {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl Quantizer {
|
|
/// Create a new quantizer.
|
|
pub fn new() -> Self {
|
|
Self { rng_state: 42 }
|
|
}
|
|
|
|
/// Quantize a model and return compression factor.
|
|
pub fn quantize(&mut self, model: &ModelInfo, config: &QuantizationConfig) -> f32 {
|
|
// Calculate bit reduction
|
|
let original_bits = model.precision_bits as f32;
|
|
let target_bits = self.precision_to_bits(config.precision);
|
|
|
|
// Base compression from bit reduction
|
|
let mut compression = original_bits / target_bits;
|
|
|
|
// Adjust for quantization method overhead
|
|
compression *= self.method_efficiency(config.method);
|
|
|
|
// Adjust for group quantization (adds scale/zero-point storage)
|
|
if let Some(group_size) = config.group_size {
|
|
let group_overhead = 1.0 + (2.0 * 16.0) / (group_size as f32 * target_bits);
|
|
compression /= group_overhead;
|
|
}
|
|
|
|
// Skip layers reduce effective compression
|
|
if !config.skip_layers.is_empty() {
|
|
let skip_fraction = config.skip_layers.len() as f32 / model.num_layers as f32;
|
|
compression = compression * (1.0 - skip_fraction) + skip_fraction;
|
|
}
|
|
|
|
compression.max(1.0)
|
|
}
|
|
|
|
/// Convert precision enum to bits.
|
|
fn precision_to_bits(&self, precision: QuantPrecision) -> f32 {
|
|
match precision {
|
|
QuantPrecision::Int2 => 2.0,
|
|
QuantPrecision::Int3 => 3.0,
|
|
QuantPrecision::Int4 => 4.0,
|
|
QuantPrecision::Int8 => 8.0,
|
|
QuantPrecision::FP8 => 8.0,
|
|
QuantPrecision::FP16 => 16.0,
|
|
QuantPrecision::BF16 => 16.0,
|
|
QuantPrecision::Mixed => 6.0, // Average of mixed precision
|
|
}
|
|
}
|
|
|
|
/// Get efficiency factor for quantization method.
|
|
fn method_efficiency(&self, method: QuantMethod) -> f32 {
|
|
match method {
|
|
QuantMethod::PTQ => 0.98, // Post-training has some overhead
|
|
QuantMethod::QAT => 0.99, // Training-aware is more efficient
|
|
QuantMethod::GPTQ => 0.97, // GPTQ has good compression
|
|
QuantMethod::AWQ => 0.98, // AWQ is efficient
|
|
QuantMethod::GGML => 0.95, // GGML has metadata overhead
|
|
QuantMethod::SmoothQuant => 0.96,
|
|
QuantMethod::Dynamic => 0.94, // Dynamic has runtime overhead
|
|
QuantMethod::Static => 0.99,
|
|
}
|
|
}
|
|
|
|
/// Estimate accuracy degradation from quantization.
|
|
pub fn estimate_degradation(&self, config: &QuantizationConfig) -> f64 {
|
|
let _target_bits = self.precision_to_bits(config.precision);
|
|
|
|
// Rough estimate: lower bits = higher degradation
|
|
let base_degradation = match config.precision {
|
|
QuantPrecision::Int2 => 0.15,
|
|
QuantPrecision::Int3 => 0.08,
|
|
QuantPrecision::Int4 => 0.03,
|
|
QuantPrecision::Int8 => 0.01,
|
|
QuantPrecision::FP8 => 0.01,
|
|
QuantPrecision::FP16 => 0.001,
|
|
QuantPrecision::BF16 => 0.002,
|
|
QuantPrecision::Mixed => 0.02,
|
|
};
|
|
|
|
// Better methods reduce degradation
|
|
let method_factor = match config.method {
|
|
QuantMethod::GPTQ | QuantMethod::AWQ => 0.7,
|
|
QuantMethod::QAT => 0.5,
|
|
_ => 1.0,
|
|
};
|
|
|
|
// Per-channel reduces degradation
|
|
let channel_factor = if config.per_channel { 0.8 } else { 1.0 };
|
|
|
|
// Calibration helps
|
|
let calib_factor = 1.0 - (config.calibration_size as f64 / 10000.0).min(0.3);
|
|
|
|
base_degradation * method_factor * channel_factor * calib_factor
|
|
}
|
|
|
|
/// Simulate calibration data collection.
|
|
pub fn calibrate(&mut self, num_samples: usize) -> Vec<f64> {
|
|
let mut calibration_data = Vec::with_capacity(num_samples);
|
|
|
|
for _ in 0..num_samples {
|
|
// Simulate collecting activation statistics
|
|
calibration_data.push(self.random_normal() * 0.1);
|
|
}
|
|
|
|
calibration_data
|
|
}
|
|
|
|
/// Random number.
|
|
fn random(&mut self) -> f64 {
|
|
self.rng_state = self
|
|
.rng_state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(self.rng_state >> 11) as f64 / (1u64 << 53) as f64
|
|
}
|
|
|
|
/// Random normal.
|
|
fn random_normal(&mut self) -> f64 {
|
|
let u1 = self.random() + 1e-10;
|
|
let u2 = self.random();
|
|
(-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_quantizer_creation() {
|
|
let quantizer = Quantizer::new();
|
|
assert_eq!(quantizer.rng_state, 42);
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_int4() {
|
|
let mut quantizer = Quantizer::new();
|
|
let model = forge_shared::sample_model_info();
|
|
let config = QuantizationConfig {
|
|
precision: QuantPrecision::Int4,
|
|
..Default::default()
|
|
};
|
|
|
|
let compression = quantizer.quantize(&model, &config);
|
|
assert!(compression > 2.0); // FP16 -> INT4 should give ~4x compression
|
|
assert!(compression < 5.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_quantize_int8() {
|
|
let mut quantizer = Quantizer::new();
|
|
let model = forge_shared::sample_model_info();
|
|
let config = QuantizationConfig {
|
|
precision: QuantPrecision::Int8,
|
|
..Default::default()
|
|
};
|
|
|
|
let compression = quantizer.quantize(&model, &config);
|
|
assert!(compression > 1.5); // FP16 -> INT8 should give ~2x compression
|
|
assert!(compression < 2.5);
|
|
}
|
|
|
|
#[test]
|
|
fn test_estimate_degradation() {
|
|
let quantizer = Quantizer::new();
|
|
|
|
let config_int4 = QuantizationConfig {
|
|
precision: QuantPrecision::Int4,
|
|
..Default::default()
|
|
};
|
|
let config_int8 = QuantizationConfig {
|
|
precision: QuantPrecision::Int8,
|
|
..Default::default()
|
|
};
|
|
|
|
let deg_int4 = quantizer.estimate_degradation(&config_int4);
|
|
let deg_int8 = quantizer.estimate_degradation(&config_int8);
|
|
|
|
// INT4 should have higher degradation than INT8
|
|
assert!(deg_int4 > deg_int8);
|
|
}
|
|
|
|
#[test]
|
|
fn test_calibrate() {
|
|
let mut quantizer = Quantizer::new();
|
|
let data = quantizer.calibrate(100);
|
|
|
|
assert_eq!(data.len(), 100);
|
|
for val in &data {
|
|
assert!(val.is_finite());
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_group_quantization_overhead() {
|
|
let mut quantizer = Quantizer::new();
|
|
let model = forge_shared::sample_model_info();
|
|
|
|
let config_no_group = QuantizationConfig {
|
|
precision: QuantPrecision::Int4,
|
|
group_size: None,
|
|
..Default::default()
|
|
};
|
|
let config_with_group = QuantizationConfig {
|
|
precision: QuantPrecision::Int4,
|
|
group_size: Some(128),
|
|
..Default::default()
|
|
};
|
|
|
|
let compression_no_group = quantizer.quantize(&model, &config_no_group);
|
|
let compression_with_group = quantizer.quantize(&model, &config_with_group);
|
|
|
|
// Group quantization adds overhead, so compression should be slightly less
|
|
assert!(compression_no_group >= compression_with_group);
|
|
}
|
|
}
|