Skip to main content

pictor_model/
calibration.rs

1//! Post-Training Quantization (PTQ) calibration pipeline.
2//!
3//! PTQ calibration:
4//! 1. Run calibration data through the model (no gradient)
5//! 2. Collect activation statistics per layer (min, max, percentiles)
6//! 3. Compute optimal quantization scales
7//! 4. Store scales for use during static-quantized inference
8//!
9//! Supported calibration methods:
10//! - MinMax: scale = max(|x|) / clip_val (simple, fast)
11//! - Percentile: scale = p99(|x|) / clip_val (robust to outliers)
12//! - ACIQ: Analytical Clipping for Integer Quantization (Bell et al. 2019)
13//!   optimal_clip = 2.83 * std_dev for normal distributions
14//! - MSE: Mean squared error minimization (approximated via grid search)
15
16use std::collections::HashMap;
17
18// ─── Calibration method ───────────────────────────────────────────────────────
19
20/// Calibration method for computing quantization scale.
21#[derive(Debug, Clone, Copy, PartialEq)]
22pub enum CalibMethod {
23    /// Use the global maximum absolute value.
24    MinMax,
25    /// Use a high percentile of absolute values.
26    Percentile(f32), // e.g., 0.9999
27    /// ACIQ: analytical optimal clipping based on std_dev.
28    Aciq,
29    /// Mean squared error (MSE) minimization (approximated).
30    Mse,
31}
32
33// ─── Layer calibration statistics ─────────────────────────────────────────────
34
35/// Statistics collected from calibration activations for one layer.
36#[derive(Debug, Clone)]
37pub struct LayerCalibStats {
38    pub layer_name: String,
39    pub num_samples: usize,
40    pub running_min: f32,
41    pub running_max: f32,
42    pub running_mean: f32,
43    pub running_var: f32, // Welford online variance accumulator (M2)
44    /// Histogram of |x| values (for percentile computation), 256 bins.
45    histogram: Vec<u64>,
46    histogram_max: f32,
47}
48
49impl LayerCalibStats {
50    /// Create a new empty stats collector for the given layer.
51    pub fn new(layer_name: impl Into<String>) -> Self {
52        Self {
53            layer_name: layer_name.into(),
54            num_samples: 0,
55            running_min: f32::INFINITY,
56            running_max: f32::NEG_INFINITY,
57            running_mean: 0.0,
58            running_var: 0.0,
59            histogram: vec![0u64; 256],
60            histogram_max: 0.0,
61        }
62    }
63
64    /// Update statistics with a new batch of activations.
65    ///
66    /// Uses Welford's online algorithm for numerically stable mean/variance.
67    pub fn update(&mut self, activations: &[f32]) {
68        for &x in activations {
69            if !x.is_finite() {
70                continue;
71            }
72
73            // Update min/max
74            if x < self.running_min {
75                self.running_min = x;
76            }
77            if x > self.running_max {
78                self.running_max = x;
79            }
80
81            // Welford online mean/variance
82            self.num_samples += 1;
83            let n = self.num_samples as f32;
84            let delta = x - self.running_mean;
85            self.running_mean += delta / n;
86            let delta2 = x - self.running_mean;
87            self.running_var += delta * delta2;
88
89            // Update histogram max (using abs value)
90            let abs_x = x.abs();
91            if abs_x > self.histogram_max {
92                // Rebuild histogram with new max
93                self.rebuild_histogram_max(abs_x);
94            }
95        }
96
97        // Re-insert all new activations into histogram
98        for &x in activations {
99            if !x.is_finite() {
100                continue;
101            }
102            let abs_x = x.abs();
103            self.insert_histogram(abs_x);
104        }
105    }
106
107    /// Rebuild histogram bins when a new maximum is encountered.
108    fn rebuild_histogram_max(&mut self, new_max: f32) {
109        // Scale existing histogram to new range
110        if self.histogram_max > 0.0 && new_max > self.histogram_max {
111            let scale = self.histogram_max / new_max;
112            let mut new_hist = vec![0u64; 256];
113            for (old_bin, &count) in self.histogram.iter().enumerate() {
114                if count == 0 {
115                    continue;
116                }
117                // Map old bin center to new bin
118                let old_frac = (old_bin as f32 + 0.5) / 256.0;
119                let new_frac = old_frac * scale;
120                let new_bin = (new_frac * 256.0) as usize;
121                let new_bin = new_bin.min(255);
122                new_hist[new_bin] += count;
123            }
124            self.histogram = new_hist;
125        }
126        self.histogram_max = new_max;
127    }
128
129    /// Insert a single absolute value into the histogram.
130    fn insert_histogram(&mut self, abs_x: f32) {
131        if self.histogram_max <= 0.0 {
132            return;
133        }
134        let frac = abs_x / self.histogram_max;
135        let bin = (frac * 256.0) as usize;
136        let bin = bin.min(255);
137        self.histogram[bin] += 1;
138    }
139
140    /// Standard deviation computed from Welford variance accumulator.
141    pub fn std_dev(&self) -> f32 {
142        if self.num_samples < 2 {
143            return 0.0;
144        }
145        let variance = self.running_var / self.num_samples as f32;
146        variance.max(0.0).sqrt()
147    }
148
149    /// Percentile of absolute values (0.0 - 1.0).
150    ///
151    /// Uses the histogram for efficient O(256) computation.
152    pub fn percentile_abs(&self, p: f32) -> f32 {
153        if self.num_samples == 0 || self.histogram_max <= 0.0 {
154            return 0.0;
155        }
156
157        let p_clamped = p.clamp(0.0, 1.0);
158        let target_count = (p_clamped * self.num_samples as f32).ceil() as u64;
159        let target_count = target_count.max(1);
160
161        let mut cumulative = 0u64;
162        for (bin_idx, &count) in self.histogram.iter().enumerate() {
163            cumulative += count;
164            if cumulative >= target_count {
165                // Return upper edge of this bin
166                let upper = (bin_idx as f32 + 1.0) / 256.0 * self.histogram_max;
167                return upper.min(self.histogram_max);
168            }
169        }
170
171        self.histogram_max
172    }
173
174    /// ACIQ optimal clipping: 2.83 * std_dev (Laplacian assumption).
175    ///
176    /// Reference: Bell et al. 2019, "Accurate Post Training Quantization".
177    pub fn aciq_clip(&self) -> f32 {
178        2.83 * self.std_dev()
179    }
180
181    /// Compute final quantization scale for INT8 (clip_val = 127).
182    pub fn compute_scale(&self, method: CalibMethod) -> f32 {
183        const CLIP_VAL: f32 = 127.0;
184        self.compute_scale_with_clip(method, CLIP_VAL)
185    }
186
187    /// Compute final quantization scale for INT4 (clip_val = 7).
188    pub fn compute_scale_int4(&self, method: CalibMethod) -> f32 {
189        const CLIP_VAL: f32 = 7.0;
190        self.compute_scale_with_clip(method, CLIP_VAL)
191    }
192
193    /// Internal: compute scale for the given clip value.
194    fn compute_scale_with_clip(&self, method: CalibMethod, clip_val: f32) -> f32 {
195        if self.num_samples == 0 {
196            return 0.0;
197        }
198
199        let abs_max = self.running_min.abs().max(self.running_max.abs());
200
201        let clipping_value = match method {
202            CalibMethod::MinMax => abs_max,
203            CalibMethod::Percentile(p) => {
204                let pv = self.percentile_abs(p);
205                if pv <= 0.0 {
206                    abs_max
207                } else {
208                    pv
209                }
210            }
211            CalibMethod::Aciq => {
212                let aciq = self.aciq_clip();
213                if aciq <= 0.0 {
214                    abs_max
215                } else {
216                    aciq.min(abs_max)
217                }
218            }
219            CalibMethod::Mse => {
220                // Grid search: try fractions of abs_max, pick the one minimizing MSE proxy.
221                // We approximate MSE using histogram statistics.
222                self.mse_optimal_clip(abs_max, clip_val)
223            }
224        };
225
226        if clipping_value <= 0.0 {
227            return 0.0;
228        }
229
230        clipping_value / clip_val
231    }
232
233    /// Approximate MSE-optimal clipping via grid search on the histogram.
234    ///
235    /// For each candidate clip value, computes:
236    ///   MSE ≈ clipping_error + rounding_error
237    fn mse_optimal_clip(&self, abs_max: f32, clip_val: f32) -> f32 {
238        if abs_max <= 0.0 || self.num_samples == 0 {
239            return abs_max;
240        }
241
242        let n_steps = 100usize;
243        let mut best_clip = abs_max;
244        let mut best_mse = f32::INFINITY;
245
246        for step in 1..=n_steps {
247            let frac = step as f32 / n_steps as f32;
248            let candidate_clip = abs_max * frac;
249            if candidate_clip <= 0.0 {
250                continue;
251            }
252
253            // Rounding error variance: (clip / clip_val)^2 / 3
254            let scale = candidate_clip / clip_val;
255            let rounding_var = scale * scale / 3.0;
256
257            // Clipping error: estimate fraction of values > candidate_clip
258            // from histogram, compute their average squared distance.
259            let clip_frac = candidate_clip / self.histogram_max;
260            let clip_bin = (clip_frac * 256.0) as usize;
261
262            let mut clip_error_sq = 0.0f32;
263            for (bin_idx, &count) in self.histogram.iter().enumerate().skip(clip_bin.min(255)) {
264                if count == 0 {
265                    continue;
266                }
267                // Bin center value
268                let bin_center = (bin_idx as f32 + 0.5) / 256.0 * self.histogram_max;
269                let excess = (bin_center - candidate_clip).max(0.0);
270                clip_error_sq += count as f32 * excess * excess;
271            }
272
273            let n = self.num_samples as f32;
274            let clip_mse = if n > 0.0 { clip_error_sq / n } else { 0.0 };
275            let total_mse = rounding_var + clip_mse;
276
277            if total_mse < best_mse {
278                best_mse = total_mse;
279                best_clip = candidate_clip;
280            }
281        }
282
283        best_clip
284    }
285
286    /// Summary statistics as a struct.
287    pub fn summary(&self) -> CalibSummary {
288        let p99 = self.percentile_abs(0.99);
289        let p9999 = self.percentile_abs(0.9999);
290        let int8_scale = self.compute_scale(CalibMethod::Percentile(0.9999));
291
292        let (min, max) = if self.num_samples == 0 {
293            (0.0, 0.0)
294        } else {
295            (self.running_min, self.running_max)
296        };
297
298        CalibSummary {
299            layer_name: self.layer_name.clone(),
300            num_samples: self.num_samples,
301            min,
302            max,
303            mean: self.running_mean,
304            std_dev: self.std_dev(),
305            p99,
306            p9999,
307            suggested_int8_scale: int8_scale,
308        }
309    }
310}
311
312// ─── Calibration summary ──────────────────────────────────────────────────────
313
314/// Summary statistics for a single layer's calibration.
315#[derive(Debug, Clone)]
316pub struct CalibSummary {
317    pub layer_name: String,
318    pub num_samples: usize,
319    pub min: f32,
320    pub max: f32,
321    pub mean: f32,
322    pub std_dev: f32,
323    pub p99: f32,
324    pub p9999: f32,
325    pub suggested_int8_scale: f32,
326}
327
328impl CalibSummary {
329    /// One-line human-readable summary.
330    pub fn summary_line(&self) -> String {
331        format!(
332            "[{}] n={} min={:.4} max={:.4} mean={:.4} std={:.4} p99={:.4} p9999={:.4} scale_int8={:.6}",
333            self.layer_name,
334            self.num_samples,
335            self.min,
336            self.max,
337            self.mean,
338            self.std_dev,
339            self.p99,
340            self.p9999,
341            self.suggested_int8_scale,
342        )
343    }
344}
345
346// ─── Calibration database ─────────────────────────────────────────────────────
347
348/// Calibration database: stores stats for all layers.
349pub struct CalibrationDb {
350    layers: HashMap<String, LayerCalibStats>,
351    method: CalibMethod,
352}
353
354impl CalibrationDb {
355    /// Create a new calibration database with the given method.
356    pub fn new(method: CalibMethod) -> Self {
357        Self {
358            layers: HashMap::new(),
359            method,
360        }
361    }
362
363    /// Create with MinMax calibration.
364    pub fn new_minmax() -> Self {
365        Self::new(CalibMethod::MinMax)
366    }
367
368    /// Create with Percentile calibration.
369    pub fn new_percentile(p: f32) -> Self {
370        Self::new(CalibMethod::Percentile(p.clamp(0.0, 1.0)))
371    }
372
373    /// Record activations for a named layer.
374    pub fn record(&mut self, layer_name: &str, activations: &[f32]) {
375        let stats = self
376            .layers
377            .entry(layer_name.to_owned())
378            .or_insert_with(|| LayerCalibStats::new(layer_name));
379        stats.update(activations);
380    }
381
382    /// Get stats for a layer.
383    pub fn get_stats(&self, layer_name: &str) -> Option<&LayerCalibStats> {
384        self.layers.get(layer_name)
385    }
386
387    /// Compute scale for a layer using the configured method.
388    pub fn scale_for_layer(&self, layer_name: &str) -> Option<f32> {
389        self.layers
390            .get(layer_name)
391            .map(|s| s.compute_scale(self.method))
392    }
393
394    /// Export all scales as a map: layer_name → scale.
395    pub fn export_scales(&self) -> HashMap<String, f32> {
396        self.layers
397            .iter()
398            .map(|(name, stats)| (name.clone(), stats.compute_scale(self.method)))
399            .collect()
400    }
401
402    /// Export summaries for all layers.
403    pub fn summaries(&self) -> Vec<CalibSummary> {
404        let mut summaries: Vec<CalibSummary> = self.layers.values().map(|s| s.summary()).collect();
405        summaries.sort_by(|a, b| a.layer_name.cmp(&b.layer_name));
406        summaries
407    }
408
409    /// Total number of calibrated layers.
410    pub fn num_layers(&self) -> usize {
411        self.layers.len()
412    }
413
414    /// Generate a calibration report.
415    pub fn report(&self) -> String {
416        let method_str = match self.method {
417            CalibMethod::MinMax => "MinMax".to_owned(),
418            CalibMethod::Percentile(p) => format!("Percentile({:.4})", p),
419            CalibMethod::Aciq => "ACIQ".to_owned(),
420            CalibMethod::Mse => "MSE".to_owned(),
421        };
422
423        let mut lines = Vec::new();
424        lines.push("=== PTQ Calibration Report ===".to_string());
425        lines.push(format!("Method: {}", method_str));
426        lines.push(format!("Layers: {}", self.layers.len()));
427        lines.push(String::new());
428
429        let summaries = self.summaries();
430        for s in &summaries {
431            lines.push(s.summary_line());
432        }
433
434        lines.push(String::new());
435        lines.push("=== Scales ===".to_string());
436        let scales = self.export_scales();
437        let mut scale_entries: Vec<_> = scales.iter().collect();
438        scale_entries.sort_by_key(|(k, _)| k.as_str());
439        for (name, scale) in scale_entries {
440            lines.push(format!("  {}: {:.8}", name, scale));
441        }
442
443        lines.join("\n")
444    }
445}
446
447// ─── Simulation ───────────────────────────────────────────────────────────────
448
449/// Linear Congruential Generator for deterministic pseudo-random f32 in [-1, 1].
450fn lcg_f32(state: &mut u64) -> f32 {
451    *state = state
452        .wrapping_mul(6_364_136_223_846_793_005)
453        .wrapping_add(1_442_695_040_888_963_407);
454    ((*state >> 32) as f32) / (u32::MAX as f32 + 1.0) * 2.0 - 1.0
455}
456
457/// Simulate a calibration pass (for testing without a real model).
458///
459/// Generates synthetic activation patterns for each layer using a seeded LCG.
460/// Each layer gets `samples_per_layer` activation values drawn from [-1, 1],
461/// scaled by a layer-specific amplitude to simulate natural inter-layer variation.
462pub fn simulate_calibration(
463    db: &mut CalibrationDb,
464    layer_names: &[&str],
465    samples_per_layer: usize,
466    seed: u64,
467) {
468    let mut state = seed;
469
470    for (layer_idx, &layer_name) in layer_names.iter().enumerate() {
471        // Layer-specific amplitude to mimic real network activation diversity.
472        let amplitude = 1.0 + layer_idx as f32 * 0.5;
473        let mut activations = Vec::with_capacity(samples_per_layer);
474
475        for _ in 0..samples_per_layer {
476            let v = lcg_f32(&mut state) * amplitude;
477            activations.push(v);
478        }
479
480        db.record(layer_name, &activations);
481    }
482}
483
484// ─── Validation ───────────────────────────────────────────────────────────────
485
486/// Validation result for a single layer's calibration scale.
487#[derive(Debug, Clone)]
488pub struct CalibValidation {
489    pub layer_name: String,
490    pub scale: f32,
491    pub issues: Vec<String>,
492    pub is_valid: bool,
493}
494
495impl CalibValidation {
496    /// Validate the calibration scale for a layer.
497    pub fn validate(layer_name: &str, stats: &LayerCalibStats, scale: f32) -> Self {
498        let mut issues = Vec::new();
499
500        // Check: scale must be positive and finite
501        if !scale.is_finite() {
502            issues.push(format!("scale is not finite: {}", scale));
503        } else if scale <= 0.0 {
504            issues.push(format!("scale is non-positive: {}", scale));
505        }
506
507        // Check: no samples recorded
508        if stats.num_samples == 0 {
509            issues.push("no calibration samples recorded".to_owned());
510        }
511
512        // Check: scale plausibility — should be at most abs_max / 1.0
513        // (a scale larger than abs_max / 1 would map all values to near-zero)
514        if scale.is_finite() && scale > 0.0 && stats.num_samples > 0 {
515            let abs_max = stats.running_min.abs().max(stats.running_max.abs());
516            if abs_max > 0.0 {
517                // If scale is far too large (> 10x the naive minmax scale), warn
518                let minmax_scale = abs_max / 127.0;
519                if scale > minmax_scale * 10.0 {
520                    issues.push(format!(
521                        "scale {:.6} is >10x the MinMax scale {:.6} (possible overflow)",
522                        scale, minmax_scale
523                    ));
524                }
525                // If scale is far too small (< 1/1000 of minmax), warn
526                if scale < minmax_scale / 1000.0 {
527                    issues.push(format!(
528                        "scale {:.8} is <1/1000 of MinMax scale {:.6} (possible underflow)",
529                        scale, minmax_scale
530                    ));
531                }
532            } else {
533                // abs_max == 0 means all-zero activations
534                issues.push("all activations are zero".to_owned());
535            }
536        }
537
538        // Check: min <= max sanity
539        if stats.num_samples > 0 && stats.running_min > stats.running_max {
540            issues.push(format!(
541                "running_min ({}) > running_max ({})",
542                stats.running_min, stats.running_max
543            ));
544        }
545
546        let is_valid = issues.is_empty();
547
548        Self {
549            layer_name: layer_name.to_owned(),
550            scale,
551            issues,
552            is_valid,
553        }
554    }
555}
556
557/// Validate all layers in a database using its configured method.
558pub fn validate_calibration(db: &CalibrationDb) -> Vec<CalibValidation> {
559    let mut results: Vec<CalibValidation> = db
560        .layers
561        .iter()
562        .map(|(name, stats)| {
563            let scale = stats.compute_scale(db.method);
564            CalibValidation::validate(name, stats, scale)
565        })
566        .collect();
567    results.sort_by(|a, b| a.layer_name.cmp(&b.layer_name));
568    results
569}
570
571// ─── Tests ────────────────────────────────────────────────────────────────────
572
573#[cfg(test)]
574mod tests {
575    use super::*;
576
577    #[test]
578    fn layer_calib_stats_new() {
579        let stats = LayerCalibStats::new("layer0");
580        assert_eq!(stats.layer_name, "layer0");
581        assert_eq!(stats.num_samples, 0);
582        assert_eq!(stats.running_var, 0.0);
583        assert_eq!(stats.running_mean, 0.0);
584    }
585
586    #[test]
587    fn layer_calib_stats_update_single() {
588        let mut stats = LayerCalibStats::new("layer0");
589        let data: Vec<f32> = (0..100).map(|i| i as f32 * 0.01 - 0.5).collect();
590        stats.update(&data);
591        assert_eq!(stats.num_samples, data.len());
592    }
593
594    #[test]
595    fn layer_calib_stats_running_min_max() {
596        let mut stats = LayerCalibStats::new("layer0");
597        stats.update(&[-3.0, 1.0, 2.5]);
598        stats.update(&[0.0, 5.0, -1.0]);
599        assert!((stats.running_min - (-3.0)).abs() < 1e-6);
600        assert!((stats.running_max - 5.0).abs() < 1e-6);
601    }
602
603    #[test]
604    fn layer_calib_stats_std_dev() {
605        let mut stats = LayerCalibStats::new("layer0");
606        // Non-constant data should have positive std dev
607        let data: Vec<f32> = (0..1000).map(|i| (i as f32 - 500.0) * 0.01).collect();
608        stats.update(&data);
609        let sd = stats.std_dev();
610        assert!(
611            sd > 0.0,
612            "std_dev should be positive for non-constant data, got {sd}"
613        );
614    }
615
616    #[test]
617    fn layer_calib_stats_percentile_abs() {
618        let mut stats = LayerCalibStats::new("layer0");
619        let data: Vec<f32> = (1..=100).map(|i| i as f32).collect();
620        stats.update(&data);
621        // p=1.0 should return near max(|x|) = 100.0
622        let p100 = stats.percentile_abs(1.0);
623        assert!(
624            p100 >= 99.0,
625            "p=1.0 percentile should be near 100, got {p100}"
626        );
627    }
628
629    #[test]
630    fn layer_calib_stats_aciq_clip() {
631        let mut stats = LayerCalibStats::new("layer0");
632        // Normal-ish data
633        let data: Vec<f32> = (0..500).map(|i| (i as f32 - 250.0) * 0.1).collect();
634        stats.update(&data);
635        let sd = stats.std_dev();
636        let aciq = stats.aciq_clip();
637        let expected = 2.83 * sd;
638        assert!(
639            (aciq - expected).abs() < 1e-4,
640            "aciq_clip={aciq}, expected 2.83*std_dev={expected}"
641        );
642    }
643
644    #[test]
645    fn layer_calib_stats_compute_scale_minmax() {
646        let mut stats = LayerCalibStats::new("layer0");
647        let data: Vec<f32> = vec![-2.54, 1.0, 0.5, -0.3, std::f32::consts::PI];
648        stats.update(&data);
649        let scale = stats.compute_scale(CalibMethod::MinMax);
650        assert!(scale > 0.0, "MinMax scale must be positive, got {scale}");
651    }
652
653    #[test]
654    fn layer_calib_stats_compute_scale_percentile() {
655        let mut stats = LayerCalibStats::new("layer0");
656        // Data with one large outlier
657        let mut data: Vec<f32> = (0..999).map(|i| (i as f32) * 0.001).collect();
658        data.push(100.0); // outlier
659        stats.update(&data);
660
661        let scale_minmax = stats.compute_scale(CalibMethod::MinMax);
662        let scale_p99 = stats.compute_scale(CalibMethod::Percentile(0.99));
663        // Percentile-based scale should be <= MinMax scale for data with outliers
664        assert!(
665            scale_p99 <= scale_minmax,
666            "Percentile scale {scale_p99} should be <= MinMax scale {scale_minmax}"
667        );
668    }
669
670    #[test]
671    fn calib_summary_summary_line_nonempty() {
672        let mut stats = LayerCalibStats::new("fc1");
673        stats.update(&[0.1, -0.2, 0.3]);
674        let summary = stats.summary();
675        let line = summary.summary_line();
676        assert!(!line.is_empty(), "summary_line should not be empty");
677        assert!(
678            line.contains("fc1"),
679            "summary_line should contain layer name"
680        );
681    }
682
683    #[test]
684    fn calib_db_new_minmax() {
685        let db = CalibrationDb::new_minmax();
686        // Verify method is MinMax by checking that scale_for_layer uses it
687        // (internal field not public, but we can verify behavior via record/scale)
688        assert_eq!(db.num_layers(), 0);
689    }
690
691    #[test]
692    fn calib_db_record_creates_layer() {
693        let mut db = CalibrationDb::new_minmax();
694        assert_eq!(db.num_layers(), 0);
695        db.record("attn.q", &[0.1, -0.2, 0.5]);
696        assert_eq!(db.num_layers(), 1);
697        db.record("attn.k", &[0.3, 0.4]);
698        assert_eq!(db.num_layers(), 2);
699        // Recording to same layer should not increase count
700        db.record("attn.q", &[0.9, -0.1]);
701        assert_eq!(db.num_layers(), 2);
702    }
703
704    #[test]
705    fn calib_db_scale_for_unknown_layer() {
706        let db = CalibrationDb::new_minmax();
707        let result = db.scale_for_layer("nonexistent_layer");
708        assert!(result.is_none(), "Should return None for unknown layer");
709    }
710
711    #[test]
712    fn calib_db_export_scales_all_layers() {
713        let mut db = CalibrationDb::new_minmax();
714        db.record("layer_a", &[1.0, 2.0, 3.0]);
715        db.record("layer_b", &[-1.0, 0.5]);
716        db.record("layer_c", &[0.1, 0.2]);
717        let scales = db.export_scales();
718        assert_eq!(
719            scales.len(),
720            db.num_layers(),
721            "export_scales count should match num_layers"
722        );
723        assert!(scales.contains_key("layer_a"));
724        assert!(scales.contains_key("layer_b"));
725        assert!(scales.contains_key("layer_c"));
726    }
727
728    #[test]
729    fn calib_db_report_nonempty() {
730        let mut db = CalibrationDb::new_percentile(0.999);
731        db.record("block0.ffn", &[0.1, -0.5, 0.3, 0.9]);
732        let report = db.report();
733        assert!(
734            !report.is_empty(),
735            "report() should return a non-empty string"
736        );
737        assert!(
738            report.contains("block0.ffn"),
739            "report should contain layer name"
740        );
741    }
742
743    #[test]
744    fn simulate_calibration_fills_db() {
745        let mut db = CalibrationDb::new_minmax();
746        let layer_names = ["layer0", "layer1", "layer2", "layer3"];
747        simulate_calibration(&mut db, &layer_names, 256, 42);
748        assert_eq!(db.num_layers(), layer_names.len());
749        for &name in &layer_names {
750            let stats = db.get_stats(name).expect("layer should exist");
751            assert_eq!(stats.num_samples, 256);
752        }
753    }
754
755    #[test]
756    fn simulate_calibration_deterministic() {
757        let layer_names = ["attn.q", "attn.k", "attn.v", "ffn.up"];
758        let seed = 12345u64;
759
760        let mut db1 = CalibrationDb::new_minmax();
761        simulate_calibration(&mut db1, &layer_names, 128, seed);
762        let scales1 = db1.export_scales();
763
764        let mut db2 = CalibrationDb::new_minmax();
765        simulate_calibration(&mut db2, &layer_names, 128, seed);
766        let scales2 = db2.export_scales();
767
768        for name in &layer_names {
769            let s1 = scales1[*name];
770            let s2 = scales2[*name];
771            assert!(
772                (s1 - s2).abs() < 1e-8,
773                "scales should be identical for same seed: {name}: {s1} vs {s2}"
774            );
775        }
776    }
777
778    #[test]
779    fn validate_calibration_valid_layer() {
780        let mut stats = LayerCalibStats::new("block0.attn");
781        // Typical activation values
782        let data: Vec<f32> = (0..200).map(|i| (i as f32 - 100.0) * 0.05).collect();
783        stats.update(&data);
784        let scale = stats.compute_scale(CalibMethod::MinMax);
785        let val = CalibValidation::validate("block0.attn", &stats, scale);
786        assert!(
787            val.is_valid,
788            "should be valid for reasonable data: issues={:?}",
789            val.issues
790        );
791        assert!(val.issues.is_empty());
792    }
793
794    #[test]
795    fn validate_calibration_all_valid() {
796        let mut db = CalibrationDb::new_minmax();
797        let layer_names = [
798            "embed",
799            "block0.attn.q",
800            "block0.attn.k",
801            "block0.ffn.up",
802            "block0.ffn.down",
803            "lm_head",
804        ];
805        simulate_calibration(&mut db, &layer_names, 512, 999);
806        let validations = validate_calibration(&db);
807        assert_eq!(validations.len(), layer_names.len());
808        for v in &validations {
809            assert!(
810                v.is_valid,
811                "layer '{}' should be valid after simulate_calibration: issues={:?}",
812                v.layer_name, v.issues
813            );
814        }
815    }
816}