trustformers_optim/
quantized.rs

1use anyhow::Result;
2use std::collections::HashMap;
3use trustformers_core::tensor::Tensor;
4
5#[derive(Debug, Clone)]
6pub struct QuantizationConfig {
7    pub scale: f32,
8    pub zero_point: i8,
9    pub min_val: f32,
10    pub max_val: f32,
11}
12
13impl QuantizationConfig {
14    pub fn new(min_val: f32, max_val: f32) -> Self {
15        let scale = (max_val - min_val) / 255.0;
16        let zero_point = (-min_val / scale).round().clamp(-128.0, 127.0) as i8;
17
18        Self {
19            scale,
20            zero_point,
21            min_val,
22            max_val,
23        }
24    }
25
26    pub fn quantize(&self, value: f32) -> i8 {
27        let quantized = ((value - self.min_val) / self.scale).round() - 128.0;
28        quantized.clamp(-128.0, 127.0) as i8
29    }
30
31    pub fn dequantize(&self, quantized: i8) -> f32 {
32        (quantized as f32 + 128.0) * self.scale + self.min_val
33    }
34}
35
36#[derive(Debug, Clone)]
37pub struct QuantizedState {
38    pub data: Vec<i8>,
39    pub config: QuantizationConfig,
40}
41
42impl QuantizedState {
43    pub fn new(values: &[f32]) -> Self {
44        let min_val = values.iter().fold(f32::INFINITY, |a, &b| a.min(b));
45        let max_val = values.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
46
47        let config = QuantizationConfig::new(min_val, max_val);
48        let data: Vec<i8> = values.iter().map(|&v| config.quantize(v)).collect();
49
50        Self { data, config }
51    }
52
53    pub fn to_f32(&self) -> Vec<f32> {
54        self.data.iter().map(|&q| self.config.dequantize(q)).collect()
55    }
56
57    pub fn update(&mut self, new_values: &[f32]) {
58        let min_val = new_values.iter().fold(f32::INFINITY, |a, &b| a.min(b));
59        let max_val = new_values.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
60
61        self.config = QuantizationConfig::new(min_val, max_val);
62        self.data = new_values.iter().map(|&v| self.config.quantize(v)).collect();
63    }
64}
65
66#[derive(Debug)]
67pub struct Adam8bit {
68    pub learning_rate: f32,
69    pub beta1: f32,
70    pub beta2: f32,
71    pub epsilon: f32,
72    pub weight_decay: f32,
73    pub step: usize,
74    pub momentum_states: HashMap<String, QuantizedState>,
75    pub variance_states: HashMap<String, QuantizedState>,
76}
77
78impl Default for Adam8bit {
79    fn default() -> Self {
80        Self {
81            learning_rate: 1e-3,
82            beta1: 0.9,
83            beta2: 0.999,
84            epsilon: 1e-8,
85            weight_decay: 0.0,
86            step: 0,
87            momentum_states: HashMap::new(),
88            variance_states: HashMap::new(),
89        }
90    }
91}
92
93impl Adam8bit {
94    pub fn new(learning_rate: f32) -> Self {
95        Self {
96            learning_rate,
97            ..Default::default()
98        }
99    }
100
101    pub fn with_config(
102        learning_rate: f32,
103        beta1: f32,
104        beta2: f32,
105        epsilon: f32,
106        weight_decay: f32,
107    ) -> Self {
108        Self {
109            learning_rate,
110            beta1,
111            beta2,
112            epsilon,
113            weight_decay,
114            step: 0,
115            momentum_states: HashMap::new(),
116            variance_states: HashMap::new(),
117        }
118    }
119
120    pub fn step(
121        &mut self,
122        parameters: &mut HashMap<String, Tensor>,
123        gradients: &HashMap<String, Tensor>,
124    ) -> Result<()> {
125        self.step += 1;
126
127        let bias_correction1 = 1.0 - self.beta1.powi(self.step as i32);
128        let bias_correction2 = 1.0 - self.beta2.powi(self.step as i32);
129
130        for (name, param) in parameters.iter_mut() {
131            let grad = gradients
132                .get(name)
133                .ok_or_else(|| anyhow::anyhow!("Missing gradient for parameter: {}", name))?;
134
135            let mut param_data = param.data()?;
136            let grad_data = grad.data()?;
137
138            if param_data.len() != grad_data.len() {
139                return Err(anyhow::anyhow!(
140                    "Parameter and gradient size mismatch for: {}",
141                    name
142                ));
143            }
144
145            if !self.momentum_states.contains_key(name) {
146                let zeros = vec![0.0; param_data.len()];
147                self.momentum_states.insert(name.clone(), QuantizedState::new(&zeros));
148                self.variance_states.insert(name.clone(), QuantizedState::new(&zeros));
149            }
150
151            let momentum_state = self.momentum_states.get_mut(name).ok_or_else(|| {
152                anyhow::anyhow!("momentum_state should exist after initialization")
153            })?;
154            let variance_state = self.variance_states.get_mut(name).ok_or_else(|| {
155                anyhow::anyhow!("variance_state should exist after initialization")
156            })?;
157
158            let mut momentum = momentum_state.to_f32();
159            let mut variance = variance_state.to_f32();
160
161            for i in 0..param_data.len() {
162                let mut grad_val = grad_data[i];
163
164                if self.weight_decay > 0.0 {
165                    grad_val += self.weight_decay * param_data[i];
166                }
167
168                momentum[i] = self.beta1 * momentum[i] + (1.0 - self.beta1) * grad_val;
169                variance[i] = self.beta2 * variance[i] + (1.0 - self.beta2) * grad_val * grad_val;
170
171                let corrected_momentum = momentum[i] / bias_correction1;
172                let corrected_variance = variance[i] / bias_correction2;
173
174                param_data[i] -= self.learning_rate * corrected_momentum
175                    / (corrected_variance.sqrt() + self.epsilon);
176            }
177
178            momentum_state.update(&momentum);
179            variance_state.update(&variance);
180
181            // Update the parameter tensor with modified data
182            *param = Tensor::new(param_data)?;
183        }
184
185        Ok(())
186    }
187
188    pub fn memory_usage(&self) -> usize {
189        let mut total = 0;
190        for state in self.momentum_states.values() {
191            total += state.data.len();
192        }
193        for state in self.variance_states.values() {
194            total += state.data.len();
195        }
196        total
197    }
198
199    pub fn memory_savings_vs_fp32(&self) -> f32 {
200        let quantized_size = self.memory_usage();
201        let fp32_equivalent = quantized_size * 4;
202        1.0 - (quantized_size as f32 / fp32_equivalent as f32)
203    }
204}
205
206#[derive(Debug)]
207pub struct AdamW8bit {
208    pub learning_rate: f32,
209    pub beta1: f32,
210    pub beta2: f32,
211    pub epsilon: f32,
212    pub weight_decay: f32,
213    pub step: usize,
214    pub momentum_states: HashMap<String, QuantizedState>,
215    pub variance_states: HashMap<String, QuantizedState>,
216}
217
218impl Default for AdamW8bit {
219    fn default() -> Self {
220        Self {
221            learning_rate: 1e-3,
222            beta1: 0.9,
223            beta2: 0.999,
224            epsilon: 1e-8,
225            weight_decay: 1e-2,
226            step: 0,
227            momentum_states: HashMap::new(),
228            variance_states: HashMap::new(),
229        }
230    }
231}
232
233impl AdamW8bit {
234    pub fn new(learning_rate: f32) -> Self {
235        Self {
236            learning_rate,
237            ..Default::default()
238        }
239    }
240
241    pub fn with_config(
242        learning_rate: f32,
243        beta1: f32,
244        beta2: f32,
245        epsilon: f32,
246        weight_decay: f32,
247    ) -> Self {
248        Self {
249            learning_rate,
250            beta1,
251            beta2,
252            epsilon,
253            weight_decay,
254            step: 0,
255            momentum_states: HashMap::new(),
256            variance_states: HashMap::new(),
257        }
258    }
259
260    pub fn step(
261        &mut self,
262        parameters: &mut HashMap<String, Tensor>,
263        gradients: &HashMap<String, Tensor>,
264    ) -> Result<()> {
265        self.step += 1;
266
267        let bias_correction1 = 1.0 - self.beta1.powi(self.step as i32);
268        let bias_correction2 = 1.0 - self.beta2.powi(self.step as i32);
269
270        for (name, param) in parameters.iter_mut() {
271            let grad = gradients
272                .get(name)
273                .ok_or_else(|| anyhow::anyhow!("Missing gradient for parameter: {}", name))?;
274
275            let mut param_data = param.data()?;
276            let grad_data = grad.data()?;
277
278            if param_data.len() != grad_data.len() {
279                return Err(anyhow::anyhow!(
280                    "Parameter and gradient size mismatch for: {}",
281                    name
282                ));
283            }
284
285            if !self.momentum_states.contains_key(name) {
286                let zeros = vec![0.0; param_data.len()];
287                self.momentum_states.insert(name.clone(), QuantizedState::new(&zeros));
288                self.variance_states.insert(name.clone(), QuantizedState::new(&zeros));
289            }
290
291            let momentum_state = self.momentum_states.get_mut(name).ok_or_else(|| {
292                anyhow::anyhow!("momentum_state should exist after initialization")
293            })?;
294            let variance_state = self.variance_states.get_mut(name).ok_or_else(|| {
295                anyhow::anyhow!("variance_state should exist after initialization")
296            })?;
297
298            let mut momentum = momentum_state.to_f32();
299            let mut variance = variance_state.to_f32();
300
301            for i in 0..param_data.len() {
302                let grad_val = grad_data[i];
303
304                momentum[i] = self.beta1 * momentum[i] + (1.0 - self.beta1) * grad_val;
305                variance[i] = self.beta2 * variance[i] + (1.0 - self.beta2) * grad_val * grad_val;
306
307                let corrected_momentum = momentum[i] / bias_correction1;
308                let corrected_variance = variance[i] / bias_correction2;
309
310                let update = corrected_momentum / (corrected_variance.sqrt() + self.epsilon);
311
312                param_data[i] = param_data[i] * (1.0 - self.learning_rate * self.weight_decay)
313                    - self.learning_rate * update;
314            }
315
316            momentum_state.update(&momentum);
317            variance_state.update(&variance);
318
319            // Update the parameter tensor with modified data
320            *param = Tensor::new(param_data)?;
321        }
322
323        Ok(())
324    }
325
326    pub fn memory_usage(&self) -> usize {
327        let mut total = 0;
328        for state in self.momentum_states.values() {
329            total += state.data.len();
330        }
331        for state in self.variance_states.values() {
332            total += state.data.len();
333        }
334        total
335    }
336
337    pub fn memory_savings_vs_fp32(&self) -> f32 {
338        let quantized_size = self.memory_usage();
339        let fp32_equivalent = quantized_size * 4;
340        1.0 - (quantized_size as f32 / fp32_equivalent as f32)
341    }
342}
343
344#[cfg(test)]
345mod tests {
346    use super::*;
347    use approx::assert_abs_diff_eq;
348
349    #[test]
350    fn test_quantization_config() {
351        let config = QuantizationConfig::new(-1.0, 1.0);
352
353        // Test that values map correctly to the quantized range
354        assert_eq!(config.quantize(-1.0), -128);
355        assert_eq!(config.quantize(1.0), 127);
356
357        // Test middle value
358        let mid_quantized = config.quantize(0.0);
359        assert!((-1..=1).contains(&mid_quantized));
360
361        // Test round-trip accuracy
362        let original = 0.5;
363        let quantized = config.quantize(original);
364        let reconstructed = config.dequantize(quantized);
365        assert_abs_diff_eq!(original, reconstructed, epsilon = 0.02);
366    }
367
368    #[test]
369    fn test_quantized_state() {
370        let values = vec![0.1, -0.5, 0.8, -0.2];
371        let state = QuantizedState::new(&values);
372
373        let reconstructed = state.to_f32();
374
375        // Test that we have the right number of values
376        assert_eq!(values.len(), reconstructed.len());
377
378        // Test that quantization preserves relative ordering
379        assert!(reconstructed[2] > reconstructed[0]); // 0.8 > 0.1
380        assert!(reconstructed[1] < reconstructed[0]); // -0.5 < 0.1
381
382        // Test approximate reconstruction (quantization introduces some error)
383        for (orig, recon) in values.iter().zip(reconstructed.iter()) {
384            assert_abs_diff_eq!(orig, recon, epsilon = 0.1);
385        }
386    }
387
388    #[test]
389    fn test_adam8bit_creation() {
390        let optimizer = Adam8bit::new(0.001);
391        assert_eq!(optimizer.learning_rate, 0.001);
392        assert_eq!(optimizer.beta1, 0.9);
393        assert_eq!(optimizer.beta2, 0.999);
394        assert_eq!(optimizer.step, 0);
395    }
396
397    #[test]
398    fn test_adam8bit_memory_usage() {
399        let mut optimizer = Adam8bit::new(0.001);
400
401        let mut parameters = HashMap::new();
402        let mut gradients = HashMap::new();
403
404        let param_data = vec![1.0, 2.0, 3.0, 4.0];
405        let grad_data = vec![0.1, 0.2, 0.3, 0.4];
406
407        parameters.insert(
408            "layer1".to_string(),
409            Tensor::new(param_data).expect("Failed to create tensor"),
410        );
411        gradients.insert(
412            "layer1".to_string(),
413            Tensor::new(grad_data).expect("Failed to create tensor"),
414        );
415
416        optimizer.step(&mut parameters, &gradients).expect("Step failed");
417
418        assert_eq!(optimizer.memory_usage(), 8);
419        assert!(optimizer.memory_savings_vs_fp32() > 0.7);
420    }
421
422    #[test]
423    fn test_adamw8bit_creation() {
424        let optimizer = AdamW8bit::new(0.001);
425        assert_eq!(optimizer.learning_rate, 0.001);
426        assert_eq!(optimizer.weight_decay, 0.01);
427    }
428}
trustformers_optim/quantized.rs

trustformers_optim/
quantized.rs
