Initial commit

crates/training/rtx-model-merging/src/config.rs

@@ -0,0 +1,607 @@
+//! Configuration system for model merging operations
+
+use crate::error::{MergeError, Result};
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+/// Main configuration for model merging operations
+#[derive(Debug, Clone, Serialize, Deserialize, Default)]
+pub struct MergeConfig {
+    /// Model loader configuration
+    pub loader: ModelLoaderConfig,
+    /// Validation configuration
+    pub validation: ValidationConfig,
+    /// Evaluation configuration
+    pub evaluation: EvaluationConfig,
+    /// Performance optimization settings
+    pub performance: PerformanceConfig,
+    /// Resource limits
+    pub resources: ResourceConfig,
+}
+
+/// Configuration for model loading
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ModelLoaderConfig {
+    /// Supported model formats
+    pub supported_formats: Vec<String>,
+    /// Memory mapping for large models
+    pub use_memory_mapping: bool,
+    /// Streaming threshold in MB
+    pub streaming_threshold_mb: usize,
+    /// Lazy loading enabled
+    pub lazy_loading: bool,
+    /// Cache loaded models
+    pub cache_models: bool,
+    /// Maximum cache size in MB
+    pub max_cache_size_mb: usize,
+    /// Parallel loading workers
+    pub parallel_workers: usize,
+}
+
+impl Default for ModelLoaderConfig {
+    fn default() -> Self {
+        Self {
+            supported_formats: vec![
+                "safetensors".to_string(),
+                "pytorch".to_string(),
+                "onnx".to_string(),
+                "huggingface".to_string(),
+            ],
+            use_memory_mapping: true,
+            streaming_threshold_mb: 1024,
+            lazy_loading: true,
+            cache_models: true,
+            max_cache_size_mb: 8192,
+            parallel_workers: 4,
+        }
+    }
+}
+
+/// Configuration for model validation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ValidationConfig {
+    /// Enable architecture validation
+    pub validate_architecture: bool,
+    /// Enable parameter compatibility checks
+    pub validate_parameters: bool,
+    /// Enable numerical stability checks
+    pub validate_numerical_stability: bool,
+    /// Enable memory usage validation
+    pub validate_memory_usage: bool,
+    /// Tolerance for numerical comparisons
+    pub numerical_tolerance: f32,
+    /// Maximum allowed parameter deviation
+    pub max_parameter_deviation: f32,
+    /// Validation timeout in seconds
+    pub timeout_seconds: u64,
+}
+
+impl Default for ValidationConfig {
+    fn default() -> Self {
+        Self {
+            validate_architecture: true,
+            validate_parameters: true,
+            validate_numerical_stability: true,
+            validate_memory_usage: true,
+            numerical_tolerance: 1e-6,
+            max_parameter_deviation: 0.1,
+            timeout_seconds: 300,
+        }
+    }
+}
+
+/// Configuration for performance evaluation
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EvaluationConfig {
+    /// Enable performance benchmarking
+    pub enable_benchmarking: bool,
+    /// Benchmark datasets
+    pub benchmark_datasets: Vec<String>,
+    /// Evaluation metrics to compute
+    pub metrics: Vec<String>,
+    /// Evaluation batch size
+    pub batch_size: usize,
+    /// Number of evaluation samples
+    pub num_samples: usize,
+    /// Enable A/B testing
+    pub enable_ab_testing: bool,
+    /// Statistical significance threshold
+    pub significance_threshold: f32,
+}
+
+impl Default for EvaluationConfig {
+    fn default() -> Self {
+        Self {
+            enable_benchmarking: true,
+            benchmark_datasets: vec!["validation".to_string(), "test".to_string()],
+            metrics: vec![
+                "accuracy".to_string(),
+                "loss".to_string(),
+                "f1_score".to_string(),
+                "perplexity".to_string(),
+            ],
+            batch_size: 32,
+            num_samples: 1000,
+            enable_ab_testing: true,
+            significance_threshold: 0.05,
+        }
+    }
+}
+
+/// Configuration for performance optimization
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PerformanceConfig {
+    /// Enable GPU acceleration
+    pub use_gpu: bool,
+    /// GPU device IDs to use
+    pub gpu_devices: Vec<usize>,
+    /// Enable mixed precision
+    pub mixed_precision: bool,
+    /// Enable parallel processing
+    pub parallel_processing: bool,
+    /// Number of worker threads
+    pub num_threads: usize,
+    /// Batch processing size
+    pub batch_size: usize,
+    /// Enable vectorized operations
+    pub vectorized_ops: bool,
+    /// Cache intermediate results
+    pub cache_intermediates: bool,
+}
+
+impl Default for PerformanceConfig {
+    fn default() -> Self {
+        Self {
+            use_gpu: true,
+            gpu_devices: vec![0],
+            mixed_precision: true,
+            parallel_processing: true,
+            num_threads: std::thread::available_parallelism()
+                .map(std::num::NonZero::get)
+                .unwrap_or(4),
+            batch_size: 64,
+            vectorized_ops: true,
+            cache_intermediates: true,
+        }
+    }
+}
+
+/// Configuration for resource limits
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ResourceConfig {
+    /// Maximum memory usage in MB
+    pub max_memory_mb: Option<usize>,
+    /// Maximum GPU memory usage in MB
+    pub max_gpu_memory_mb: Option<usize>,
+    /// Maximum disk usage for temporary files in MB
+    pub max_disk_usage_mb: Option<usize>,
+    /// Operation timeout in seconds
+    pub operation_timeout_sec: u64,
+    /// Enable resource monitoring
+    pub monitor_resources: bool,
+    /// Resource check interval in seconds
+    pub monitor_interval_sec: u64,
+}
+
+impl Default for ResourceConfig {
+    fn default() -> Self {
+        Self {
+            max_memory_mb: None, // No limit by default
+            max_gpu_memory_mb: None,
+            max_disk_usage_mb: Some(10240), // 10GB default
+            operation_timeout_sec: 3600,    // 1 hour
+            monitor_resources: true,
+            monitor_interval_sec: 30,
+        }
+    }
+}
+
+/// Merge strategy configurations
+#[derive(Debug, Clone, Serialize, Deserialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum MergeStrategy {
+    /// TIES (Task-specific Interference Elimination) merging
+    Ties(TiesConfig),
+    /// DARE (Drop And REscale) merging
+    Dare(DareConfig),
+    /// SLERP (Spherical Linear Interpolation) merging
+    Slerp(SlerpConfig),
+    /// Task arithmetic merging
+    TaskArithmetic(TaskArithmeticConfig),
+    /// Fisher information weighted merging
+    Fisher(FisherConfig),
+    /// Model soup ensemble merging
+    ModelSoup(ModelSoupConfig),
+    /// Frankenmerging (layer-wise merging)
+    Frankenmerge(FrankenmergeConfig),
+    /// Progressive merging with validation
+    Progressive(ProgressiveConfig),
+}
+
+/// Configuration for TIES merging
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TiesConfig {
+    /// Sign consistency threshold
+    pub sign_threshold: f32,
+    /// Magnitude threshold for parameter selection
+    pub magnitude_threshold: f32,
+    /// Rescaling method
+    pub rescale_method: RescaleMethod,
+    /// Enable sign voting
+    pub enable_sign_voting: bool,
+    /// Voting threshold
+    pub voting_threshold: f32,
+    /// Density parameter
+    pub density: f32,
+}
+
+impl Default for TiesConfig {
+    fn default() -> Self {
+        Self {
+            sign_threshold: 0.5,
+            magnitude_threshold: 0.1,
+            rescale_method: RescaleMethod::Magnitude,
+            enable_sign_voting: true,
+            voting_threshold: 0.5,
+            density: 0.8,
+        }
+    }
+}
+
+/// Configuration for DARE merging
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DareConfig {
+    /// Drop probability
+    pub drop_probability: f32,
+    /// Rescaling factor
+    pub rescale_factor: f32,
+    /// Random seed for reproducibility
+    pub seed: Option<u64>,
+    /// Enable adaptive dropping
+    pub adaptive_dropping: bool,
+    /// Importance threshold
+    pub importance_threshold: f32,
+}
+
+impl Default for DareConfig {
+    fn default() -> Self {
+        Self {
+            drop_probability: 0.1,
+            rescale_factor: 1.0,
+            seed: None,
+            adaptive_dropping: false,
+            importance_threshold: 0.01,
+        }
+    }
+}
+
+/// Configuration for SLERP merging
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct SlerpConfig {
+    /// Interpolation parameter (0.0 to 1.0)
+    pub t: f32,
+    /// Enable quaternion-based interpolation
+    pub use_quaternions: bool,
+    /// Normalization method
+    pub normalization: NormalizationMethod,
+    /// Enable adaptive interpolation
+    pub adaptive_t: bool,
+    /// Parameter-specific interpolation weights
+    pub parameter_weights: Option<HashMap<String, f32>>,
+}
+
+impl Default for SlerpConfig {
+    fn default() -> Self {
+        Self {
+            t: 0.5,
+            use_quaternions: false,
+            normalization: NormalizationMethod::L2,
+            adaptive_t: false,
+            parameter_weights: None,
+        }
+    }
+}
+
+/// Configuration for task arithmetic
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TaskArithmeticConfig {
+    /// Task vector scaling factors
+    pub scaling_factors: Vec<f32>,
+    /// Enable sign preservation
+    pub preserve_signs: bool,
+    /// Magnitude clipping threshold
+    pub magnitude_clip: Option<f32>,
+    /// Enable selective arithmetic
+    pub selective_arithmetic: bool,
+    /// Parameter selection criteria
+    pub selection_criteria: Vec<String>,
+}
+
+impl Default for TaskArithmeticConfig {
+    fn default() -> Self {
+        Self {
+            scaling_factors: vec![1.0],
+            preserve_signs: true,
+            magnitude_clip: None,
+            selective_arithmetic: false,
+            selection_criteria: vec![],
+        }
+    }
+}
+
+/// Configuration for Fisher information merging
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FisherConfig {
+    /// Fisher information computation method
+    pub fisher_method: FisherMethod,
+    /// Regularization strength
+    pub regularization: f32,
+    /// Enable diagonal approximation
+    pub diagonal_fisher: bool,
+    /// Number of samples for Fisher estimation
+    pub num_samples: usize,
+    /// Enable empirical Fisher
+    pub empirical_fisher: bool,
+}
+
+impl Default for FisherConfig {
+    fn default() -> Self {
+        Self {
+            fisher_method: FisherMethod::Empirical,
+            regularization: 1e-6,
+            diagonal_fisher: true,
+            num_samples: 1000,
+            empirical_fisher: true,
+        }
+    }
+}
+
+/// Configuration for model soup merging
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ModelSoupConfig {
+    /// Ensemble weights
+    pub weights: Vec<f32>,
+    /// Soup method
+    pub soup_method: SoupMethod,
+    /// Enable greedy soup construction
+    pub greedy_soup: bool,
+    /// Validation metric for greedy selection
+    pub validation_metric: String,
+    /// Maximum models in soup
+    pub max_models: usize,
+}
+
+impl Default for ModelSoupConfig {
+    fn default() -> Self {
+        Self {
+            weights: vec![],
+            soup_method: SoupMethod::Uniform,
+            greedy_soup: false,
+            validation_metric: "accuracy".to_string(),
+            max_models: 10,
+        }
+    }
+}
+
+/// Configuration for Frankenmerging
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct FrankenmergeConfig {
+    /// Layer assignment strategy
+    pub layer_assignment: LayerAssignmentStrategy,
+    /// Enable cross-architecture merging
+    pub cross_architecture: bool,
+    /// Layer compatibility checks
+    pub compatibility_checks: bool,
+    /// Enable layer interpolation
+    pub layer_interpolation: bool,
+    /// Interpolation weights per layer
+    pub layer_weights: HashMap<String, f32>,
+}
+
+impl Default for FrankenmergeConfig {
+    fn default() -> Self {
+        Self {
+            layer_assignment: LayerAssignmentStrategy::BestPerformance,
+            cross_architecture: false,
+            compatibility_checks: true,
+            layer_interpolation: false,
+            layer_weights: HashMap::new(),
+        }
+    }
+}
+
+/// Configuration for progressive merging
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ProgressiveConfig {
+    /// Number of progressive steps
+    pub num_steps: usize,
+    /// Validation frequency
+    pub validation_frequency: usize,
+    /// Early stopping patience
+    pub early_stopping_patience: Option<usize>,
+    /// Validation metric threshold
+    pub metric_threshold: f32,
+    /// Enable rollback on failure
+    pub enable_rollback: bool,
+}
+
+impl Default for ProgressiveConfig {
+    fn default() -> Self {
+        Self {
+            num_steps: 10,
+            validation_frequency: 1,
+            early_stopping_patience: Some(3),
+            metric_threshold: 0.01,
+            enable_rollback: true,
+        }
+    }
+}
+
+/// Rescaling methods for TIES
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub enum RescaleMethod {
+    /// Rescale by magnitude
+    Magnitude,
+    /// Rescale by sign consistency
+    SignConsistency,
+    /// No rescaling
+    None,
+}
+
+/// Normalization methods for SLERP
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub enum NormalizationMethod {
+    /// L2 normalization
+    L2,
+    /// L1 normalization
+    L1,
+    /// No normalization
+    None,
+}
+
+/// Fisher information computation methods
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub enum FisherMethod {
+    /// Empirical Fisher information
+    Empirical,
+    /// True Fisher information
+    True,
+    /// Diagonal Fisher approximation
+    Diagonal,
+}
+
+/// Model soup methods
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub enum SoupMethod {
+    /// Uniform averaging
+    Uniform,
+    /// Weighted averaging
+    Weighted,
+    /// Greedy soup construction
+    Greedy,
+}
+
+/// Layer assignment strategies for Frankenmerging
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub enum LayerAssignmentStrategy {
+    /// Assign based on best performance
+    BestPerformance,
+    /// Round-robin assignment
+    RoundRobin,
+    /// Random assignment
+    Random,
+    /// Performance-weighted assignment
+    PerformanceWeighted,
+}
+
+/// Load configuration from file
+pub fn load_config<P: AsRef<std::path::Path>>(path: P) -> Result<MergeConfig> {
+    let content = std::fs::read_to_string(path.as_ref())
+        .map_err(|e| MergeError::io(path.as_ref().display().to_string(), e))?;
+
+    let extension = path
+        .as_ref()
+        .extension()
+        .and_then(|s| s.to_str())
+        .unwrap_or("");
+
+    match extension {
+        "toml" => toml::from_str(&content).map_err(MergeError::from),
+        "json" => serde_json::from_str(&content).map_err(MergeError::from),
+        _ => Err(MergeError::unsupported_format(format!(
+            "Configuration format: {extension}"
+        ))),
+    }
+}
+
+/// Save configuration to file
+pub fn save_config<P: AsRef<std::path::Path>>(config: &MergeConfig, path: P) -> Result<()> {
+    let extension = path
+        .as_ref()
+        .extension()
+        .and_then(|s| s.to_str())
+        .unwrap_or("");
+
+    let content = match extension {
+        "toml" => toml::to_string_pretty(config)
+            .map_err(|e| MergeError::internal(format!("TOML serialization error: {e}")))?,
+        "json" => serde_json::to_string_pretty(config).map_err(MergeError::from)?,
+        _ => {
+            return Err(MergeError::unsupported_format(format!(
+                "Configuration format: {extension}"
+            )));
+        }
+    };
+
+    std::fs::write(path.as_ref(), content)
+        .map_err(|e| MergeError::io(path.as_ref().display().to_string(), e))?;
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tempfile::NamedTempFile;
+
+    #[test]
+    fn test_default_config() {
+        let config = MergeConfig::default();
+        assert!(config.loader.use_memory_mapping);
+        assert!(config.validation.validate_architecture);
+        assert!(config.performance.use_gpu);
+    }
+
+    #[test]
+    fn test_ties_config() {
+        let ties = TiesConfig::default();
+        assert_eq!(ties.sign_threshold, 0.5);
+        assert_eq!(ties.density, 0.8);
+    }
+
+    #[test]
+    fn test_config_serialization() -> Result<()> {
+        let config = MergeConfig::default();
+        let json_str = serde_json::to_string(&config)?;
+        let deserialized: MergeConfig = serde_json::from_str(&json_str)?;
+
+        assert_eq!(
+            config.loader.use_memory_mapping,
+            deserialized.loader.use_memory_mapping
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_config_file_operations() -> Result<()> {
+        let config = MergeConfig::default();
+
+        let temp_file = NamedTempFile::new()
+            .map_err(|e| MergeError::internal(format!("Failed to create temp file: {}", e)))?;
+        let mut temp_path = temp_file.path().to_path_buf();
+        temp_path.set_extension("json");
+
+        save_config(&config, &temp_path)?;
+        let loaded_config = load_config(&temp_path)?;
+
+        assert_eq!(
+            config.loader.parallel_workers,
+            loaded_config.loader.parallel_workers
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_merge_strategy_serialization() -> Result<()> {
+        let strategy = MergeStrategy::Ties(TiesConfig::default());
+        let json_str = serde_json::to_string(&strategy)?;
+        let deserialized: MergeStrategy = serde_json::from_str(&json_str)?;
+
+        match deserialized {
+            MergeStrategy::Ties(_) => (),
+            _ => panic!("Wrong strategy type"),
+        }
+
+        Ok(())
+    }
+}