torsh_nn/quantization/
schemes.rs

1//! Different quantization schemes and their implementations
2
3use crate::quantization::{CalibrationMethod, QuantizationParams, QuantizationScheme};
4use torsh_core::{
5    dtype::DType,
6    error::{Result, TorshError},
7};
8use torsh_tensor::Tensor;
9
10/// Post-training quantization (PTQ) implementation
11pub struct PostTrainingQuantization {
12    scheme: QuantizationScheme,
13    target_dtype: DType,
14    calibration_method: CalibrationMethod,
15}
16
17impl PostTrainingQuantization {
18    /// Create a new PTQ quantizer
19    pub fn new(
20        scheme: QuantizationScheme,
21        target_dtype: DType,
22        calibration_method: CalibrationMethod,
23    ) -> Self {
24        Self {
25            scheme,
26            target_dtype,
27            calibration_method,
28        }
29    }
30
31    /// Apply post-training quantization to a tensor
32    pub fn quantize_tensor(&self, tensor: &Tensor) -> Result<(Tensor, QuantizationParams)> {
33        let data = tensor.to_vec()?;
34        let params = self.calculate_params(&data)?;
35        let quantized = params.quantize(tensor)?;
36        Ok((quantized, params))
37    }
38
39    /// Calculate quantization parameters for the given data
40    fn calculate_params(&self, data: &[f32]) -> Result<QuantizationParams> {
41        let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
42        let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
43
44        match self.scheme {
45            QuantizationScheme::Symmetric => {
46                let scale = self.calculate_symmetric_scale(min_val, max_val)?;
47                Ok(QuantizationParams::symmetric(
48                    scale,
49                    DType::F32,
50                    self.target_dtype,
51                ))
52            }
53            QuantizationScheme::Asymmetric => {
54                let (scale, zero_point) = self.calculate_asymmetric_params(min_val, max_val)?;
55                Ok(QuantizationParams::asymmetric(
56                    scale,
57                    zero_point,
58                    DType::F32,
59                    self.target_dtype,
60                ))
61            }
62            QuantizationScheme::Dynamic => {
63                // Dynamic quantization uses runtime statistics
64                self.dynamic_quantization_params(data)
65            }
66            QuantizationScheme::KLDivergence => self.kl_divergence_params(data),
67            QuantizationScheme::Percentile(percentile) => self.percentile_params(data, percentile),
68        }
69    }
70
71    /// Calculate symmetric quantization scale
72    fn calculate_symmetric_scale(&self, min_val: f32, max_val: f32) -> Result<f32> {
73        let max_abs = max_val.abs().max(min_val.abs());
74        let max_quant = match self.target_dtype {
75            DType::I8 => 127.0,
76            DType::U8 => 127.0, // Use symmetric range for U8 too
77            DType::I16 => 32767.0,
78            _ => {
79                return Err(TorshError::InvalidArgument(
80                    "Unsupported quantization dtype".to_string(),
81                ))
82            }
83        };
84
85        if max_abs == 0.0 {
86            return Ok(1.0); // Avoid division by zero
87        }
88
89        Ok(max_abs / max_quant)
90    }
91
92    /// Calculate asymmetric quantization parameters
93    fn calculate_asymmetric_params(&self, min_val: f32, max_val: f32) -> Result<(f32, i32)> {
94        let (qmin, qmax) = match self.target_dtype {
95            DType::U8 => (0.0, 255.0),
96            DType::I8 => (-128.0, 127.0),
97            DType::I16 => (-32768.0, 32767.0),
98            _ => {
99                return Err(TorshError::InvalidArgument(
100                    "Unsupported quantization dtype".to_string(),
101                ))
102            }
103        };
104
105        let scale = (max_val - min_val) / (qmax - qmin);
106        let zero_point = (qmin - min_val / scale).round() as i32;
107
108        Ok((scale, zero_point))
109    }
110
111    /// Dynamic quantization parameters
112    fn dynamic_quantization_params(&self, data: &[f32]) -> Result<QuantizationParams> {
113        // For dynamic quantization, we calculate parameters based on runtime statistics
114        // This is similar to symmetric but might use different calibration methods
115        use crate::quantization::calibration::calculate_optimal_scale;
116
117        let scale = calculate_optimal_scale(data, &self.calibration_method, self.target_dtype)?;
118        Ok(QuantizationParams::symmetric(
119            scale,
120            DType::F32,
121            self.target_dtype,
122        ))
123    }
124
125    /// KL divergence-based quantization parameters
126    fn kl_divergence_params(&self, data: &[f32]) -> Result<QuantizationParams> {
127        let scale = self.find_optimal_scale_kl(data)?;
128        Ok(QuantizationParams::symmetric(
129            scale,
130            DType::F32,
131            self.target_dtype,
132        ))
133    }
134
135    /// Percentile-based quantization parameters
136    fn percentile_params(&self, data: &[f32], percentile: f32) -> Result<QuantizationParams> {
137        let mut sorted_abs: Vec<f32> = data.iter().map(|&x| x.abs()).collect();
138        sorted_abs.sort_by(|a, b| {
139            a.partial_cmp(b)
140                .expect("data comparison should not involve NaN")
141        });
142
143        let index = ((percentile / 100.0) * sorted_abs.len() as f32) as usize;
144        let max_val = sorted_abs
145            .get(index.min(sorted_abs.len() - 1))
146            .copied()
147            .unwrap_or(0.0);
148
149        let scale = self.calculate_symmetric_scale(-max_val, max_val)?;
150        Ok(QuantizationParams::symmetric(
151            scale,
152            DType::F32,
153            self.target_dtype,
154        ))
155    }
156
157    /// Find optimal scale using KL divergence
158    fn find_optimal_scale_kl(&self, data: &[f32]) -> Result<f32> {
159        let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b.abs()));
160        let max_quant = match self.target_dtype {
161            DType::I8 => 127.0,
162            DType::U8 => 255.0,
163            _ => {
164                return Err(TorshError::InvalidArgument(
165                    "Unsupported dtype for KL divergence".to_string(),
166                ))
167            }
168        };
169
170        let base_scale = max_val / max_quant;
171        let mut best_scale = base_scale;
172        let mut best_divergence = f32::INFINITY;
173
174        // Try different scales and find the one with minimum KL divergence
175        for multiplier in [0.5, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2, 1.5, 2.0] {
176            let scale = base_scale * multiplier;
177            let divergence = self.calculate_kl_divergence(data, scale)?;
178
179            if divergence < best_divergence {
180                best_divergence = divergence;
181                best_scale = scale;
182            }
183        }
184
185        Ok(best_scale)
186    }
187
188    /// Calculate KL divergence for a given scale
189    fn calculate_kl_divergence(&self, data: &[f32], scale: f32) -> Result<f32> {
190        let max_quant = match self.target_dtype {
191            DType::I8 => 127,
192            DType::U8 => 255,
193            _ => return Err(TorshError::InvalidArgument("Unsupported dtype".to_string())),
194        };
195
196        // Create histogram of original data
197        let num_bins = 256;
198        let min_val = data.iter().fold(f32::INFINITY, |a, &b| a.min(b));
199        let max_val = data.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
200        let bin_width = (max_val - min_val) / num_bins as f32;
201
202        let mut original_hist = vec![0u32; num_bins];
203        let mut quantized_hist = vec![0u32; num_bins];
204
205        for &value in data {
206            // Original histogram
207            let bin_idx = ((value - min_val) / bin_width) as usize;
208            let bin_idx = bin_idx.min(num_bins - 1);
209            original_hist[bin_idx] += 1;
210
211            // Quantized histogram
212            let quantized = ((value / scale).round() as i32).clamp(-max_quant, max_quant);
213            let dequantized = quantized as f32 * scale;
214            let quant_bin_idx = ((dequantized - min_val) / bin_width) as usize;
215            let quant_bin_idx = quant_bin_idx.min(num_bins - 1);
216            quantized_hist[quant_bin_idx] += 1;
217        }
218
219        // Calculate KL divergence
220        let total_count = data.len() as f32;
221        let mut kl_div = 0.0;
222
223        for i in 0..num_bins {
224            let p = original_hist[i] as f32 / total_count;
225            let q = quantized_hist[i] as f32 / total_count;
226
227            if p > 0.0 && q > 0.0 {
228                kl_div += p * (p / q).ln();
229            }
230        }
231
232        Ok(kl_div)
233    }
234}
235
236/// Quantization-Aware Training (QAT) implementation
237pub struct QuantizationAwareTraining {
238    fake_quantize: bool,
239    #[allow(dead_code)]
240    scheme: QuantizationScheme,
241    target_dtype: DType,
242}
243
244impl QuantizationAwareTraining {
245    /// Create a new QAT quantizer
246    pub fn new(scheme: QuantizationScheme, target_dtype: DType) -> Self {
247        Self {
248            fake_quantize: true,
249            scheme,
250            target_dtype,
251        }
252    }
253
254    /// Apply fake quantization during training
255    pub fn fake_quantize_tensor(
256        &self,
257        tensor: &Tensor,
258        params: &QuantizationParams,
259    ) -> Result<Tensor> {
260        if !self.fake_quantize {
261            return Ok(tensor.clone());
262        }
263
264        // Simulate quantization during forward pass
265        let quantized = params.quantize(tensor)?;
266        let dequantized = params.dequantize(&quantized)?;
267
268        Ok(dequantized)
269    }
270
271    /// Enable/disable fake quantization
272    pub fn set_fake_quantize(&mut self, enabled: bool) {
273        self.fake_quantize = enabled;
274    }
275
276    /// Create learnable quantization parameters
277    pub fn learnable_params(
278        &self,
279        initial_scale: f32,
280        initial_zero_point: i32,
281    ) -> LearnableQuantParams {
282        LearnableQuantParams::new(initial_scale, initial_zero_point, self.target_dtype)
283    }
284}
285
286/// Learnable quantization parameters for QAT
287#[derive(Debug, Clone)]
288pub struct LearnableQuantParams {
289    scale: Tensor,
290    zero_point: Tensor,
291    target_dtype: DType,
292}
293
294impl LearnableQuantParams {
295    /// Create new learnable parameters
296    pub fn new(initial_scale: f32, initial_zero_point: i32, target_dtype: DType) -> Self {
297        let scale = Tensor::from_data(
298            vec![initial_scale],
299            vec![1],
300            torsh_core::device::DeviceType::Cpu,
301        )
302        .unwrap_or_else(|_| {
303            Tensor::zeros(&[1], torsh_core::device::DeviceType::Cpu)
304                .expect("Failed to create fallback zeros tensor")
305        });
306        let zero_point = Tensor::from_data(
307            vec![initial_zero_point as f32],
308            vec![1],
309            torsh_core::device::DeviceType::Cpu,
310        )
311        .unwrap_or_else(|_| {
312            Tensor::zeros(&[1], torsh_core::device::DeviceType::Cpu)
313                .expect("Failed to create fallback zeros tensor")
314        });
315
316        Self {
317            scale,
318            zero_point,
319            target_dtype,
320        }
321    }
322
323    /// Get current quantization parameters
324    pub fn current_params(&self) -> Result<QuantizationParams> {
325        let scale_val = self.scale.to_vec()?[0];
326        let zero_point_val = self.zero_point.to_vec()?[0] as i32;
327
328        Ok(QuantizationParams::asymmetric(
329            scale_val,
330            zero_point_val,
331            DType::F32,
332            self.target_dtype,
333        ))
334    }
335
336    /// Update parameters during training
337    pub fn update_scale(&mut self, new_scale: f32) -> Result<()> {
338        self.scale = Tensor::from_data(
339            vec![new_scale],
340            vec![1],
341            torsh_core::device::DeviceType::Cpu,
342        )?;
343        Ok(())
344    }
345
346    /// Update zero point during training
347    pub fn update_zero_point(&mut self, new_zero_point: i32) -> Result<()> {
348        self.zero_point = Tensor::from_data(
349            vec![new_zero_point as f32],
350            vec![1],
351            torsh_core::device::DeviceType::Cpu,
352        )?;
353        Ok(())
354    }
355}
356
357/// Block-wise quantization for large models
358pub struct BlockWiseQuantization {
359    block_size: usize,
360    scheme: QuantizationScheme,
361    target_dtype: DType,
362}
363
364impl BlockWiseQuantization {
365    /// Create a new block-wise quantizer
366    pub fn new(block_size: usize, scheme: QuantizationScheme, target_dtype: DType) -> Self {
367        Self {
368            block_size,
369            scheme,
370            target_dtype,
371        }
372    }
373
374    /// Quantize tensor in blocks
375    pub fn quantize_blocks(&self, tensor: &Tensor) -> Result<(Tensor, Vec<QuantizationParams>)> {
376        let data = tensor.to_vec()?;
377        let mut quantized_data = Vec::new();
378        let mut block_params = Vec::new();
379
380        for chunk in data.chunks(self.block_size) {
381            let ptq = PostTrainingQuantization::new(
382                self.scheme.clone(),
383                self.target_dtype,
384                CalibrationMethod::MinMax,
385            );
386
387            let block_tensor = Tensor::from_data(
388                chunk.to_vec(),
389                vec![chunk.len()],
390                torsh_core::device::DeviceType::Cpu,
391            )?;
392            let (quantized_block, params) = ptq.quantize_tensor(&block_tensor)?;
393
394            let block_data: Vec<u8> = match self.target_dtype {
395                DType::I8 => quantized_block
396                    .to_vec()?
397                    .into_iter()
398                    .map(|x: f32| x as i8 as u8)
399                    .collect(),
400                DType::U8 => quantized_block
401                    .to_vec()?
402                    .into_iter()
403                    .map(|x: f32| x as u8)
404                    .collect(),
405                _ => {
406                    return Err(TorshError::InvalidArgument(
407                        "Unsupported block quantization dtype".to_string(),
408                    ))
409                }
410            };
411
412            quantized_data.extend(block_data);
413            block_params.push(params);
414        }
415
416        let float_data: Vec<f32> = quantized_data.into_iter().map(|x| x as f32).collect();
417        let quantized_tensor = Tensor::from_data(
418            float_data,
419            tensor.shape().dims().to_vec(),
420            torsh_core::device::DeviceType::Cpu,
421        )?;
422        Ok((quantized_tensor, block_params))
423    }
424}
425
426/// Mixed precision quantization
427pub struct MixedPrecisionQuantization {
428    layer_configs: Vec<LayerQuantConfig>,
429}
430
431/// Configuration for per-layer quantization
432#[derive(Debug, Clone)]
433pub struct LayerQuantConfig {
434    pub layer_name: String,
435    pub weight_dtype: DType,
436    pub activation_dtype: DType,
437    pub scheme: QuantizationScheme,
438    pub per_channel: bool,
439}
440
441impl MixedPrecisionQuantization {
442    /// Create a new mixed precision quantizer
443    pub fn new(layer_configs: Vec<LayerQuantConfig>) -> Self {
444        Self { layer_configs }
445    }
446
447    /// Get quantization config for a specific layer
448    pub fn get_layer_config(&self, layer_name: &str) -> Option<&LayerQuantConfig> {
449        self.layer_configs
450            .iter()
451            .find(|config| config.layer_name == layer_name)
452    }
453
454    /// Apply layer-specific quantization
455    pub fn quantize_layer(
456        &self,
457        layer_name: &str,
458        weights: &Tensor,
459        activations: &Tensor,
460    ) -> Result<(Tensor, Tensor, QuantizationParams, QuantizationParams)> {
461        let config = self.get_layer_config(layer_name).ok_or_else(|| {
462            TorshError::InvalidArgument(format!(
463                "No quantization config found for layer: {}",
464                layer_name
465            ))
466        })?;
467
468        // Quantize weights
469        let weight_ptq = PostTrainingQuantization::new(
470            config.scheme.clone(),
471            config.weight_dtype,
472            CalibrationMethod::MinMax,
473        );
474        let (quantized_weights, weight_params) = weight_ptq.quantize_tensor(weights)?;
475
476        // Quantize activations
477        let activation_ptq = PostTrainingQuantization::new(
478            config.scheme.clone(),
479            config.activation_dtype,
480            CalibrationMethod::MinMax,
481        );
482        let (quantized_activations, activation_params) =
483            activation_ptq.quantize_tensor(activations)?;
484
485        Ok((
486            quantized_weights,
487            quantized_activations,
488            weight_params,
489            activation_params,
490        ))
491    }
492}
493
494/// Utility functions for quantization schemes
495pub mod utils {
496    use super::*;
497
498    /// Create a simple INT8 symmetric quantization scheme
499    pub fn int8_symmetric() -> PostTrainingQuantization {
500        PostTrainingQuantization::new(
501            QuantizationScheme::Symmetric,
502            DType::I8,
503            CalibrationMethod::MinMax,
504        )
505    }
506
507    /// Create a UINT8 asymmetric quantization scheme
508    pub fn uint8_asymmetric() -> PostTrainingQuantization {
509        PostTrainingQuantization::new(
510            QuantizationScheme::Asymmetric,
511            DType::U8,
512            CalibrationMethod::MinMax,
513        )
514    }
515
516    /// Create a dynamic quantization scheme
517    pub fn dynamic_int8() -> PostTrainingQuantization {
518        PostTrainingQuantization::new(
519            QuantizationScheme::Dynamic,
520            DType::I8,
521            CalibrationMethod::Entropy,
522        )
523    }
524
525    /// Create a KL divergence-based quantization scheme
526    pub fn kl_divergence_int8() -> PostTrainingQuantization {
527        PostTrainingQuantization::new(
528            QuantizationScheme::KLDivergence,
529            DType::I8,
530            CalibrationMethod::Entropy,
531        )
532    }
533
534    /// Create a percentile-based quantization scheme
535    pub fn percentile_int8(percentile: f32) -> PostTrainingQuantization {
536        PostTrainingQuantization::new(
537            QuantizationScheme::Percentile(percentile),
538            DType::I8,
539            CalibrationMethod::MinMax,
540        )
541    }
542}
543
544#[cfg(test)]
545mod tests {
546    use super::*;
547
548    #[test]
549    fn test_post_training_quantization() {
550        let data = vec![1.0, -1.0, 0.5, -0.5, 2.0, -2.0];
551        let tensor = Tensor::from_data(data, vec![6], torsh_core::device::DeviceType::Cpu)
552            .expect("Tensor should succeed");
553
554        let ptq = PostTrainingQuantization::new(
555            QuantizationScheme::Symmetric,
556            DType::I8,
557            CalibrationMethod::MinMax,
558        );
559
560        let result = ptq.quantize_tensor(&tensor);
561        assert!(result.is_ok());
562
563        let (_quantized, params) = result.expect("operation should succeed");
564        // Note: quantized tensor will have F32 dtype due to our implementation
565        assert!(params.scale > 0.0);
566    }
567
568    #[test]
569    fn test_qat_fake_quantization() {
570        let tensor_data = vec![0.1f32; 10];
571        let tensor = Tensor::from_data(tensor_data, vec![10], torsh_core::device::DeviceType::Cpu)
572            .expect("Tensor should succeed");
573        let params = QuantizationParams::symmetric(0.1, DType::F32, DType::I8);
574
575        let qat = QuantizationAwareTraining::new(QuantizationScheme::Symmetric, DType::I8);
576        let fake_quantized = qat
577            .fake_quantize_tensor(&tensor, &params)
578            .expect("fake quantization should succeed");
579
580        assert_eq!(fake_quantized.shape().dims(), tensor.shape().dims());
581        // Note: tensor dtype will be F32 due to our implementation
582    }
583
584    #[test]
585    fn test_block_wise_quantization() {
586        let data: Vec<f32> = (0..100).map(|x| x as f32 / 10.0).collect();
587        let tensor = Tensor::from_data(data, vec![100], torsh_core::device::DeviceType::Cpu)
588            .expect("Tensor should succeed");
589
590        let block_quant = BlockWiseQuantization::new(25, QuantizationScheme::Symmetric, DType::I8);
591
592        let result = block_quant.quantize_blocks(&tensor);
593        assert!(result.is_ok());
594
595        let (quantized, params) = result.expect("operation should succeed");
596        assert_eq!(quantized.shape().dims(), tensor.shape().dims());
597        assert_eq!(params.len(), 4); // 100 / 25 = 4 blocks
598    }
599
600    #[test]
601    fn test_mixed_precision_quantization() {
602        let layer_configs = vec![
603            LayerQuantConfig {
604                layer_name: "conv1".to_string(),
605                weight_dtype: DType::I8,
606                activation_dtype: DType::I8,
607                scheme: QuantizationScheme::Symmetric,
608                per_channel: true,
609            },
610            LayerQuantConfig {
611                layer_name: "fc1".to_string(),
612                weight_dtype: DType::I8,
613                activation_dtype: DType::U8,
614                scheme: QuantizationScheme::Asymmetric,
615                per_channel: false,
616            },
617        ];
618
619        let mixed_prec = MixedPrecisionQuantization::new(layer_configs);
620        let config = mixed_prec.get_layer_config("conv1");
621        assert!(config.is_some());
622        assert_eq!(
623            config.expect("operation should succeed").weight_dtype,
624            DType::I8
625        );
626    }
627}
torsh_nn/quantization/schemes.rs

torsh_nn/quantization/
schemes.rs
