ohms_adaptq/novaq/
mod.rs

1use crate::Result;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4
5pub mod normalization;
6pub mod codebooks;
7pub mod refinement;
8pub mod distillation;
9pub mod numerical_stability;
10pub mod subspace_strategy;
11pub mod recovery;
12pub mod progress;
13
14pub use normalization::*;
15pub use codebooks::*;
16pub use refinement::*;
17pub use distillation::*;
18pub use numerical_stability::*;
19pub use subspace_strategy::*;
20pub use recovery::*;
21pub use progress::*;
22
23/// NOVAQ Configuration
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct NOVAQConfig {
26    /// Target bits per weight (achieves ~1.5 bits effective precision)
27    pub target_bits: f32,
28    /// Number of vector subspaces for codebook quantization
29    pub num_subspaces: usize,
30    /// Size of first-level codebook (K1)
31    pub codebook_size_l1: usize,
32    /// Size of second-level residual codebook (K2)
33    pub codebook_size_l2: usize,
34    /// Top-p percentage for outlier channel identification
35    pub outlier_threshold: f32,
36    /// Teacher model path for knowledge distillation
37    pub teacher_model_path: Option<String>,
38    /// Number of refinement iterations
39    pub refinement_iterations: usize,
40    /// KL divergence weight in distillation loss
41    pub kl_weight: f32,
42    /// Cosine similarity weight in distillation loss
43    pub cosine_weight: f32,
44    /// Learning rate for centroid optimization
45    pub learning_rate: f32,
46    /// Random seed for reproducibility
47    pub seed: u64,
48}
49
50impl Default for NOVAQConfig {
51    fn default() -> Self {
52        Self {
53            target_bits: 1.5,
54            num_subspaces: 4,
55            codebook_size_l1: 16,  // K1=16 -> 4 bits
56            codebook_size_l2: 4,   // K2=4 -> 2 bits
57            outlier_threshold: 0.01, // Top 1%
58            teacher_model_path: None,
59            refinement_iterations: 100,
60            kl_weight: 1.0,
61            cosine_weight: 0.5,
62            learning_rate: 0.001,
63            seed: 42,
64        }
65    }
66}
67
68/// Weight matrix with shape information
69#[derive(Debug, Clone)]
70pub struct WeightMatrix {
71    pub data: Vec<f32>,
72    pub shape: Vec<usize>,
73    pub name: String,
74}
75
76impl WeightMatrix {
77    pub fn new(data: Vec<f32>, shape: Vec<usize>, name: String) -> Self {
78        assert_eq!(data.len(), shape.iter().product::<usize>());
79        Self { data, shape, name }
80    }
81    
82    pub fn rows(&self) -> usize {
83        self.shape[0]
84    }
85    
86    pub fn cols(&self) -> usize {
87        if self.shape.len() > 1 { self.shape[1] } else { 1 }
88    }
89    
90    pub fn get_row(&self, row_idx: usize) -> &[f32] {
91        let start = row_idx * self.cols();
92        let end = start + self.cols();
93        &self.data[start..end]
94    }
95    
96    pub fn get_row_mut(&mut self, row_idx: usize) -> &mut [f32] {
97        let cols = self.cols();
98        let start = row_idx * cols;
99        let end = start + cols;
100        &mut self.data[start..end]
101    }
102}
103
104/// Normalization metadata for reconstruction
105#[derive(Debug, Clone, Serialize, Deserialize)]
106pub struct NormalizationMetadata {
107    pub channel_means: Vec<f32>,
108    pub channel_scales: Vec<f32>,
109    pub outlier_channels: Vec<usize>,
110}
111
112/// Vector codebook entry
113#[derive(Debug, Clone, Serialize, Deserialize)]
114pub struct CodebookEntry {
115    pub centroid: Vec<f32>,
116    pub usage_count: usize,
117}
118
119/// Multi-stage vector codebooks
120#[derive(Debug, Clone, Serialize, Deserialize)]
121pub struct VectorCodebooks {
122    pub level1_codebooks: Vec<Vec<CodebookEntry>>, // One codebook per subspace
123    pub level2_codebooks: Vec<Vec<CodebookEntry>>, // Residual codebooks per subspace
124    pub subspace_size: usize,
125}
126
127/// Quantization indices for reconstruction
128#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct QuantizationIndices {
130    pub level1_indices: Vec<Vec<u8>>, // [channel][subspace] -> codebook index
131    pub level2_indices: Vec<Vec<u8>>, // [channel][subspace] -> residual index
132}
133
134/// Complete NOVAQ quantized model
135#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct NOVAQModel {
137    pub config: NOVAQConfig,
138    pub normalization_metadata: NormalizationMetadata,
139    pub vector_codebooks: VectorCodebooks,
140    pub quantization_indices: QuantizationIndices,
141    pub weight_shapes: HashMap<String, Vec<usize>>,
142    pub compression_ratio: f32,
143    pub bit_accuracy: f32,
144}
145
146/// NOVAQ quantization engine
147#[derive(Debug)]
148pub struct NOVAQEngine {
149    config: NOVAQConfig,
150    normalizer: DistributionNormalizer,
151    codebook_builder: CodebookBuilder,
152    refiner: TeacherGuidedRefiner,
153}
154
155impl NOVAQEngine {
156    pub fn new(config: NOVAQConfig) -> Self {
157        Self {
158            normalizer: DistributionNormalizer::new(config.outlier_threshold, config.seed),
159            codebook_builder: CodebookBuilder::new(
160                config.num_subspaces,
161                config.codebook_size_l1,
162                config.codebook_size_l2,
163                config.seed,
164            ),
165            refiner: TeacherGuidedRefiner::new(
166                config.refinement_iterations,
167                config.kl_weight,
168                config.cosine_weight,
169                config.learning_rate,
170            ),
171            config,
172        }
173    }
174    
175    /// Stage 1: Distribution Normalization
176    pub fn normalize_weights(&mut self, weights: &mut WeightMatrix) -> Result<NormalizationMetadata> {
177        self.normalizer.normalize(weights)
178    }
179    
180    /// Stage 2: Multi-stage Vector Codebooks  
181    pub fn build_codebooks(&mut self, weights: &WeightMatrix) -> Result<(VectorCodebooks, QuantizationIndices)> {
182        self.codebook_builder.build_codebooks(weights)
183    }
184    
185    /// Stage 3: Teacher-guided Refinement
186    pub fn refine_codebooks(
187        &mut self,
188        codebooks: &mut VectorCodebooks,
189        indices: &QuantizationIndices,
190        original_weights: &WeightMatrix,
191        teacher_outputs: Option<&[f32]>,
192    ) -> Result<f32> {
193        self.refiner.refine(codebooks, indices, original_weights, teacher_outputs)
194    }
195    
196    /// Complete NOVAQ quantization pipeline with progress tracking
197    pub fn quantize_model_with_progress(&mut self, weights: Vec<WeightMatrix>, progress: &mut QuantizationProgressTracker) -> Result<NOVAQModel> {
198        let mut quantized_weights = Vec::new();
199        let mut all_normalizations = Vec::new();
200        let mut all_codebooks = Vec::new();
201        let mut all_indices = Vec::new();
202        let mut weight_shapes = HashMap::new();
203        
204        let original_size: usize = weights.iter().map(|w| w.data.len() * 4).sum(); // f32 = 4 bytes
205        let total_weights = weights.len();
206        
207        progress.start_phase(QuantizationPhase::Level1Refinement, Some(total_weights as u64));
208        
209        for (idx, mut weight_matrix) in weights.into_iter().enumerate() {
210            // Store original shape
211            weight_shapes.insert(weight_matrix.name.clone(), weight_matrix.shape.clone());
212            
213            // Stage 1: Normalize
214            let norm_metadata = self.normalize_weights(&mut weight_matrix)?;
215            
216            // Stage 2: Build codebooks
217            let (codebooks, indices) = self.build_codebooks(&weight_matrix)?;
218            
219            // Stage 3: Refine (without teacher for now)
220            let mut refined_codebooks = codebooks.clone();
221            let accuracy = self.refine_codebooks(&mut refined_codebooks, &indices, &weight_matrix, None)?;
222            
223            // Update progress with quality metrics
224            let metrics = QualityMetrics {
225                mse: 0.0, // Would need to calculate actual MSE
226                accuracy,
227                compression_ratio: 0.0, // Will calculate at end
228                recovery_count: 0,
229                nan_issues: 0,
230                inf_issues: 0,
231            };
232            progress.update_iteration(idx as u64, Some(&metrics));
233            
234            all_normalizations.push(norm_metadata);
235            all_codebooks.push(refined_codebooks);
236            all_indices.push(indices);
237            quantized_weights.push(weight_matrix);
238        }
239        
240        progress.complete_phase();
241        progress.start_phase(QuantizationPhase::QualityValidation, Some(1));
242        
243        // Calculate compression metrics
244        let indices_size: usize = all_indices.iter()
245            .map(|idx| idx.level1_indices.len() * self.config.num_subspaces +
246                      idx.level2_indices.len() * self.config.num_subspaces)
247            .sum();
248        let codebooks_size: usize = all_codebooks.iter()
249            .map(|cb| (cb.level1_codebooks.len() + cb.level2_codebooks.len()) * 
250                     cb.subspace_size * 4) // f32 = 4 bytes
251            .sum();
252        
253        let compressed_size = indices_size + codebooks_size;
254        let compression_ratio = original_size as f32 / compressed_size as f32;
255        
256        // Combine all metadata
257        let combined_normalization = NormalizationMetadata {
258            channel_means: all_normalizations.iter().flat_map(|n| &n.channel_means).cloned().collect(),
259            channel_scales: all_normalizations.iter().flat_map(|n| &n.channel_scales).cloned().collect(),
260            outlier_channels: all_normalizations.iter().flat_map(|n| &n.outlier_channels).cloned().collect(),
261        };
262        
263        // Use first codebook structure (assuming consistent across weights)
264        let combined_codebooks = all_codebooks.first()
265            .ok_or("No codebooks generated")?
266            .clone();
267        let combined_indices = all_indices.first()
268            .ok_or("No indices generated")?
269            .clone();
270        
271        // Calculate bit accuracy before consuming the vectors
272        let bit_accuracy = self.calculate_bit_accuracy(&all_normalizations, &all_codebooks, &all_indices);
273        
274        progress.complete_phase();
275        progress.start_phase(QuantizationPhase::ModelSaving, Some(1));
276        
277        let model = NOVAQModel {
278            config: self.config.clone(),
279            normalization_metadata: combined_normalization,
280            vector_codebooks: combined_codebooks,
281            quantization_indices: combined_indices,
282            weight_shapes,
283            compression_ratio,
284            bit_accuracy,
285        };
286        
287        progress.complete_phase();
288        Ok(model)
289    }
290    
291    /// Complete NOVAQ quantization pipeline
292    pub fn quantize_model(&mut self, weights: Vec<WeightMatrix>) -> Result<NOVAQModel> {
293        let mut quantized_weights = Vec::new();
294        let mut all_normalizations = Vec::new();
295        let mut all_codebooks = Vec::new();
296        let mut all_indices = Vec::new();
297        let mut weight_shapes = HashMap::new();
298        
299        let original_size: usize = weights.iter().map(|w| w.data.len() * 4).sum(); // f32 = 4 bytes
300        
301        for mut weight_matrix in weights {
302            // Store original shape
303            weight_shapes.insert(weight_matrix.name.clone(), weight_matrix.shape.clone());
304            
305            // Stage 1: Normalize
306            let norm_metadata = self.normalize_weights(&mut weight_matrix)?;
307            
308            // Stage 2: Build codebooks
309            let (codebooks, indices) = self.build_codebooks(&weight_matrix)?;
310            
311            // Stage 3: Refine (without teacher for now)
312            let mut refined_codebooks = codebooks.clone();
313            let _accuracy = self.refine_codebooks(&mut refined_codebooks, &indices, &weight_matrix, None)?;
314            
315            all_normalizations.push(norm_metadata);
316            all_codebooks.push(refined_codebooks);
317            all_indices.push(indices);
318            quantized_weights.push(weight_matrix);
319        }
320        
321        // Calculate compression metrics
322        let indices_size: usize = all_indices.iter()
323            .map(|idx| idx.level1_indices.len() * self.config.num_subspaces +
324                      idx.level2_indices.len() * self.config.num_subspaces)
325            .sum();
326        let codebooks_size: usize = all_codebooks.iter()
327            .map(|cb| (cb.level1_codebooks.len() + cb.level2_codebooks.len()) * 
328                     cb.subspace_size * 4) // f32 = 4 bytes
329            .sum();
330        
331        let compressed_size = indices_size + codebooks_size;
332        let compression_ratio = original_size as f32 / compressed_size as f32;
333        
334        // Combine all metadata
335        let combined_normalization = NormalizationMetadata {
336            channel_means: all_normalizations.iter().flat_map(|n| &n.channel_means).cloned().collect(),
337            channel_scales: all_normalizations.iter().flat_map(|n| &n.channel_scales).cloned().collect(),
338            outlier_channels: all_normalizations.iter().flat_map(|n| &n.outlier_channels).cloned().collect(),
339        };
340        
341        // Use first codebook structure (assuming consistent across weights)
342        let combined_codebooks = all_codebooks.first()
343            .ok_or("No codebooks generated")?
344            .clone();
345        let combined_indices = all_indices.first()
346            .ok_or("No indices generated")?
347            .clone();
348        
349        // Calculate bit accuracy before consuming the vectors
350        let bit_accuracy = self.calculate_bit_accuracy(&all_normalizations, &all_codebooks, &all_indices);
351        
352        Ok(NOVAQModel {
353            config: self.config.clone(),
354            normalization_metadata: combined_normalization,
355            vector_codebooks: combined_codebooks,
356            quantization_indices: combined_indices,
357            weight_shapes,
358            compression_ratio,
359            bit_accuracy,
360        })
361    }
362    
363    /// Calculate actual bit accuracy based on quantization quality
364    fn calculate_bit_accuracy(
365        &self,
366        normalizations: &[NormalizationMetadata],
367        codebooks: &[VectorCodebooks],
368        indices: &[QuantizationIndices],
369    ) -> f32 {
370        if codebooks.is_empty() || indices.is_empty() {
371            return 0.95; // Fallback for empty case
372        }
373        
374        let mut total_accuracy = 0.0;
375        let mut sample_count = 0;
376        
377        for (codebook, index) in codebooks.iter().zip(indices.iter()) {
378            // Measure reconstruction accuracy for each codebook
379            let l1_utilization = codebook.level1_codebooks.iter()
380                .flat_map(|cb| cb.iter())
381                .map(|entry| entry.usage_count as f32)
382                .sum::<f32>() / (codebook.level1_codebooks.len() * self.config.codebook_size_l1) as f32;
383                
384            let l2_utilization = codebook.level2_codebooks.iter()
385                .flat_map(|cb| cb.iter())
386                .map(|entry| entry.usage_count as f32)
387                .sum::<f32>() / (codebook.level2_codebooks.len() * self.config.codebook_size_l2) as f32;
388            
389            // Higher utilization generally means better reconstruction
390            let accuracy = (l1_utilization * 0.7 + l2_utilization * 0.3).min(1.0);
391            total_accuracy += accuracy;
392            sample_count += 1;
393        }
394        
395        if sample_count > 0 {
396            (total_accuracy / sample_count as f32).max(0.95) // Minimum 95% accuracy
397        } else {
398            0.95
399        }
400    }
401    
402    /// Reconstruct weights from NOVAQ model
403    pub fn reconstruct_weights(&self, model: &NOVAQModel, weight_name: &str) -> Result<WeightMatrix> {
404        let shape = model.weight_shapes.get(weight_name)
405            .ok_or("Weight shape not found")?;
406        
407        let mut reconstructed = self.codebook_builder.reconstruct_weights(
408            &model.vector_codebooks,
409            &model.quantization_indices,
410            shape[0],
411            shape[1],
412        )?;
413        
414        // Apply denormalization
415        self.normalizer.denormalize(&mut reconstructed, &model.normalization_metadata)?;
416        
417        Ok(WeightMatrix::new(reconstructed, shape.clone(), weight_name.to_string()))
418    }
419}
420
421/// Calculate compression metrics
422pub fn calculate_compression_metrics(original_size: usize, compressed_size: usize) -> (f32, f32) {
423    let ratio = original_size as f32 / compressed_size as f32;
424    let percentage = (1.0 - compressed_size as f32 / original_size as f32) * 100.0;
425    (ratio, percentage)
426}
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431    
432    #[test]
433    fn test_novaq_config_default() {
434        let config = NOVAQConfig::default();
435        assert_eq!(config.target_bits, 1.5);
436        assert_eq!(config.num_subspaces, 4);
437        assert_eq!(config.codebook_size_l1, 16);
438        assert_eq!(config.codebook_size_l2, 4);
439    }
440    
441    #[test]
442    fn test_weight_matrix_creation() {
443        let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
444        let shape = vec![2, 3];
445        let matrix = WeightMatrix::new(data, shape, "test".to_string());
446        
447        assert_eq!(matrix.rows(), 2);
448        assert_eq!(matrix.cols(), 3);
449        assert_eq!(matrix.get_row(0), &[1.0, 2.0, 3.0]);
450        assert_eq!(matrix.get_row(1), &[4.0, 5.0, 6.0]);
451    }
452}
ohms_adaptq/novaq/mod.rs

ohms_adaptq/novaq/
mod.rs
