trustformers_wasm/optimization/quantization/
config.rs

1//! Quantization configuration and supporting types
2
3use serde::{Deserialize, Serialize};
4use std::vec::Vec;
5use wasm_bindgen::prelude::*;
6
7/// Quantization strategies
8#[wasm_bindgen]
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
10pub enum QuantizationStrategy {
11    /// No quantization
12    None,
13    /// Dynamic quantization - quantize weights only
14    Dynamic,
15    /// Static quantization - quantize weights and activations
16    Static,
17    /// Post-training quantization
18    PostTraining,
19    /// Quantization-aware training (requires pre-quantized model)
20    QAT,
21    /// AWQ (Activation-aware Weight Quantization) - preserves important weights
22    AWQ,
23    /// GPTQ (Gradient-based Post-Training Quantization) - uses second-order information
24    GPTQ,
25    /// SmoothQuant - balances weights and activations difficulty
26    SmoothQuant,
27    /// LLM.int8() - mixed-precision quantization for large models
28    LLMInt8,
29    /// QLoRA - Quantized Low-Rank Adaptation
30    QLoRA,
31    /// GGML-style quantization for efficient inference
32    GGML,
33    /// Adaptive bitwidth quantization with dynamic allocation
34    AdaptiveBitwidth,
35    /// Outlier-aware quantization for handling activation spikes
36    OutlierAware,
37    /// HQQ (Half-Quadratic Quantization) - superior quality quantization using half-quadratic optimization
38    HQQ,
39    /// SpQR (Sparse-Quantized Representation) - ultra-sparse models with mixed precision
40    SpQR,
41    /// AQLM (Additive Quantization for Language Models) - additive quantization for transformers
42    AQLM,
43}
44
45/// Quantization precision levels
46#[wasm_bindgen]
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
48pub enum QuantizationPrecision {
49    /// 16-bit floating point
50    FP16,
51    /// 8-bit floating point (E4M3 or E5M2 format)
52    FP8,
53    /// 8-bit integer
54    INT8,
55    /// 4-bit integer
56    INT4,
57    /// 2-bit integer (experimental)
58    INT2,
59    /// 1-bit binary quantization
60    INT1,
61    /// Mixed precision (FP16 for outliers, INT8 for normal weights)
62    Mixed,
63    /// Adaptive precision based on layer importance
64    Adaptive,
65}
66
67impl QuantizationPrecision {
68    /// Number of bits used to represent one quantized element for this precision.
69    ///
70    /// `Mixed` (FP16 for a minority of outlier weights, INT8 for the rest) and `Adaptive`
71    /// (a runtime-selected, per-layer bit-width) do not have one single fixed bit-width by
72    /// design. Both nominally report 8 bits here as a documented approximation used for
73    /// size/statistics purposes: for `Mixed`, INT8 is the majority-path width (outliers are
74    /// expected to be a small minority of elements); for `Adaptive`, INT8 is the
75    /// default/starting width used before any runtime adaptation narrows individual layers
76    /// further. This is intentionally explicit (not a silently-defaulted wildcard arm).
77    pub fn bits(&self) -> u32 {
78        match self {
79            QuantizationPrecision::FP16 => 16,
80            QuantizationPrecision::FP8 | QuantizationPrecision::INT8 => 8,
81            QuantizationPrecision::INT4 => 4,
82            QuantizationPrecision::INT2 => 2,
83            QuantizationPrecision::INT1 => 1,
84            QuantizationPrecision::Mixed | QuantizationPrecision::Adaptive => 8,
85        }
86    }
87
88    /// Storage size in bytes per element (fractional for sub-byte widths like INT4/INT2/INT1).
89    pub fn bytes_per_element(&self) -> f32 {
90        self.bits() as f32 / 8.0
91    }
92}
93
94/// Quantization configuration
95#[wasm_bindgen]
96#[derive(Debug, Clone)]
97pub struct QuantizationConfig {
98    strategy: QuantizationStrategy,
99    precision: QuantizationPrecision,
100    target_size_mb: f32,
101    performance_threshold: f32,
102    accuracy_threshold: f32,
103    auto_select: bool,
104}
105
106#[wasm_bindgen]
107impl QuantizationConfig {
108    /// Create a new quantization configuration
109    #[wasm_bindgen(constructor)]
110    pub fn new(strategy: QuantizationStrategy, precision: QuantizationPrecision) -> Self {
111        Self {
112            strategy,
113            precision,
114            target_size_mb: 50.0,       // Default target size
115            performance_threshold: 2.0, // 2x speedup minimum
116            accuracy_threshold: 0.95,   // 95% accuracy retention minimum
117            auto_select: false,
118        }
119    }
120
121    /// Create an automatic configuration that selects best settings
122    pub fn auto() -> Self {
123        Self {
124            strategy: QuantizationStrategy::Dynamic,
125            precision: QuantizationPrecision::INT8,
126            target_size_mb: 10.0,
127            performance_threshold: 1.5,
128            accuracy_threshold: 0.90,
129            auto_select: true,
130        }
131    }
132
133    /// Create a configuration optimized for mobile devices
134    pub fn mobile() -> Self {
135        Self {
136            strategy: QuantizationStrategy::PostTraining,
137            precision: QuantizationPrecision::INT8,
138            target_size_mb: 5.0,
139            performance_threshold: 3.0,
140            accuracy_threshold: 0.85,
141            auto_select: false,
142        }
143    }
144
145    /// Create a configuration for desktop/high-performance devices
146    pub fn desktop() -> Self {
147        Self {
148            strategy: QuantizationStrategy::Dynamic,
149            precision: QuantizationPrecision::FP16,
150            target_size_mb: 100.0,
151            performance_threshold: 1.2,
152            accuracy_threshold: 0.98,
153            auto_select: false,
154        }
155    }
156
157    /// Create a configuration for ultra-low latency inference
158    pub fn ultra_fast() -> Self {
159        Self {
160            strategy: QuantizationStrategy::GGML,
161            precision: QuantizationPrecision::FP8,
162            target_size_mb: 15.0,
163            performance_threshold: 4.0,
164            accuracy_threshold: 0.88,
165            auto_select: false,
166        }
167    }
168
169    /// Create a configuration for fine-tuning with QLoRA
170    pub fn qlora() -> Self {
171        Self {
172            strategy: QuantizationStrategy::QLoRA,
173            precision: QuantizationPrecision::Mixed,
174            target_size_mb: 8.0,
175            performance_threshold: 2.5,
176            accuracy_threshold: 0.92,
177            auto_select: false,
178        }
179    }
180
181    /// Create a configuration with adaptive bitwidth for optimal efficiency
182    pub fn adaptive() -> Self {
183        Self {
184            strategy: QuantizationStrategy::AdaptiveBitwidth,
185            precision: QuantizationPrecision::Adaptive,
186            target_size_mb: 12.0,
187            performance_threshold: 3.0,
188            accuracy_threshold: 0.93,
189            auto_select: true,
190        }
191    }
192
193    /// Create a configuration for models with activation outliers
194    pub fn outlier_aware() -> Self {
195        Self {
196            strategy: QuantizationStrategy::OutlierAware,
197            precision: QuantizationPrecision::Mixed,
198            target_size_mb: 20.0,
199            performance_threshold: 2.0,
200            accuracy_threshold: 0.96,
201            auto_select: false,
202        }
203    }
204
205    /// Set target model size in MB
206    pub fn set_target_size_mb(mut self, size_mb: f32) -> Self {
207        self.target_size_mb = size_mb;
208        self
209    }
210
211    /// Set performance threshold (minimum speedup factor)
212    pub fn set_performance_threshold(mut self, threshold: f32) -> Self {
213        self.performance_threshold = threshold;
214        self
215    }
216
217    /// Set accuracy threshold (minimum accuracy retention)
218    pub fn set_accuracy_threshold(mut self, threshold: f32) -> Self {
219        self.accuracy_threshold = threshold;
220        self
221    }
222
223    /// Enable automatic strategy selection
224    pub fn enable_auto_select(mut self) -> Self {
225        self.auto_select = true;
226        self
227    }
228
229    // Getters for private fields
230    #[wasm_bindgen(getter)]
231    pub fn strategy(&self) -> QuantizationStrategy {
232        self.strategy
233    }
234
235    #[wasm_bindgen(getter)]
236    pub fn precision(&self) -> QuantizationPrecision {
237        self.precision
238    }
239
240    #[wasm_bindgen(getter)]
241    pub fn target_size_mb(&self) -> f32 {
242        self.target_size_mb
243    }
244
245    #[wasm_bindgen(getter)]
246    pub fn performance_threshold(&self) -> f32 {
247        self.performance_threshold
248    }
249
250    #[wasm_bindgen(getter)]
251    pub fn accuracy_threshold(&self) -> f32 {
252        self.accuracy_threshold
253    }
254
255    #[wasm_bindgen(getter)]
256    pub fn auto_select(&self) -> bool {
257        self.auto_select
258    }
259}
260
261/// Quantization statistics
262#[wasm_bindgen]
263#[derive(Debug, Clone, Serialize, Deserialize)]
264pub struct QuantizationStats {
265    original_size_bytes: usize,
266    quantized_size_bytes: usize,
267    compression_ratio: f32,
268    size_reduction_percent: f32,
269    estimated_speedup: f32,
270    strategy_used: QuantizationStrategy,
271    precision_used: QuantizationPrecision,
272}
273
274#[wasm_bindgen]
275impl QuantizationStats {
276    #[wasm_bindgen(getter)]
277    pub fn original_size_bytes(&self) -> usize {
278        self.original_size_bytes
279    }
280
281    #[wasm_bindgen(getter)]
282    pub fn quantized_size_bytes(&self) -> usize {
283        self.quantized_size_bytes
284    }
285
286    #[wasm_bindgen(getter)]
287    pub fn compression_ratio(&self) -> f32 {
288        self.compression_ratio
289    }
290
291    #[wasm_bindgen(getter)]
292    pub fn size_reduction_percent(&self) -> f32 {
293        self.size_reduction_percent
294    }
295
296    #[wasm_bindgen(getter)]
297    pub fn estimated_speedup(&self) -> f32 {
298        self.estimated_speedup
299    }
300
301    #[wasm_bindgen(getter)]
302    pub fn strategy_used(&self) -> QuantizationStrategy {
303        self.strategy_used
304    }
305
306    #[wasm_bindgen(getter)]
307    pub fn precision_used(&self) -> QuantizationPrecision {
308        self.precision_used
309    }
310}
311
312impl QuantizationStats {
313    /// Create new quantization statistics
314    pub fn new(
315        original_size_bytes: usize,
316        quantized_size_bytes: usize,
317        compression_ratio: f32,
318        size_reduction_percent: f32,
319        estimated_speedup: f32,
320        strategy_used: QuantizationStrategy,
321        precision_used: QuantizationPrecision,
322    ) -> Self {
323        Self {
324            original_size_bytes,
325            quantized_size_bytes,
326            compression_ratio,
327            size_reduction_percent,
328            estimated_speedup,
329            strategy_used,
330            precision_used,
331        }
332    }
333}
334
335/// Runtime performance monitoring for adaptive quantization
336#[derive(Debug, Clone)]
337pub struct RuntimeMonitor {
338    pub inference_times: Vec<f64>,
339    pub memory_usage: Vec<usize>,
340    pub accuracy_scores: Vec<f32>,
341    pub thermal_state: ThermalState,
342    pub adaptation_history: Vec<AdaptationEvent>,
343}
344
345/// Current device thermal state
346#[derive(Debug, Clone, Copy, PartialEq, Eq)]
347pub enum ThermalState {
348    Nominal,  // Normal operating temperature
349    Fair,     // Slightly elevated
350    Serious,  // High temperature, throttling recommended
351    Critical, // Very high temperature, aggressive throttling needed
352}
353
354/// Adaptive quantization state that changes at runtime
355#[derive(Debug, Clone)]
356pub struct AdaptiveQuantizationState {
357    pub current_strategy: QuantizationStrategy,
358    pub current_precision: QuantizationPrecision,
359    pub adaptation_rate: f32,
360    pub performance_target: f32,
361    pub accuracy_target: f32,
362    pub last_adaptation: f64,
363    pub confidence_score: f32,
364}
365
366/// Record of quantization adaptations
367#[derive(Debug, Clone)]
368pub struct AdaptationEvent {
369    pub timestamp: f64,
370    pub trigger: AdaptationTrigger,
371    pub old_strategy: QuantizationStrategy,
372    pub new_strategy: QuantizationStrategy,
373    pub old_precision: QuantizationPrecision,
374    pub new_precision: QuantizationPrecision,
375    pub improvement_ratio: f32,
376}
377
378/// What triggered the adaptation
379#[derive(Debug, Clone, Copy, PartialEq, Eq)]
380pub enum AdaptationTrigger {
381    PerformanceDrop,     // Inference too slow
382    MemoryPressure,      // Running out of memory
383    AccuracyDrop,        // Model accuracy below threshold
384    ThermalThrottling,   // Device overheating
385    BatteryOptimization, // Low battery, need efficiency
386    WorkloadChange,      // Different type of inference requests
387}
388
389/// Device capabilities for quantization optimization
390#[derive(Debug, Clone)]
391pub struct DeviceCapabilities {
392    pub supports_int8: bool,
393    pub supports_int4: bool,
394    pub supports_fp16: bool,
395    pub memory_bandwidth_gb_s: f32,
396    pub compute_capability: ComputeCapability,
397}
398
399#[derive(Debug, Clone, Copy, PartialEq, Eq)]
400pub enum ComputeCapability {
401    Low,    // Basic CPU
402    Medium, // High-end CPU or integrated GPU
403    High,   // Dedicated GPU
404}
405
406#[cfg(test)]
407mod tests {
408    use super::*;
409
410    #[test]
411    fn test_bits_fp16() {
412        assert_eq!(QuantizationPrecision::FP16.bits(), 16);
413    }
414
415    #[test]
416    fn test_bits_fp8() {
417        assert_eq!(QuantizationPrecision::FP8.bits(), 8);
418    }
419
420    #[test]
421    fn test_bits_int8() {
422        assert_eq!(QuantizationPrecision::INT8.bits(), 8);
423    }
424
425    #[test]
426    fn test_bits_int4() {
427        assert_eq!(QuantizationPrecision::INT4.bits(), 4);
428    }
429
430    #[test]
431    fn test_bits_int2() {
432        assert_eq!(QuantizationPrecision::INT2.bits(), 2);
433    }
434
435    #[test]
436    fn test_bits_int1() {
437        assert_eq!(QuantizationPrecision::INT1.bits(), 1);
438    }
439
440    #[test]
441    fn test_bits_mixed_is_nominally_eight() {
442        assert_eq!(QuantizationPrecision::Mixed.bits(), 8);
443    }
444
445    #[test]
446    fn test_bits_adaptive_is_nominally_eight() {
447        assert_eq!(QuantizationPrecision::Adaptive.bits(), 8);
448    }
449
450    #[test]
451    fn test_bytes_per_element_int8_is_one_byte() {
452        assert_eq!(QuantizationPrecision::INT8.bytes_per_element(), 1.0);
453    }
454
455    #[test]
456    fn test_bytes_per_element_int4_is_half_byte() {
457        assert_eq!(QuantizationPrecision::INT4.bytes_per_element(), 0.5);
458    }
459
460    #[test]
461    fn test_bytes_per_element_int1_is_eighth_byte() {
462        assert_eq!(QuantizationPrecision::INT1.bytes_per_element(), 0.125);
463    }
464
465    #[test]
466    fn test_bytes_per_element_fp16_is_two_bytes() {
467        assert_eq!(QuantizationPrecision::FP16.bytes_per_element(), 2.0);
468    }
469}
trustformers_wasm/optimization/quantization/config.rs

trustformers_wasm/optimization/quantization/
config.rs
