//! Shared types for FoundationForge - Model Compression & Distillation Hub. //! //! This crate defines the IPC types for model compression, quantization, //! pruning, and knowledge distillation. use serde::{Deserialize, Serialize}; // ============================================================================ // Model Types // ============================================================================ /// Model architecture type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum ModelArchitecture { /// Transformer (BERT, GPT, etc.). Transformer, /// Vision Transformer. ViT, /// Convolutional neural network. CNN, /// Mixture of Experts. MoE, /// Recurrent neural network. RNN, /// Multi-layer perceptron. MLP, /// Diffusion model. Diffusion, /// Custom architecture. Custom, } /// Model information. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ModelInfo { /// Model name. pub name: String, /// Architecture type. pub architecture: ModelArchitecture, /// Number of parameters. pub num_parameters: u64, /// Number of layers. pub num_layers: usize, /// Hidden dimension. pub hidden_dim: usize, /// Vocabulary size (for language models). pub vocab_size: Option, /// Image size (for vision models). pub image_size: Option<(usize, usize)>, /// Original precision (bits). pub precision_bits: u8, /// Model size in bytes. pub size_bytes: u64, } impl Default for ModelInfo { fn default() -> Self { Self { name: "model".to_string(), architecture: ModelArchitecture::Transformer, num_parameters: 7_000_000_000, num_layers: 32, hidden_dim: 4096, vocab_size: Some(32000), image_size: None, precision_bits: 16, size_bytes: 14_000_000_000, } } } // ============================================================================ // Quantization Types // ============================================================================ /// Quantization method. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum QuantMethod { /// Post-Training Quantization. PTQ, /// Quantization-Aware Training. QAT, /// GPTQ (Generative Pre-trained Transformer Quantization). GPTQ, /// AWQ (Activation-aware Weight Quantization). AWQ, /// GGML/GGUF quantization. GGML, /// SmoothQuant. SmoothQuant, /// Dynamic quantization. Dynamic, /// Static quantization. Static, } /// Quantization precision. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum QuantPrecision { /// 2-bit quantization. Int2, /// 3-bit quantization. Int3, /// 4-bit quantization. Int4, /// 8-bit quantization. Int8, /// FP8 (E4M3 or E5M2). FP8, /// FP16 (half precision). FP16, /// BF16 (brain float). BF16, /// Mixed precision. Mixed, } /// Quantization configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct QuantizationConfig { /// Quantization method. pub method: QuantMethod, /// Target precision. pub precision: QuantPrecision, /// Group size for grouped quantization. pub group_size: Option, /// Whether to quantize activations. pub quantize_activations: bool, /// Whether to quantize embeddings. pub quantize_embeddings: bool, /// Calibration dataset size. pub calibration_size: usize, /// Use symmetric quantization. pub symmetric: bool, /// Per-channel quantization. pub per_channel: bool, /// Layers to skip (keep in higher precision). pub skip_layers: Vec, } impl Default for QuantizationConfig { fn default() -> Self { Self { method: QuantMethod::GPTQ, precision: QuantPrecision::Int4, group_size: Some(128), quantize_activations: false, quantize_embeddings: false, calibration_size: 512, symmetric: false, per_channel: true, skip_layers: vec![], } } } // ============================================================================ // Pruning Types // ============================================================================ /// Pruning method. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum PruneMethod { /// Magnitude-based pruning. Magnitude, /// Movement pruning. Movement, /// Structured pruning (remove entire neurons/heads). Structured, /// Unstructured pruning (individual weights). Unstructured, /// Layer-wise pruning. LayerWise, /// Global pruning. Global, /// Lottery ticket hypothesis. LotteryTicket, /// N:M sparsity pattern. NM, } /// Pruning schedule. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub enum PruneSchedule { /// One-shot pruning. OneShot, /// Gradual pruning. Gradual, /// Cubic schedule. Cubic, /// Linear schedule. Linear, /// Exponential schedule. Exponential, } /// Pruning configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PruningConfig { /// Pruning method. pub method: PruneMethod, /// Target sparsity (0.0 to 1.0). pub target_sparsity: f32, /// Pruning schedule. pub schedule: PruneSchedule, /// Number of pruning steps. pub num_steps: usize, /// Initial sparsity. pub initial_sparsity: f32, /// Final sparsity. pub final_sparsity: f32, /// Retraining epochs after pruning. pub retrain_epochs: usize, /// N for N:M sparsity. pub n: Option, /// M for N:M sparsity. pub m: Option, /// Layers to exclude from pruning. pub exclude_layers: Vec, } impl Default for PruningConfig { fn default() -> Self { Self { method: PruneMethod::Magnitude, target_sparsity: 0.5, schedule: PruneSchedule::Gradual, num_steps: 10, initial_sparsity: 0.0, final_sparsity: 0.5, retrain_epochs: 3, n: None, m: None, exclude_layers: vec!["embedding".to_string(), "lm_head".to_string()], } } } // ============================================================================ // Distillation Types // ============================================================================ /// Distillation loss type. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum DistillLoss { /// KL divergence on logits. KL, /// Mean squared error on logits. MSE, /// Cosine similarity on hidden states. Cosine, /// Attention transfer. AttentionTransfer, /// Feature map matching. FeatureMatching, /// Contrastive distillation. Contrastive, /// Combined loss. Combined, } /// Distillation configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DistillationConfig { /// Distillation loss type. pub loss_type: DistillLoss, /// Temperature for softmax. pub temperature: f32, /// Alpha for balancing hard/soft targets. pub alpha: f32, /// Use intermediate layer matching. pub intermediate_matching: bool, /// Layer mapping (teacher layer -> student layer). pub layer_mapping: Vec<(usize, usize)>, /// Number of distillation epochs. pub epochs: usize, /// Batch size. pub batch_size: usize, /// Learning rate. pub learning_rate: f64, /// Use progressive distillation. pub progressive: bool, } impl Default for DistillationConfig { fn default() -> Self { Self { loss_type: DistillLoss::KL, temperature: 4.0, alpha: 0.5, intermediate_matching: true, layer_mapping: vec![], epochs: 10, batch_size: 32, learning_rate: 5e-5, progressive: false, } } } /// Student model configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct StudentConfig { /// Number of layers. pub num_layers: usize, /// Hidden dimension. pub hidden_dim: usize, /// Number of attention heads. pub num_heads: usize, /// Intermediate dimension. pub intermediate_dim: usize, /// Initialize from teacher. pub init_from_teacher: bool, /// Layer indices to copy from teacher. pub copy_layers: Vec, } impl Default for StudentConfig { fn default() -> Self { Self { num_layers: 6, hidden_dim: 768, num_heads: 12, intermediate_dim: 3072, init_from_teacher: true, copy_layers: vec![0, 5, 11, 17, 23, 31], } } } // ============================================================================ // Export Types // ============================================================================ /// Export format. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] pub enum ExportFormat { /// ONNX format. ONNX, /// TensorFlow Lite. TFLite, /// Core ML. CoreML, /// GGUF/GGML format. GGUF, /// SafeTensors format. SafeTensors, /// PyTorch format. PyTorch, /// Custom binary format. Custom, } /// Export configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ExportConfig { /// Export format. pub format: ExportFormat, /// Optimize for inference. pub optimize: bool, /// Include tokenizer. pub include_tokenizer: bool, /// Target device. pub target_device: TargetDevice, /// Output path. pub output_path: String, } impl Default for ExportConfig { fn default() -> Self { Self { format: ExportFormat::SafeTensors, optimize: true, include_tokenizer: true, target_device: TargetDevice::CPU, output_path: "output/model".to_string(), } } } /// Target device for optimization. #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] pub enum TargetDevice { /// CPU inference. CPU, /// NVIDIA GPU. CUDA, /// Apple Metal. Metal, /// WebGPU. WebGPU, /// Mobile (iOS/Android). Mobile, /// Edge device. Edge, } // ============================================================================ // Compression Pipeline Types // ============================================================================ /// Compression pipeline step. #[derive(Debug, Clone, Serialize, Deserialize)] pub enum CompressionStep { /// Quantization step. Quantize(QuantizationConfig), /// Pruning step. Prune(PruningConfig), /// Distillation step. Distill(DistillationConfig), /// Layer removal. RemoveLayers(Vec), /// Vocabulary pruning. PruneVocab { keep_tokens: usize }, /// Head pruning. PruneHeads { target_heads: usize }, } /// Compression pipeline configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PipelineConfig { /// Pipeline name. pub name: String, /// Compression steps. pub steps: Vec, /// Evaluation dataset. pub eval_dataset: Option, /// Target size (bytes). pub target_size: Option, /// Target accuracy degradation. pub max_accuracy_loss: f32, /// Export config. pub export: ExportConfig, } impl Default for PipelineConfig { fn default() -> Self { Self { name: "default".to_string(), steps: vec![CompressionStep::Quantize(QuantizationConfig::default())], eval_dataset: None, target_size: None, max_accuracy_loss: 0.02, export: ExportConfig::default(), } } } // ============================================================================ // Progress Types // ============================================================================ /// Compression progress. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CompressionProgress { /// Current step index. pub step: usize, /// Total steps. pub total_steps: usize, /// Current step name. pub step_name: String, /// Step progress (0.0 to 1.0). pub step_progress: f32, /// Current model size (bytes). pub current_size: u64, /// Original model size (bytes). pub original_size: u64, /// Compression ratio. pub compression_ratio: f32, /// Current perplexity/accuracy. pub current_metric: Option, /// Original perplexity/accuracy. pub original_metric: Option, /// Elapsed time (seconds). pub elapsed_seconds: f64, } /// Compression result. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct CompressionResult { /// Original model info. pub original_model: ModelInfo, /// Compressed model info. pub compressed_model: ModelInfo, /// Compression ratio (original/compressed). pub compression_ratio: f32, /// Speedup ratio. pub speedup_ratio: Option, /// Original metric (perplexity, accuracy, etc.). pub original_metric: f64, /// Compressed metric. pub compressed_metric: f64, /// Metric degradation. pub metric_degradation: f64, /// Total compression time. pub compression_time: f64, /// Steps applied. pub steps_applied: Vec, /// Output path. pub output_path: String, } // ============================================================================ // Benchmark Types // ============================================================================ /// Benchmark configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BenchmarkConfig { /// Benchmark dataset. pub dataset: String, /// Number of samples. pub num_samples: usize, /// Batch size. pub batch_size: usize, /// Measure latency. pub measure_latency: bool, /// Measure memory. pub measure_memory: bool, /// Number of warmup runs. pub warmup_runs: usize, /// Number of benchmark runs. pub benchmark_runs: usize, } impl Default for BenchmarkConfig { fn default() -> Self { Self { dataset: "wikitext".to_string(), num_samples: 1000, batch_size: 1, measure_latency: true, measure_memory: true, warmup_runs: 3, benchmark_runs: 10, } } } /// Benchmark result. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct BenchmarkResult { /// Perplexity (for language models). pub perplexity: Option, /// Accuracy (for classification). pub accuracy: Option, /// Average latency (ms). pub avg_latency_ms: Option, /// P50 latency (ms). pub p50_latency_ms: Option, /// P99 latency (ms). pub p99_latency_ms: Option, /// Throughput (tokens/second or samples/second). pub throughput: Option, /// Memory usage (bytes). pub memory_bytes: Option, /// Peak memory usage (bytes). pub peak_memory_bytes: Option, } // ============================================================================ // Sample Functions // ============================================================================ /// Create a sample model info. #[must_use] pub fn sample_model_info() -> ModelInfo { ModelInfo { name: "llama-7b".to_string(), architecture: ModelArchitecture::Transformer, num_parameters: 7_000_000_000, num_layers: 32, hidden_dim: 4096, vocab_size: Some(32000), image_size: None, precision_bits: 16, size_bytes: 14_000_000_000, } } /// Create a sample quantization config. #[must_use] pub fn sample_quantization_config() -> QuantizationConfig { QuantizationConfig { method: QuantMethod::GPTQ, precision: QuantPrecision::Int4, group_size: Some(128), quantize_activations: false, quantize_embeddings: false, calibration_size: 512, symmetric: false, per_channel: true, skip_layers: vec![], } } /// Create a sample pruning config. #[must_use] pub fn sample_pruning_config() -> PruningConfig { PruningConfig { method: PruneMethod::Magnitude, target_sparsity: 0.5, schedule: PruneSchedule::Gradual, num_steps: 10, initial_sparsity: 0.0, final_sparsity: 0.5, retrain_epochs: 3, n: None, m: None, exclude_layers: vec![], } } /// Create a sample distillation config. #[must_use] pub fn sample_distillation_config() -> DistillationConfig { DistillationConfig { loss_type: DistillLoss::KL, temperature: 4.0, alpha: 0.5, intermediate_matching: true, layer_mapping: vec![(0, 0), (10, 2), (21, 4), (31, 5)], epochs: 10, batch_size: 32, learning_rate: 5e-5, progressive: false, } } /// Create a sample pipeline config. #[must_use] pub fn sample_pipeline_config() -> PipelineConfig { PipelineConfig { name: "quant-prune-distill".to_string(), steps: vec![ CompressionStep::Prune(sample_pruning_config()), CompressionStep::Quantize(sample_quantization_config()), ], eval_dataset: Some("wikitext".to_string()), target_size: Some(4_000_000_000), max_accuracy_loss: 0.02, export: ExportConfig::default(), } } // ============================================================================ // Tests // ============================================================================ #[cfg(test)] mod tests { use super::*; #[test] fn test_model_info() { let info = sample_model_info(); assert_eq!(info.num_parameters, 7_000_000_000); assert_eq!(info.architecture, ModelArchitecture::Transformer); } #[test] fn test_quantization_config() { let config = sample_quantization_config(); assert_eq!(config.method, QuantMethod::GPTQ); assert_eq!(config.precision, QuantPrecision::Int4); } #[test] fn test_pruning_config() { let config = sample_pruning_config(); assert_eq!(config.target_sparsity, 0.5); assert_eq!(config.method, PruneMethod::Magnitude); } #[test] fn test_distillation_config() { let config = sample_distillation_config(); assert_eq!(config.loss_type, DistillLoss::KL); assert!(!config.layer_mapping.is_empty()); } #[test] fn test_pipeline_config() { let config = sample_pipeline_config(); assert!(!config.steps.is_empty()); assert!(config.target_size.is_some()); } #[test] fn test_export_config() { let config = ExportConfig::default(); assert_eq!(config.format, ExportFormat::SafeTensors); assert!(config.optimize); } #[test] fn test_student_config() { let config = StudentConfig::default(); assert!(config.num_layers < 32); // Smaller than typical teacher assert!(config.init_from_teacher); } #[test] fn test_benchmark_config() { let config = BenchmarkConfig::default(); assert!(config.num_samples > 0); assert!(config.measure_latency); } #[test] fn test_serialization() { let config = sample_pipeline_config(); let json = serde_json::to_string(&config).unwrap(); let parsed: PipelineConfig = serde_json::from_str(&json).unwrap(); assert_eq!(parsed.name, config.name); } }