osobhandClaude Sonnet 5
4aaa36a57a
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU backward work. Verified formatting-only via diff sampling; no logic changed. Co-Authored-By: Claude Sonnet 5 <[email protected]>
1411 lines
47 KiB
Rust
1411 lines
47 KiB
Rust
//! Advanced Quantization Methods
|
||
//!
|
||
//! This module implements state-of-the-art quantization techniques:
|
||
//! - AWQ (Activation-aware Weight Quantization): Finds optimal scales using activation importance
|
||
//! - GPTQ (Accurate Post-Training Quantization): Hessian-based blockwise quantization
|
||
//! - SmoothQuant: Migrates quantization difficulty from activations to weights
|
||
//!
|
||
//! These methods enable INT4/INT3 weight-only quantization with minimal accuracy loss.
|
||
|
||
use crate::{
|
||
Result,
|
||
error::{CompressionError, QuantizationError},
|
||
};
|
||
use rtx_tensor::{Device, Tensor};
|
||
use serde::{Deserialize, Serialize};
|
||
use std::collections::HashMap;
|
||
|
||
// Helper to convert tensor errors
|
||
fn tensor_err(e: impl std::fmt::Display) -> CompressionError {
|
||
CompressionError::Quantization(QuantizationError::TensorError(e.to_string()))
|
||
}
|
||
|
||
// =============================================================================
|
||
// AWQ (Activation-aware Weight Quantization)
|
||
// =============================================================================
|
||
|
||
/// Configuration for AWQ quantization
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct AWQConfig {
|
||
/// Target bit width (typically 4 or 3)
|
||
pub bit_width: u8,
|
||
/// Group size for quantization (typically 128)
|
||
pub group_size: usize,
|
||
/// Number of calibration samples
|
||
pub num_calibration_samples: usize,
|
||
/// Percentile for activation importance (typically 0.9-0.99)
|
||
pub activation_percentile: f32,
|
||
/// Alpha for mixing search (0.0-1.0)
|
||
pub alpha: f32,
|
||
/// Number of search iterations for optimal scale
|
||
pub num_search_iters: usize,
|
||
/// Enable per-channel scaling
|
||
pub per_channel: bool,
|
||
/// Device for computation
|
||
pub device: Device,
|
||
}
|
||
|
||
impl Default for AWQConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
bit_width: 4,
|
||
group_size: 128,
|
||
num_calibration_samples: 128,
|
||
activation_percentile: 0.95,
|
||
alpha: 0.5,
|
||
num_search_iters: 20,
|
||
per_channel: true,
|
||
device: Device::Cpu,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl AWQConfig {
|
||
/// Create new AWQ config with specified bit width
|
||
pub fn new(bit_width: u8) -> Self {
|
||
Self {
|
||
bit_width,
|
||
..Default::default()
|
||
}
|
||
}
|
||
|
||
/// Set group size
|
||
pub fn with_group_size(mut self, group_size: usize) -> Self {
|
||
self.group_size = group_size;
|
||
self
|
||
}
|
||
|
||
/// Set calibration samples
|
||
pub fn with_calibration_samples(mut self, num_samples: usize) -> Self {
|
||
self.num_calibration_samples = num_samples;
|
||
self
|
||
}
|
||
}
|
||
|
||
/// AWQ Quantizer for activation-aware weight quantization
|
||
#[derive(Debug)]
|
||
pub struct AWQQuantizer {
|
||
config: AWQConfig,
|
||
/// Per-channel scales computed from activation importance
|
||
activation_scales: HashMap<String, Vec<f32>>,
|
||
/// Statistics collected during calibration
|
||
stats: AWQStats,
|
||
}
|
||
|
||
/// Statistics from AWQ quantization
|
||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||
pub struct AWQStats {
|
||
/// Number of layers quantized
|
||
pub layers_quantized: usize,
|
||
/// Total weights quantized
|
||
pub total_weights: usize,
|
||
/// Average quantization error (MSE)
|
||
pub avg_quant_error: f64,
|
||
/// Salient channel ratios per layer
|
||
pub salient_ratios: HashMap<String, f32>,
|
||
}
|
||
|
||
impl AWQQuantizer {
|
||
/// Create new AWQ quantizer
|
||
pub fn new(config: AWQConfig) -> Self {
|
||
Self {
|
||
config,
|
||
activation_scales: HashMap::new(),
|
||
stats: AWQStats::default(),
|
||
}
|
||
}
|
||
|
||
/// Calibrate scales using activation data
|
||
pub fn calibrate(&mut self, activations: &HashMap<String, Vec<Tensor>>) -> Result<()> {
|
||
for (layer_name, act_samples) in activations {
|
||
// Compute per-channel activation magnitudes
|
||
let channel_magnitudes = self.compute_channel_magnitudes(act_samples)?;
|
||
|
||
// Find salient channels (high activation magnitude)
|
||
let percentile_threshold =
|
||
self.compute_percentile(&channel_magnitudes, self.config.activation_percentile);
|
||
|
||
// Compute scales based on activation importance
|
||
let scales: Vec<f32> = channel_magnitudes
|
||
.iter()
|
||
.map(|&mag| {
|
||
if mag > percentile_threshold {
|
||
// Salient channel: use larger scale to preserve precision
|
||
(mag / percentile_threshold).sqrt()
|
||
} else {
|
||
1.0
|
||
}
|
||
})
|
||
.collect();
|
||
|
||
self.activation_scales.insert(layer_name.clone(), scales);
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
/// Quantize weights using AWQ method
|
||
pub fn quantize_weights(
|
||
&mut self,
|
||
layer_name: &str,
|
||
weights: &Tensor,
|
||
) -> Result<AWQQuantizedWeight> {
|
||
let shape = weights.shape().dims();
|
||
let out_features = shape[0];
|
||
let in_features = shape[1];
|
||
|
||
// Get activation scales for this layer (or use uniform scales)
|
||
let scales = self
|
||
.activation_scales
|
||
.get(layer_name)
|
||
.cloned()
|
||
.unwrap_or_else(|| vec![1.0; in_features]);
|
||
|
||
// Scale weights by activation importance
|
||
let mut weight_data = weights.to_vec().map_err(tensor_err)?;
|
||
for j in 0..out_features {
|
||
for i in 0..in_features {
|
||
let idx = j * in_features + i;
|
||
let scale = scales.get(i).copied().unwrap_or(1.0);
|
||
weight_data[idx] *= scale;
|
||
}
|
||
}
|
||
|
||
// Create scaled weights tensor for group quantization
|
||
let scaled_weights = Tensor::from_slice(
|
||
&weight_data,
|
||
&[out_features, in_features],
|
||
&self.config.device,
|
||
)
|
||
.map_err(tensor_err)?;
|
||
|
||
// Quantize scaled weights using group quantization
|
||
let (quantized, quant_scales, zeros) = self.group_quantize(&scaled_weights)?;
|
||
|
||
// Compute inverse scales for dequantization
|
||
let inv_activation_scales: Vec<f32> = scales.iter().map(|&s| 1.0 / s).collect();
|
||
|
||
self.stats.layers_quantized += 1;
|
||
self.stats.total_weights += out_features * in_features;
|
||
|
||
Ok(AWQQuantizedWeight {
|
||
quantized_data: quantized,
|
||
scales: quant_scales,
|
||
zeros,
|
||
activation_scales: inv_activation_scales,
|
||
group_size: self.config.group_size,
|
||
bit_width: self.config.bit_width,
|
||
shape: vec![out_features, in_features],
|
||
})
|
||
}
|
||
|
||
/// Group-wise quantization of weights
|
||
fn group_quantize(&self, weights: &Tensor) -> Result<(Vec<i8>, Vec<f32>, Vec<i8>)> {
|
||
let shape = weights.shape().dims();
|
||
let total_elements = shape.iter().product::<usize>();
|
||
let weight_data = weights.to_vec().map_err(tensor_err)?;
|
||
|
||
let mut quantized = Vec::with_capacity(total_elements);
|
||
let mut scales = Vec::new();
|
||
let mut zeros = Vec::new();
|
||
|
||
let num_groups = (total_elements + self.config.group_size - 1) / self.config.group_size;
|
||
|
||
for group_idx in 0..num_groups {
|
||
let start = group_idx * self.config.group_size;
|
||
let end = (start + self.config.group_size).min(total_elements);
|
||
|
||
// Find min/max in this group
|
||
let mut min_val = f32::MAX;
|
||
let mut max_val = f32::MIN;
|
||
for i in start..end {
|
||
let val = weight_data[i];
|
||
min_val = min_val.min(val);
|
||
max_val = max_val.max(val);
|
||
}
|
||
|
||
// Compute scale and zero point
|
||
let qmin = -(1 << (self.config.bit_width - 1));
|
||
let qmax = (1 << (self.config.bit_width - 1)) - 1;
|
||
|
||
let scale = (max_val - min_val) / (qmax - qmin) as f32;
|
||
let scale = if scale == 0.0 { 1.0 } else { scale };
|
||
let zero = ((qmin as f32 * scale - min_val) / scale).round() as i8;
|
||
|
||
scales.push(scale);
|
||
zeros.push(zero);
|
||
|
||
// Quantize values in group
|
||
for i in start..end {
|
||
let val = weight_data[i];
|
||
let q = ((val / scale) + zero as f32).round();
|
||
let q = q.max(qmin as f32).min(qmax as f32) as i8;
|
||
quantized.push(q);
|
||
}
|
||
}
|
||
|
||
Ok((quantized, scales, zeros))
|
||
}
|
||
|
||
/// Compute per-channel activation magnitudes
|
||
fn compute_channel_magnitudes(&self, activations: &[Tensor]) -> Result<Vec<f32>> {
|
||
if activations.is_empty() {
|
||
return Ok(Vec::new());
|
||
}
|
||
|
||
let shape = activations[0].shape().dims();
|
||
let num_channels = if shape.len() >= 2 {
|
||
shape[shape.len() - 1]
|
||
} else {
|
||
shape[0]
|
||
};
|
||
|
||
let mut magnitudes = vec![0.0f32; num_channels];
|
||
let mut counts = vec![0usize; num_channels];
|
||
|
||
for act in activations {
|
||
let flat = act.to_vec().unwrap_or_default();
|
||
for (i, &val) in flat.iter().enumerate() {
|
||
let channel_idx = i % num_channels;
|
||
magnitudes[channel_idx] += val.abs();
|
||
counts[channel_idx] += 1;
|
||
}
|
||
}
|
||
|
||
// Average magnitudes
|
||
for i in 0..num_channels {
|
||
if counts[i] > 0 {
|
||
magnitudes[i] /= counts[i] as f32;
|
||
}
|
||
}
|
||
|
||
Ok(magnitudes)
|
||
}
|
||
|
||
/// Compute percentile value
|
||
fn compute_percentile(&self, values: &[f32], percentile: f32) -> f32 {
|
||
if values.is_empty() {
|
||
return 0.0;
|
||
}
|
||
let mut sorted = values.to_vec();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||
let idx = ((sorted.len() as f32 - 1.0) * percentile) as usize;
|
||
sorted[idx.min(sorted.len() - 1)]
|
||
}
|
||
|
||
/// Get statistics
|
||
pub fn stats(&self) -> &AWQStats {
|
||
&self.stats
|
||
}
|
||
}
|
||
|
||
/// AWQ quantized weight representation
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct AWQQuantizedWeight {
|
||
/// Quantized weight data (packed integers)
|
||
pub quantized_data: Vec<i8>,
|
||
/// Per-group scales
|
||
pub scales: Vec<f32>,
|
||
/// Per-group zero points
|
||
pub zeros: Vec<i8>,
|
||
/// Activation-derived inverse scales for dequant
|
||
pub activation_scales: Vec<f32>,
|
||
/// Group size
|
||
pub group_size: usize,
|
||
/// Bit width
|
||
pub bit_width: u8,
|
||
/// Original shape [out_features, in_features]
|
||
pub shape: Vec<usize>,
|
||
}
|
||
|
||
impl AWQQuantizedWeight {
|
||
/// Dequantize weights
|
||
pub fn dequantize(&self) -> Result<Tensor> {
|
||
let out_features = self.shape[0];
|
||
let in_features = self.shape[1];
|
||
let total = out_features * in_features;
|
||
|
||
let mut dequantized = vec![0.0f32; total];
|
||
let num_groups = self.scales.len();
|
||
|
||
for (i, &qval) in self.quantized_data.iter().enumerate() {
|
||
if i >= total {
|
||
break;
|
||
}
|
||
|
||
let group_idx = i / self.group_size;
|
||
let group_idx = group_idx.min(num_groups - 1);
|
||
|
||
let scale = self.scales[group_idx];
|
||
let zero = self.zeros[group_idx] as f32;
|
||
let act_scale = self
|
||
.activation_scales
|
||
.get(i % in_features)
|
||
.copied()
|
||
.unwrap_or(1.0);
|
||
|
||
dequantized[i] = (qval as f32 - zero) * scale * act_scale;
|
||
}
|
||
|
||
Tensor::from_slice(&dequantized, &[out_features, in_features], &Device::Cpu)
|
||
.map_err(tensor_err)
|
||
}
|
||
|
||
/// Get memory footprint in bytes
|
||
pub fn memory_bytes(&self) -> usize {
|
||
let data_bytes = (self.quantized_data.len() * self.bit_width as usize + 7) / 8;
|
||
let scale_bytes = self.scales.len() * 4; // f32
|
||
let zero_bytes = self.zeros.len(); // i8
|
||
let act_scale_bytes = self.activation_scales.len() * 4; // f32
|
||
data_bytes + scale_bytes + zero_bytes + act_scale_bytes
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// GPTQ (Accurate Post-Training Quantization)
|
||
// =============================================================================
|
||
|
||
/// Configuration for GPTQ quantization
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct GPTQConfig {
|
||
/// Target bit width (typically 4 or 3)
|
||
pub bit_width: u8,
|
||
/// Group size for quantization
|
||
pub group_size: usize,
|
||
/// Block size for Hessian computation
|
||
pub block_size: usize,
|
||
/// Damping factor for Hessian (prevents numerical instability)
|
||
pub damp_percent: f32,
|
||
/// Enable static groups (same group across different rows)
|
||
pub static_groups: bool,
|
||
/// Number of calibration samples
|
||
pub num_calibration_samples: usize,
|
||
/// Enable activation order optimization
|
||
pub act_order: bool,
|
||
/// Use true sequential (quantize one column at a time)
|
||
pub true_sequential: bool,
|
||
/// Device for computation
|
||
pub device: Device,
|
||
}
|
||
|
||
impl Default for GPTQConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
bit_width: 4,
|
||
group_size: 128,
|
||
block_size: 128,
|
||
damp_percent: 0.01,
|
||
static_groups: true,
|
||
num_calibration_samples: 128,
|
||
act_order: false,
|
||
true_sequential: true,
|
||
device: Device::Cpu,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl GPTQConfig {
|
||
/// Create new GPTQ config
|
||
pub fn new(bit_width: u8) -> Self {
|
||
Self {
|
||
bit_width,
|
||
..Default::default()
|
||
}
|
||
}
|
||
|
||
/// Set group size
|
||
pub fn with_group_size(mut self, group_size: usize) -> Self {
|
||
self.group_size = group_size;
|
||
self
|
||
}
|
||
|
||
/// Enable activation order optimization
|
||
pub fn with_act_order(mut self, enabled: bool) -> Self {
|
||
self.act_order = enabled;
|
||
self
|
||
}
|
||
}
|
||
|
||
/// GPTQ Quantizer using Hessian-based optimization
|
||
#[derive(Debug)]
|
||
pub struct GPTQQuantizer {
|
||
config: GPTQConfig,
|
||
/// Cached Hessian matrices per layer
|
||
hessians: HashMap<String, Vec<Vec<f32>>>,
|
||
/// Statistics
|
||
stats: GPTQStats,
|
||
}
|
||
|
||
/// Statistics from GPTQ quantization
|
||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||
pub struct GPTQStats {
|
||
/// Number of layers quantized
|
||
pub layers_quantized: usize,
|
||
/// Total weights quantized
|
||
pub total_weights: usize,
|
||
/// Average quantization error before compensation
|
||
pub avg_error_before: f64,
|
||
/// Average quantization error after compensation
|
||
pub avg_error_after: f64,
|
||
/// Hessian computation time (ms)
|
||
pub hessian_time_ms: f64,
|
||
/// Quantization time (ms)
|
||
pub quant_time_ms: f64,
|
||
}
|
||
|
||
impl GPTQQuantizer {
|
||
/// Create new GPTQ quantizer
|
||
pub fn new(config: GPTQConfig) -> Self {
|
||
Self {
|
||
config,
|
||
hessians: HashMap::new(),
|
||
stats: GPTQStats::default(),
|
||
}
|
||
}
|
||
|
||
/// Calibrate Hessian using input activations
|
||
pub fn calibrate(&mut self, layer_name: &str, inputs: &[Tensor]) -> Result<()> {
|
||
if inputs.is_empty() {
|
||
return Ok(());
|
||
}
|
||
|
||
let in_features = inputs[0].shape().dims().last().copied().unwrap_or(1);
|
||
|
||
// Initialize Hessian as H = X^T * X (accumulated)
|
||
let mut hessian = vec![vec![0.0f32; in_features]; in_features];
|
||
|
||
for input in inputs {
|
||
let flat = input.to_vec().map_err(tensor_err)?;
|
||
let batch_size = flat.len() / in_features;
|
||
|
||
for b in 0..batch_size {
|
||
for i in 0..in_features {
|
||
for j in 0..in_features {
|
||
let xi = flat[b * in_features + i];
|
||
let xj = flat[b * in_features + j];
|
||
hessian[i][j] += xi * xj;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
// Normalize by number of samples
|
||
let num_samples = inputs.iter().map(|t| t.shape().dims()[0]).sum::<usize>();
|
||
for i in 0..in_features {
|
||
for j in 0..in_features {
|
||
hessian[i][j] /= num_samples as f32;
|
||
}
|
||
}
|
||
|
||
// Add damping to diagonal
|
||
let diag_mean: f32 =
|
||
(0..in_features).map(|i| hessian[i][i]).sum::<f32>() / in_features as f32;
|
||
let damp = self.config.damp_percent * diag_mean;
|
||
for i in 0..in_features {
|
||
hessian[i][i] += damp;
|
||
}
|
||
|
||
self.hessians.insert(layer_name.to_string(), hessian);
|
||
Ok(())
|
||
}
|
||
|
||
/// Quantize weights using GPTQ algorithm
|
||
pub fn quantize_weights(
|
||
&mut self,
|
||
layer_name: &str,
|
||
weights: &Tensor,
|
||
) -> Result<GPTQQuantizedWeight> {
|
||
let shape = weights.shape().dims();
|
||
let out_features = shape[0];
|
||
let in_features = shape[1];
|
||
|
||
// Get Hessian (or create identity if not calibrated)
|
||
let hessian = self.hessians.get(layer_name).cloned().unwrap_or_else(|| {
|
||
let mut h = vec![vec![0.0f32; in_features]; in_features];
|
||
for i in 0..in_features {
|
||
h[i][i] = 1.0;
|
||
}
|
||
h
|
||
});
|
||
|
||
// Compute Cholesky decomposition of Hessian inverse
|
||
let h_inv = self.cholesky_inverse(&hessian)?;
|
||
|
||
// Working copy of weights
|
||
let mut w = weights.to_vec().map_err(tensor_err)?;
|
||
let mut quantized = vec![0i8; out_features * in_features];
|
||
let mut scales = Vec::new();
|
||
let mut zeros = Vec::new();
|
||
|
||
// Determine column order (optionally by activation importance)
|
||
let col_order: Vec<usize> = if self.config.act_order {
|
||
// Sort by diagonal of Hessian (activation importance)
|
||
let mut indexed: Vec<_> = (0..in_features).map(|i| (i, hessian[i][i])).collect();
|
||
indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
|
||
indexed.into_iter().map(|(i, _)| i).collect()
|
||
} else {
|
||
(0..in_features).collect()
|
||
};
|
||
|
||
// Quantize column by column
|
||
for (block_start, _) in (0..in_features).step_by(self.config.block_size).enumerate() {
|
||
let block_end = (block_start + self.config.block_size).min(in_features);
|
||
|
||
for col_idx in block_start..block_end {
|
||
let col = col_order[col_idx];
|
||
|
||
// Compute group index
|
||
let group_idx = col / self.config.group_size;
|
||
|
||
// Find min/max for this column across all rows
|
||
let mut col_min = f32::MAX;
|
||
let mut col_max = f32::MIN;
|
||
for row in 0..out_features {
|
||
let val = w[row * in_features + col];
|
||
col_min = col_min.min(val);
|
||
col_max = col_max.max(val);
|
||
}
|
||
|
||
// Compute scale and zero
|
||
let qmin = -(1 << (self.config.bit_width - 1));
|
||
let qmax = (1 << (self.config.bit_width - 1)) - 1;
|
||
let scale = (col_max - col_min) / (qmax - qmin) as f32;
|
||
let scale = if scale == 0.0 { 1.0 } else { scale };
|
||
let zero = ((qmin as f32 * scale - col_min) / scale).round() as i8;
|
||
|
||
if group_idx >= scales.len() {
|
||
scales.push(scale);
|
||
zeros.push(zero);
|
||
}
|
||
|
||
// Quantize this column for each row
|
||
for row in 0..out_features {
|
||
let idx = row * in_features + col;
|
||
let val = w[idx];
|
||
let q = ((val / scale) + zero as f32).round();
|
||
let q = q.max(qmin as f32).min(qmax as f32);
|
||
quantized[idx] = q as i8;
|
||
|
||
// Compute quantization error
|
||
let dequant = (q - zero as f32) * scale;
|
||
let error = val - dequant;
|
||
|
||
// Update remaining weights to compensate for error (GPTQ key insight)
|
||
// w_remaining -= error * H_inv[col, remaining] / H_inv[col, col]
|
||
let h_diag = h_inv[col][col];
|
||
if h_diag > 1e-10 {
|
||
for remaining_col in (col_idx + 1)..in_features {
|
||
let rem = col_order[remaining_col];
|
||
let update = error * h_inv[col][rem] / h_diag;
|
||
w[row * in_features + rem] -= update;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
self.stats.layers_quantized += 1;
|
||
self.stats.total_weights += out_features * in_features;
|
||
|
||
Ok(GPTQQuantizedWeight {
|
||
quantized_data: quantized,
|
||
scales,
|
||
zeros,
|
||
group_size: self.config.group_size,
|
||
bit_width: self.config.bit_width,
|
||
shape: vec![out_features, in_features],
|
||
col_order: if self.config.act_order {
|
||
Some(col_order)
|
||
} else {
|
||
None
|
||
},
|
||
})
|
||
}
|
||
|
||
/// Compute Cholesky-based inverse of a symmetric positive definite matrix
|
||
fn cholesky_inverse(&self, matrix: &[Vec<f32>]) -> Result<Vec<Vec<f32>>> {
|
||
let n = matrix.len();
|
||
|
||
// Simple pseudo-inverse using regularized SVD-like approach
|
||
// For production, use proper linear algebra library
|
||
let mut inv = vec![vec![0.0f32; n]; n];
|
||
|
||
for i in 0..n {
|
||
let diag = matrix[i][i];
|
||
if diag > 1e-10 {
|
||
inv[i][i] = 1.0 / diag;
|
||
} else {
|
||
inv[i][i] = 1.0;
|
||
}
|
||
}
|
||
|
||
// For off-diagonal elements, use approximate inverse
|
||
for i in 0..n {
|
||
for j in 0..n {
|
||
if i != j && matrix[i][i] > 1e-10 && matrix[j][j] > 1e-10 {
|
||
inv[i][j] = -matrix[i][j] / (matrix[i][i] * matrix[j][j]).sqrt();
|
||
}
|
||
}
|
||
}
|
||
|
||
Ok(inv)
|
||
}
|
||
|
||
/// Get statistics
|
||
pub fn stats(&self) -> &GPTQStats {
|
||
&self.stats
|
||
}
|
||
}
|
||
|
||
/// GPTQ quantized weight representation
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct GPTQQuantizedWeight {
|
||
/// Quantized weight data
|
||
pub quantized_data: Vec<i8>,
|
||
/// Per-group scales
|
||
pub scales: Vec<f32>,
|
||
/// Per-group zero points
|
||
pub zeros: Vec<i8>,
|
||
/// Group size
|
||
pub group_size: usize,
|
||
/// Bit width
|
||
pub bit_width: u8,
|
||
/// Original shape
|
||
pub shape: Vec<usize>,
|
||
/// Column order (if act_order enabled)
|
||
pub col_order: Option<Vec<usize>>,
|
||
}
|
||
|
||
impl GPTQQuantizedWeight {
|
||
/// Dequantize weights
|
||
pub fn dequantize(&self) -> Result<Tensor> {
|
||
let out_features = self.shape[0];
|
||
let in_features = self.shape[1];
|
||
|
||
let mut dequantized = vec![0.0f32; out_features * in_features];
|
||
|
||
for row in 0..out_features {
|
||
for col in 0..in_features {
|
||
let idx = row * in_features + col;
|
||
let group_idx = col / self.group_size;
|
||
let group_idx = group_idx.min(self.scales.len() - 1);
|
||
|
||
let q = self.quantized_data[idx] as f32;
|
||
let scale = self.scales[group_idx];
|
||
let zero = self.zeros[group_idx] as f32;
|
||
|
||
dequantized[idx] = (q - zero) * scale;
|
||
}
|
||
}
|
||
|
||
Tensor::from_slice(&dequantized, &[out_features, in_features], &Device::Cpu)
|
||
.map_err(tensor_err)
|
||
}
|
||
|
||
/// Get compression ratio
|
||
pub fn compression_ratio(&self) -> f32 {
|
||
let original_bits = self.shape[0] * self.shape[1] * 32; // FP32
|
||
let quantized_bits = self.quantized_data.len() * self.bit_width as usize
|
||
+ self.scales.len() * 32
|
||
+ self.zeros.len() * 8;
|
||
original_bits as f32 / quantized_bits as f32
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// SmoothQuant
|
||
// =============================================================================
|
||
|
||
/// Configuration for SmoothQuant
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct SmoothQuantConfig {
|
||
/// Smoothing alpha (0.0-1.0, higher = more smoothing on activations)
|
||
pub alpha: f32,
|
||
/// Target bit width for weights
|
||
pub weight_bit_width: u8,
|
||
/// Target bit width for activations
|
||
pub activation_bit_width: u8,
|
||
/// Number of calibration samples
|
||
pub num_calibration_samples: usize,
|
||
/// Per-channel smoothing
|
||
pub per_channel: bool,
|
||
/// Minimum scale to prevent division by zero
|
||
pub min_scale: f32,
|
||
/// Device for computation
|
||
pub device: Device,
|
||
}
|
||
|
||
impl Default for SmoothQuantConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
alpha: 0.5,
|
||
weight_bit_width: 8,
|
||
activation_bit_width: 8,
|
||
num_calibration_samples: 128,
|
||
per_channel: true,
|
||
min_scale: 1e-5,
|
||
device: Device::Cpu,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl SmoothQuantConfig {
|
||
/// Create new SmoothQuant config
|
||
pub fn new(alpha: f32) -> Self {
|
||
Self {
|
||
alpha: alpha.clamp(0.0, 1.0),
|
||
..Default::default()
|
||
}
|
||
}
|
||
|
||
/// Set bit widths
|
||
pub fn with_bit_widths(mut self, weight_bits: u8, activation_bits: u8) -> Self {
|
||
self.weight_bit_width = weight_bits;
|
||
self.activation_bit_width = activation_bits;
|
||
self
|
||
}
|
||
}
|
||
|
||
/// SmoothQuant migrates quantization difficulty from activations to weights
|
||
#[derive(Debug)]
|
||
pub struct SmoothQuantQuantizer {
|
||
config: SmoothQuantConfig,
|
||
/// Smoothing scales per layer (applied to activations)
|
||
smoothing_scales: HashMap<String, Vec<f32>>,
|
||
/// Statistics
|
||
stats: SmoothQuantStats,
|
||
}
|
||
|
||
/// Statistics from SmoothQuant
|
||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||
pub struct SmoothQuantStats {
|
||
/// Number of layers smoothed
|
||
pub layers_smoothed: usize,
|
||
/// Average activation range before smoothing
|
||
pub avg_act_range_before: f64,
|
||
/// Average activation range after smoothing
|
||
pub avg_act_range_after: f64,
|
||
/// Average weight range before smoothing
|
||
pub avg_weight_range_before: f64,
|
||
/// Average weight range after smoothing
|
||
pub avg_weight_range_after: f64,
|
||
}
|
||
|
||
impl SmoothQuantQuantizer {
|
||
/// Create new SmoothQuant quantizer
|
||
pub fn new(config: SmoothQuantConfig) -> Self {
|
||
Self {
|
||
config,
|
||
smoothing_scales: HashMap::new(),
|
||
stats: SmoothQuantStats::default(),
|
||
}
|
||
}
|
||
|
||
/// Calibrate smoothing scales using activations and weights
|
||
pub fn calibrate(
|
||
&mut self,
|
||
layer_name: &str,
|
||
activations: &[Tensor],
|
||
weights: &Tensor,
|
||
) -> Result<()> {
|
||
if activations.is_empty() {
|
||
return Ok(());
|
||
}
|
||
|
||
let in_features = weights.shape().dims()[1];
|
||
|
||
// Compute per-channel activation max
|
||
let mut act_max = vec![0.0f32; in_features];
|
||
for act in activations {
|
||
let flat = act.to_vec().map_err(tensor_err)?;
|
||
let num_elements = flat.len();
|
||
for (i, &val) in flat.iter().enumerate() {
|
||
let channel = i % in_features;
|
||
act_max[channel] = act_max[channel].max(val.abs());
|
||
}
|
||
}
|
||
|
||
// Compute per-channel weight max (for each input channel)
|
||
let weight_flat = weights.to_vec().map_err(tensor_err)?;
|
||
let out_features = weights.shape().dims()[0];
|
||
let mut weight_max = vec![0.0f32; in_features];
|
||
for out_idx in 0..out_features {
|
||
for in_idx in 0..in_features {
|
||
let val = weight_flat[out_idx * in_features + in_idx];
|
||
weight_max[in_idx] = weight_max[in_idx].max(val.abs());
|
||
}
|
||
}
|
||
|
||
// Compute smoothing scales: s = act_max^alpha / weight_max^(1-alpha)
|
||
let alpha = self.config.alpha;
|
||
let scales: Vec<f32> = (0..in_features)
|
||
.map(|i| {
|
||
let a = act_max[i].max(self.config.min_scale);
|
||
let w = weight_max[i].max(self.config.min_scale);
|
||
(a.powf(alpha) / w.powf(1.0 - alpha)).max(self.config.min_scale)
|
||
})
|
||
.collect();
|
||
|
||
self.smoothing_scales.insert(layer_name.to_string(), scales);
|
||
self.stats.layers_smoothed += 1;
|
||
|
||
Ok(())
|
||
}
|
||
|
||
/// Get smoothing scales for a layer
|
||
pub fn get_scales(&self, layer_name: &str) -> Option<&Vec<f32>> {
|
||
self.smoothing_scales.get(layer_name)
|
||
}
|
||
|
||
/// Apply smoothing to activations (multiply by inverse scale)
|
||
pub fn smooth_activations(&self, layer_name: &str, activations: &Tensor) -> Result<Tensor> {
|
||
let scales = self.smoothing_scales.get(layer_name).ok_or_else(|| {
|
||
CompressionError::Quantization(QuantizationError::TensorError(format!(
|
||
"No scales for layer {}",
|
||
layer_name
|
||
)))
|
||
})?;
|
||
|
||
let shape = activations.shape().dims();
|
||
let in_features = shape[shape.len() - 1];
|
||
|
||
let mut smoothed = activations.to_vec().map_err(tensor_err)?;
|
||
for (i, val) in smoothed.iter_mut().enumerate() {
|
||
let channel = i % in_features;
|
||
let inv_scale = 1.0 / scales.get(channel).copied().unwrap_or(1.0);
|
||
*val *= inv_scale;
|
||
}
|
||
|
||
Tensor::from_slice(&smoothed, shape, &self.config.device).map_err(tensor_err)
|
||
}
|
||
|
||
/// Apply inverse smoothing to weights (multiply by scale)
|
||
pub fn smooth_weights(&self, layer_name: &str, weights: &Tensor) -> Result<Tensor> {
|
||
let scales = self.smoothing_scales.get(layer_name).ok_or_else(|| {
|
||
CompressionError::Quantization(QuantizationError::TensorError(format!(
|
||
"No scales for layer {}",
|
||
layer_name
|
||
)))
|
||
})?;
|
||
|
||
let shape = weights.shape().dims();
|
||
let out_features = shape[0];
|
||
let in_features = shape[1];
|
||
|
||
let mut smoothed = weights.to_vec().map_err(tensor_err)?;
|
||
for out_idx in 0..out_features {
|
||
for in_idx in 0..in_features {
|
||
let idx = out_idx * in_features + in_idx;
|
||
let scale = scales.get(in_idx).copied().unwrap_or(1.0);
|
||
smoothed[idx] *= scale;
|
||
}
|
||
}
|
||
|
||
Tensor::from_slice(&smoothed, shape, &self.config.device).map_err(tensor_err)
|
||
}
|
||
|
||
/// Quantize a smoothed layer (both weights and activations)
|
||
pub fn quantize_layer(
|
||
&mut self,
|
||
layer_name: &str,
|
||
weights: &Tensor,
|
||
activations: &[Tensor],
|
||
) -> Result<SmoothQuantizedLayer> {
|
||
// First calibrate
|
||
self.calibrate(layer_name, activations, weights)?;
|
||
|
||
// Smooth weights
|
||
let smoothed_weights = self.smooth_weights(layer_name, weights)?;
|
||
|
||
// Quantize smoothed weights
|
||
let quantized_weights =
|
||
self.quantize_tensor(&smoothed_weights, self.config.weight_bit_width)?;
|
||
|
||
let scales = self
|
||
.smoothing_scales
|
||
.get(layer_name)
|
||
.cloned()
|
||
.unwrap_or_default();
|
||
|
||
Ok(SmoothQuantizedLayer {
|
||
quantized_weights,
|
||
smoothing_scales: scales,
|
||
weight_bit_width: self.config.weight_bit_width,
|
||
activation_bit_width: self.config.activation_bit_width,
|
||
})
|
||
}
|
||
|
||
/// Simple tensor quantization
|
||
fn quantize_tensor(&self, tensor: &Tensor, bit_width: u8) -> Result<QuantizedTensorData> {
|
||
let flat = tensor.to_vec().map_err(tensor_err)?;
|
||
|
||
// Find min/max
|
||
let min_val = flat.iter().copied().fold(f32::MAX, f32::min);
|
||
let max_val = flat.iter().copied().fold(f32::MIN, f32::max);
|
||
|
||
let qmin = -(1i32 << (bit_width - 1));
|
||
let qmax = (1i32 << (bit_width - 1)) - 1;
|
||
|
||
let scale = (max_val - min_val) / (qmax - qmin) as f32;
|
||
let scale = if scale == 0.0 { 1.0 } else { scale };
|
||
let zero = ((qmin as f32 * scale - min_val) / scale).round() as i8;
|
||
|
||
let quantized: Vec<i8> = flat
|
||
.iter()
|
||
.map(|&v| {
|
||
let q = ((v / scale) + zero as f32).round();
|
||
q.max(qmin as f32).min(qmax as f32) as i8
|
||
})
|
||
.collect();
|
||
|
||
Ok(QuantizedTensorData {
|
||
data: quantized,
|
||
scale,
|
||
zero_point: zero,
|
||
shape: tensor.shape().dims().to_vec(),
|
||
})
|
||
}
|
||
|
||
/// Get statistics
|
||
pub fn stats(&self) -> &SmoothQuantStats {
|
||
&self.stats
|
||
}
|
||
}
|
||
|
||
/// Simple quantized tensor data
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct QuantizedTensorData {
|
||
/// Quantized data
|
||
pub data: Vec<i8>,
|
||
/// Scale factor
|
||
pub scale: f32,
|
||
/// Zero point
|
||
pub zero_point: i8,
|
||
/// Original shape
|
||
pub shape: Vec<usize>,
|
||
}
|
||
|
||
impl QuantizedTensorData {
|
||
/// Dequantize
|
||
pub fn dequantize(&self, device: Device) -> Result<Tensor> {
|
||
let dequantized: Vec<f32> = self
|
||
.data
|
||
.iter()
|
||
.map(|&q| (q as f32 - self.zero_point as f32) * self.scale)
|
||
.collect();
|
||
|
||
Tensor::from_slice(&dequantized, &self.shape, &device)
|
||
.map_err(|e| QuantizationError::TensorError(e.to_string()).into())
|
||
}
|
||
}
|
||
|
||
/// SmoothQuant quantized layer
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct SmoothQuantizedLayer {
|
||
/// Quantized weights
|
||
pub quantized_weights: QuantizedTensorData,
|
||
/// Smoothing scales for activation preprocessing
|
||
pub smoothing_scales: Vec<f32>,
|
||
/// Weight bit width
|
||
pub weight_bit_width: u8,
|
||
/// Activation bit width
|
||
pub activation_bit_width: u8,
|
||
}
|
||
|
||
impl SmoothQuantizedLayer {
|
||
/// Get inverse smoothing scales (for activation preprocessing)
|
||
pub fn inverse_scales(&self) -> Vec<f32> {
|
||
self.smoothing_scales.iter().map(|&s| 1.0 / s).collect()
|
||
}
|
||
|
||
/// Number of output features (rows of the weight matrix).
|
||
#[inline]
|
||
pub fn out_features(&self) -> usize {
|
||
self.quantized_weights.shape[0]
|
||
}
|
||
|
||
/// Number of input features (columns of the weight matrix).
|
||
#[inline]
|
||
pub fn in_features(&self) -> usize {
|
||
self.quantized_weights.shape[1]
|
||
}
|
||
|
||
/// Run the quantized linear layer forward pass on a raw f32 buffer.
|
||
///
|
||
/// Implements the SmoothQuant inference path:
|
||
/// 1. **Smooth activations** – divide each channel by its smoothing scale, migrating
|
||
/// quantisation difficulty from activations to the pre-scaled weights.
|
||
/// 2. **Quantise activations to INT8** – compute a symmetric per-tensor scale
|
||
/// `act_scale = max(|smoothed|) / 127` and round to the range `[−127, 127]`.
|
||
/// 3. **INT8 GEMM** – compute `acc[b, out] = Σ_k act_q[b,k] * (weight_q[out,k] − zp)`
|
||
/// with i32 accumulation.
|
||
/// 4. **Dequantise** – `output[b, out] = acc[b, out] as f32 * act_scale * weight_scale`.
|
||
///
|
||
/// # INT8 range
|
||
///
|
||
/// Activations are clamped to **`−127…127`** (symmetric), matching the
|
||
/// `activation_bit_width = 8` convention used throughout this crate.
|
||
///
|
||
/// # Dequantisation formula
|
||
///
|
||
/// ```text
|
||
/// output[b, out] = (Σ_k act_q[b,k] * (weight_q[out,k] − zp)) * act_scale * weight_scale
|
||
/// ```
|
||
///
|
||
/// # Arguments
|
||
///
|
||
/// * `activations` – Row-major f32 slice of shape `[batch_size, in_features]`.
|
||
/// * `batch_size` – Number of input rows.
|
||
///
|
||
/// # Returns
|
||
///
|
||
/// Row-major `Vec<f32>` of shape `[batch_size, out_features]`.
|
||
///
|
||
/// # Errors
|
||
///
|
||
/// Returns an error if the activation buffer length is inconsistent with `batch_size`.
|
||
pub fn forward_raw(&self, activations: &[f32], batch_size: usize) -> Result<Vec<f32>> {
|
||
let in_f = self.in_features();
|
||
let out_f = self.out_features();
|
||
|
||
let expected_len = batch_size * in_f;
|
||
if activations.len() != expected_len {
|
||
return Err(CompressionError::Quantization(
|
||
QuantizationError::InvalidConfig(format!(
|
||
"SmoothQuantizedLayer::forward_raw: activation buffer length {} \
|
||
does not match batch_size={} × in_features={}",
|
||
activations.len(),
|
||
batch_size,
|
||
in_f
|
||
)),
|
||
));
|
||
}
|
||
|
||
let zp = self.quantized_weights.zero_point as i32;
|
||
let weight_scale = self.quantized_weights.scale;
|
||
|
||
// ------------------------------------------------------------------
|
||
// Step 1 + 2: smooth then quantise activations to INT8
|
||
//
|
||
// We compute smoothed values first (in place into a temporary buffer),
|
||
// then find max_abs and derive act_scale before quantising.
|
||
// ------------------------------------------------------------------
|
||
let mut smoothed = vec![0.0f32; batch_size * in_f];
|
||
for b in 0..batch_size {
|
||
let row = b * in_f;
|
||
for c in 0..in_f {
|
||
// Divide by the per-channel smoothing scale (≥ ε for numerical safety).
|
||
let s = self.smoothing_scales[c].max(f32::EPSILON);
|
||
smoothed[row + c] = activations[row + c] / s;
|
||
}
|
||
}
|
||
|
||
// Per-tensor activation scale: max absolute value over the whole batch.
|
||
let max_abs = smoothed.iter().map(|v| v.abs()).fold(0.0f32, f32::max);
|
||
|
||
// Guard against zero-tensor inputs; any non-zero scale works here.
|
||
let act_scale = if max_abs < f32::EPSILON {
|
||
1.0f32
|
||
} else {
|
||
max_abs / 127.0f32
|
||
};
|
||
|
||
// Quantise each smoothed value to INT8 in range [−127, 127].
|
||
let act_q: Vec<i8> = smoothed
|
||
.iter()
|
||
.map(|&v| {
|
||
let q = (v / act_scale).round();
|
||
q.clamp(-127.0, 127.0) as i8
|
||
})
|
||
.collect();
|
||
|
||
// ------------------------------------------------------------------
|
||
// Step 3: INT8 GEMM with i32 accumulation
|
||
//
|
||
// output_i32[b, out_row] = Σ_k act_q[b,k] * (weight_q[out_row,k] − zp)
|
||
// ------------------------------------------------------------------
|
||
let mut output = vec![0.0f32; batch_size * out_f];
|
||
for b in 0..batch_size {
|
||
let act_row = b * in_f;
|
||
for out_row in 0..out_f {
|
||
let weight_row = out_row * in_f;
|
||
let mut acc = 0i32;
|
||
for k in 0..in_f {
|
||
let wq = self.quantized_weights.data[weight_row + k] as i32 - zp;
|
||
let aq = act_q[act_row + k] as i32;
|
||
acc += aq * wq;
|
||
}
|
||
// Step 4: dequantise
|
||
output[b * out_f + out_row] = acc as f32 * act_scale * weight_scale;
|
||
}
|
||
}
|
||
|
||
Ok(output)
|
||
}
|
||
}
|
||
|
||
// =============================================================================
|
||
// Tests
|
||
// =============================================================================
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn test_awq_config_default() {
|
||
let config = AWQConfig::default();
|
||
assert_eq!(config.bit_width, 4);
|
||
assert_eq!(config.group_size, 128);
|
||
}
|
||
|
||
#[test]
|
||
fn test_awq_quantizer_basic() {
|
||
let config = AWQConfig::new(4).with_group_size(64);
|
||
let quantizer = AWQQuantizer::new(config);
|
||
assert_eq!(quantizer.stats().layers_quantized, 0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_gptq_config_default() {
|
||
let config = GPTQConfig::default();
|
||
assert_eq!(config.bit_width, 4);
|
||
assert_eq!(config.group_size, 128);
|
||
assert!(config.true_sequential);
|
||
}
|
||
|
||
#[test]
|
||
fn test_gptq_quantizer_basic() {
|
||
let config = GPTQConfig::new(4).with_act_order(true);
|
||
let quantizer = GPTQQuantizer::new(config);
|
||
assert!(quantizer.config.act_order);
|
||
}
|
||
|
||
#[test]
|
||
fn test_smoothquant_config_default() {
|
||
let config = SmoothQuantConfig::default();
|
||
assert_eq!(config.alpha, 0.5);
|
||
assert_eq!(config.weight_bit_width, 8);
|
||
assert_eq!(config.activation_bit_width, 8);
|
||
}
|
||
|
||
#[test]
|
||
fn test_smoothquant_alpha_clamping() {
|
||
let config1 = SmoothQuantConfig::new(1.5);
|
||
assert_eq!(config1.alpha, 1.0);
|
||
|
||
let config2 = SmoothQuantConfig::new(-0.5);
|
||
assert_eq!(config2.alpha, 0.0);
|
||
}
|
||
|
||
#[test]
|
||
fn test_quantized_tensor_dequantize() {
|
||
let data = QuantizedTensorData {
|
||
data: vec![0, 10, -10, 5],
|
||
scale: 0.1,
|
||
zero_point: 0,
|
||
shape: vec![2, 2],
|
||
};
|
||
|
||
let dequant = data.dequantize(Device::Cpu).unwrap();
|
||
let flat: Vec<f32> = dequant.to_vec().unwrap();
|
||
|
||
assert!((flat[0] - 0.0).abs() < 0.01);
|
||
assert!((flat[1] - 1.0).abs() < 0.01);
|
||
assert!((flat[2] - -1.0).abs() < 0.01);
|
||
assert!((flat[3] - 0.5).abs() < 0.01);
|
||
}
|
||
|
||
#[test]
|
||
fn test_awq_quantized_weight_memory() {
|
||
let weight = AWQQuantizedWeight {
|
||
quantized_data: vec![0; 1024],
|
||
scales: vec![1.0; 8],
|
||
zeros: vec![0; 8],
|
||
activation_scales: vec![1.0; 256],
|
||
group_size: 128,
|
||
bit_width: 4,
|
||
shape: vec![4, 256],
|
||
};
|
||
|
||
let bytes = weight.memory_bytes();
|
||
assert!(bytes > 0);
|
||
assert!(bytes < 1024 * 4 + 100); // Less than full FP32 + overhead
|
||
}
|
||
|
||
#[test]
|
||
fn test_gptq_compression_ratio() {
|
||
let weight = GPTQQuantizedWeight {
|
||
quantized_data: vec![0; 1024],
|
||
scales: vec![1.0; 8],
|
||
zeros: vec![0; 8],
|
||
group_size: 128,
|
||
bit_width: 4,
|
||
shape: vec![32, 32],
|
||
col_order: None,
|
||
};
|
||
|
||
let ratio = weight.compression_ratio();
|
||
assert!(ratio > 1.0); // Should be compressed
|
||
assert!(ratio < 10.0); // But not impossibly so
|
||
}
|
||
|
||
// -----------------------------------------------------------------------
|
||
// Helper: build a SmoothQuantizedLayer with fully controlled state
|
||
//
|
||
// Parameters:
|
||
// out_features, in_features — weight matrix shape
|
||
// weight_data — INT8 quantized weights, row-major [out, in]
|
||
// weight_scale — single global scale for the weight tensor
|
||
// weight_zp — zero point for weight dequantization
|
||
// smoothing — per-channel smoothing scales, len = in_features
|
||
// -----------------------------------------------------------------------
|
||
fn make_smooth_layer(
|
||
out_features: usize,
|
||
in_features: usize,
|
||
weight_data: Vec<i8>,
|
||
weight_scale: f32,
|
||
weight_zp: i8,
|
||
smoothing: Vec<f32>,
|
||
) -> SmoothQuantizedLayer {
|
||
SmoothQuantizedLayer {
|
||
quantized_weights: QuantizedTensorData {
|
||
data: weight_data,
|
||
scale: weight_scale,
|
||
zero_point: weight_zp,
|
||
shape: vec![out_features, in_features],
|
||
},
|
||
smoothing_scales: smoothing,
|
||
weight_bit_width: 8,
|
||
activation_bit_width: 8,
|
||
}
|
||
}
|
||
|
||
/// Output length must equal batch_size × out_features.
|
||
#[test]
|
||
fn test_smoothquant_forward_raw_shape() {
|
||
let out_f = 3usize;
|
||
let in_f = 4usize;
|
||
let layer = make_smooth_layer(
|
||
out_f,
|
||
in_f,
|
||
vec![0i8; out_f * in_f],
|
||
1.0,
|
||
0,
|
||
vec![1.0f32; in_f],
|
||
);
|
||
let activations = vec![0.0f32; 2 * in_f]; // batch=2
|
||
let out = layer.forward_raw(&activations, 2).unwrap();
|
||
assert_eq!(
|
||
out.len(),
|
||
2 * out_f,
|
||
"output length must be batch * out_features"
|
||
);
|
||
}
|
||
|
||
/// Identity layer (W=I, smoothing=1, zp=0, weight_scale=1) reproduces activations.
|
||
///
|
||
/// With weight_scale=1 and act_scale derived from the input, roundtrip quantisation
|
||
/// introduces < 1% relative error for values in [1, 127].
|
||
#[test]
|
||
fn test_smoothquant_forward_raw_identity_layer() {
|
||
// 4×4 identity weight matrix (INT8), zero-point=0, weight_scale=1.
|
||
// smoothing_scales = [1.0; 4] → no smoothing effect.
|
||
// Activations: [1, 2, 3, 4] (batch=1).
|
||
// Expected output ≈ [1, 2, 3, 4] (within quantisation tolerance).
|
||
let in_f = 4usize;
|
||
let out_f = 4usize;
|
||
let identity: Vec<i8> = vec![
|
||
1, 0, 0, 0, //
|
||
0, 1, 0, 0, //
|
||
0, 0, 1, 0, //
|
||
0, 0, 0, 1, //
|
||
];
|
||
let layer = make_smooth_layer(out_f, in_f, identity, 1.0, 0, vec![1.0f32; in_f]);
|
||
let activations = vec![1.0f32, 2.0, 3.0, 4.0];
|
||
let out = layer.forward_raw(&activations, 1).unwrap();
|
||
|
||
assert_eq!(out.len(), out_f);
|
||
// Tolerance: INT8 quantisation of activations loses at most 1/127 ≈ 0.8%.
|
||
for (i, (&computed, &expected)) in out.iter().zip(activations.iter()).enumerate() {
|
||
let rel_err = (computed - expected).abs() / (expected.abs().max(1e-6_f32));
|
||
assert!(
|
||
rel_err < 0.02,
|
||
"element {i}: got {computed:.4}, expected {expected:.4}, rel_err={rel_err:.4}"
|
||
);
|
||
}
|
||
}
|
||
|
||
/// smoothing_scales=[2.0]*n divides each activation channel by 2 before quantisation.
|
||
///
|
||
/// With W=I, this halves the output compared to smoothing_scales=[1.0].
|
||
#[test]
|
||
fn test_smoothquant_forward_raw_scale_effect() {
|
||
let in_f = 4usize;
|
||
let out_f = 4usize;
|
||
let identity: Vec<i8> = vec![
|
||
1, 0, 0, 0, //
|
||
0, 1, 0, 0, //
|
||
0, 0, 1, 0, //
|
||
0, 0, 0, 1, //
|
||
];
|
||
|
||
// Layer with smoothing_scales = [1.0; 4] (baseline)
|
||
let layer_no_smooth =
|
||
make_smooth_layer(out_f, in_f, identity.clone(), 1.0, 0, vec![1.0f32; in_f]);
|
||
// Layer with smoothing_scales = [2.0; 4] (halves activations)
|
||
let layer_smooth = make_smooth_layer(out_f, in_f, identity, 1.0, 0, vec![2.0f32; in_f]);
|
||
|
||
let activations = vec![10.0f32, 20.0, 30.0, 40.0];
|
||
|
||
let out_baseline = layer_no_smooth.forward_raw(&activations, 1).unwrap();
|
||
let out_halved = layer_smooth.forward_raw(&activations, 1).unwrap();
|
||
|
||
for (i, (&baseline, &halved)) in out_baseline.iter().zip(out_halved.iter()).enumerate() {
|
||
let ratio = baseline / halved;
|
||
assert!(
|
||
(ratio - 2.0).abs() < 0.05,
|
||
"element {i}: baseline={baseline:.4} / halved={halved:.4} = ratio {ratio:.4}, expected ~2.0"
|
||
);
|
||
}
|
||
}
|
||
|
||
/// Manually computed reference: 2×2 layer, known weights/scales/input.
|
||
///
|
||
/// Setup:
|
||
/// W (INT8) = [[2, 0], [0, 2]], weight_scale=0.5, weight_zp=0
|
||
/// smoothing_scales = [1.0, 1.0]
|
||
/// activations = [10.0, 20.0] (batch=1)
|
||
///
|
||
/// Step-by-step:
|
||
/// 1. Smooth: [10/1, 20/1] = [10.0, 20.0]
|
||
/// 2. act_scale = max(|10|, |20|) / 127 = 20/127 ≈ 0.15748
|
||
/// 3. Quantise: act_q = round([10/0.15748, 20/0.15748]).clamp(-127,127)
|
||
/// ≈ round([63.5, 127.0]) = [64, 127] (or [63, 127] depending on rounding)
|
||
/// 4. INT8 GEMM (weight_zp=0):
|
||
/// out[0] = 2*act_q[0] + 0*act_q[1] = 2*64 = 128 (or 2*63=126)
|
||
/// out[1] = 0*act_q[0] + 2*act_q[1] = 2*127 = 254
|
||
/// 5. Dequant: output * act_scale * weight_scale
|
||
/// out[0] ≈ 128 * 0.15748 * 0.5 ≈ 10.08 (expected ~10.0)
|
||
/// out[1] ≈ 254 * 0.15748 * 0.5 ≈ 20.0
|
||
///
|
||
/// Tolerance: 2% (quantisation rounding).
|
||
#[test]
|
||
fn test_smoothquant_forward_raw_matches_manual() {
|
||
let layer = make_smooth_layer(2, 2, vec![2i8, 0, 0, 2], 0.5, 0, vec![1.0f32, 1.0]);
|
||
let activations = vec![10.0f32, 20.0];
|
||
let out = layer.forward_raw(&activations, 1).unwrap();
|
||
|
||
assert_eq!(out.len(), 2);
|
||
assert!(
|
||
(out[0] - 10.0).abs() < 0.5,
|
||
"out[0] ≈ 10.0, got {:.4}",
|
||
out[0]
|
||
);
|
||
assert!(
|
||
(out[1] - 20.0).abs() < 0.5,
|
||
"out[1] ≈ 20.0, got {:.4}",
|
||
out[1]
|
||
);
|
||
}
|
||
}
|