Files
rustytorch/demos/forge-shared/src/lib.rs
T
2026-03-04 00:08:42 +00:00

718 lines
19 KiB
Rust

//! 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<usize>,
/// 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<usize>,
/// 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<String>,
}
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<usize>,
/// M for N:M sparsity.
pub m: Option<usize>,
/// Layers to exclude from pruning.
pub exclude_layers: Vec<String>,
}
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<usize>,
}
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<usize>),
/// 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<CompressionStep>,
/// Evaluation dataset.
pub eval_dataset: Option<String>,
/// Target size (bytes).
pub target_size: Option<u64>,
/// 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<f64>,
/// Original perplexity/accuracy.
pub original_metric: Option<f64>,
/// 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<f32>,
/// 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<String>,
/// 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<f64>,
/// Accuracy (for classification).
pub accuracy: Option<f64>,
/// Average latency (ms).
pub avg_latency_ms: Option<f64>,
/// P50 latency (ms).
pub p50_latency_ms: Option<f64>,
/// P99 latency (ms).
pub p99_latency_ms: Option<f64>,
/// Throughput (tokens/second or samples/second).
pub throughput: Option<f64>,
/// Memory usage (bytes).
pub memory_bytes: Option<u64>,
/// Peak memory usage (bytes).
pub peak_memory_bytes: Option<u64>,
}
// ============================================================================
// 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);
}
}