//! 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, /// 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, /// Evaluation metrics to compute pub metrics: Vec, /// 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, /// 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, /// Maximum GPU memory usage in MB pub max_gpu_memory_mb: Option, /// Maximum disk usage for temporary files in MB pub max_disk_usage_mb: Option, /// 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, /// 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>, } 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, /// Enable sign preservation pub preserve_signs: bool, /// Magnitude clipping threshold pub magnitude_clip: Option, /// Enable selective arithmetic pub selective_arithmetic: bool, /// Parameter selection criteria pub selection_criteria: Vec, } 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, /// 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, } 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, /// 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>(path: P) -> Result { 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>(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(()) } }