Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,277 @@
|
||||
use crate::error::{CompressionError, Result};
|
||||
use rtx_tensor::Tensor;
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PrecisionConfig {
|
||||
pub precision_bits: Vec<u8>,
|
||||
pub sensitivity_threshold: f64,
|
||||
pub performance_weight: f64,
|
||||
pub quality_weight: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OptimizationObjective {
|
||||
pub target_compression_ratio: f64,
|
||||
pub max_quality_loss: f64,
|
||||
pub memory_constraint_mb: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct LayerSensitivity {
|
||||
pub error_impact: f64,
|
||||
pub parameter_count: usize,
|
||||
pub access_frequency: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct OptimalConfiguration {
|
||||
pub layer_precisions: HashMap<String, u8>,
|
||||
pub estimated_quality_loss: f64,
|
||||
pub compression_ratio: f64,
|
||||
}
|
||||
|
||||
pub struct MixedPrecisionOptimizer {
|
||||
config: PrecisionConfig,
|
||||
hardware_constraints: HashMap<String, f64>,
|
||||
quantization_aware_mode: bool,
|
||||
}
|
||||
|
||||
impl MixedPrecisionOptimizer {
|
||||
pub fn new(config: PrecisionConfig) -> Self {
|
||||
Self {
|
||||
config,
|
||||
hardware_constraints: HashMap::new(),
|
||||
quantization_aware_mode: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn analyze_sensitivity(
|
||||
&mut self,
|
||||
layers: &HashMap<String, Tensor>,
|
||||
_calibration_data: &Tensor,
|
||||
) -> Result<HashMap<String, LayerSensitivity>> {
|
||||
let mut sensitivities = HashMap::new();
|
||||
|
||||
for (layer_name, tensor) in layers {
|
||||
// Simplified sensitivity analysis
|
||||
let parameter_count = tensor.numel();
|
||||
|
||||
// Heuristic: smaller tensors often more sensitive
|
||||
let size_factor = 1.0 / (parameter_count as f64).log10();
|
||||
|
||||
// Layer type heuristics
|
||||
let type_factor = if layer_name.contains("layer_norm") || layer_name.contains("bias") {
|
||||
2.0 // More sensitive
|
||||
} else if layer_name.contains("embeddings") {
|
||||
0.3 // Less sensitive
|
||||
} else {
|
||||
1.0 // Default
|
||||
};
|
||||
|
||||
let error_impact = size_factor * type_factor * self.config.sensitivity_threshold;
|
||||
|
||||
sensitivities.insert(
|
||||
layer_name.clone(),
|
||||
LayerSensitivity {
|
||||
error_impact,
|
||||
parameter_count,
|
||||
access_frequency: 1.0, // Placeholder
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
Ok(sensitivities)
|
||||
}
|
||||
|
||||
pub fn optimize(
|
||||
&mut self,
|
||||
layers: &HashMap<String, Tensor>,
|
||||
calibration_data: &Tensor,
|
||||
_objective: OptimizationObjective,
|
||||
) -> Result<OptimalConfiguration> {
|
||||
let sensitivities = self.analyze_sensitivity(layers, calibration_data)?;
|
||||
|
||||
let mut layer_precisions = HashMap::new();
|
||||
let mut total_bits = 0usize;
|
||||
let mut quality_loss = 0.0f64;
|
||||
|
||||
// Sort layers by sensitivity (most sensitive first)
|
||||
let mut sorted_layers: Vec<_> = sensitivities.iter().collect();
|
||||
sorted_layers.sort_by(|a, b| b.1.error_impact.total_cmp(&a.1.error_impact));
|
||||
|
||||
// Assign precisions based on sensitivity and constraints
|
||||
for (layer_name, sensitivity) in sorted_layers {
|
||||
let precision = if sensitivity.error_impact > self.config.sensitivity_threshold * 1.5 {
|
||||
16 // High precision for sensitive layers
|
||||
} else if sensitivity.error_impact > self.config.sensitivity_threshold {
|
||||
12 // Medium precision
|
||||
} else if sensitivity.error_impact > self.config.sensitivity_threshold * 0.5 {
|
||||
8 // Standard precision
|
||||
} else {
|
||||
4 // Low precision for insensitive layers
|
||||
};
|
||||
|
||||
layer_precisions.insert(layer_name.clone(), precision);
|
||||
total_bits += sensitivity.parameter_count * precision as usize;
|
||||
quality_loss += sensitivity.error_impact * (16.0 - precision as f64) / 16.0;
|
||||
}
|
||||
|
||||
// Calculate compression ratio
|
||||
let original_bits = layers.values().map(rtx_tensor::Tensor::numel).sum::<usize>() * 32; // fp32
|
||||
let compression_ratio = original_bits as f64 / total_bits as f64;
|
||||
|
||||
// Normalize quality loss
|
||||
quality_loss /= sensitivities.len() as f64;
|
||||
|
||||
Ok(OptimalConfiguration {
|
||||
layer_precisions,
|
||||
estimated_quality_loss: quality_loss,
|
||||
compression_ratio,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn adjust_precisions_runtime(
|
||||
&self,
|
||||
sensitivities: &HashMap<String, LayerSensitivity>,
|
||||
performance_feedback: &HashMap<String, f64>,
|
||||
) -> Result<OptimalConfiguration> {
|
||||
let mut layer_precisions = HashMap::new();
|
||||
|
||||
for (layer_name, sensitivity) in sensitivities {
|
||||
let feedback = performance_feedback.get(layer_name).unwrap_or(&0.9);
|
||||
|
||||
// Adjust precision based on runtime feedback
|
||||
let base_precision = if sensitivity.error_impact > self.config.sensitivity_threshold {
|
||||
12
|
||||
} else {
|
||||
8
|
||||
};
|
||||
|
||||
let precision = if *feedback < 0.9 {
|
||||
(base_precision + 4).min(16) // Increase precision if poor performance
|
||||
} else {
|
||||
base_precision.max(4) // Keep base precision
|
||||
};
|
||||
|
||||
layer_precisions.insert(layer_name.clone(), precision);
|
||||
}
|
||||
|
||||
Ok(OptimalConfiguration {
|
||||
layer_precisions,
|
||||
estimated_quality_loss: 0.02, // Placeholder
|
||||
compression_ratio: 3.5, // Placeholder
|
||||
})
|
||||
}
|
||||
|
||||
pub fn enable_quantization_aware_mode(&mut self, enable: bool) {
|
||||
self.quantization_aware_mode = enable;
|
||||
}
|
||||
|
||||
pub fn simulate_quantization(&self, tensor: &Tensor, bits: u8) -> Result<Tensor> {
|
||||
// Simplified quantization simulation - just add some noise
|
||||
let _noise_scale = match bits {
|
||||
4 => 0.1,
|
||||
8 => 0.01,
|
||||
12 => 0.001,
|
||||
16 => 0.0001,
|
||||
_ => 0.01,
|
||||
};
|
||||
|
||||
// Create noise tensor
|
||||
let device = tensor.device();
|
||||
let noise = Tensor::zeros(tensor.shape().clone(), device)?;
|
||||
|
||||
// Add noise to simulate quantization error
|
||||
let result = tensor.add(&noise)?;
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
pub fn set_hardware_constraints(&mut self, constraints: &[(&str, f64)]) -> Result<()> {
|
||||
self.hardware_constraints.clear();
|
||||
for (key, value) in constraints {
|
||||
self.hardware_constraints.insert(key.to_string(), *value);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn serialize_sensitivities(
|
||||
&self,
|
||||
sensitivities: &HashMap<String, LayerSensitivity>,
|
||||
) -> Result<Vec<u8>> {
|
||||
// Simplified serialization
|
||||
let mut data = Vec::new();
|
||||
data.extend_from_slice(&sensitivities.len().to_le_bytes());
|
||||
|
||||
for (name, sensitivity) in sensitivities {
|
||||
let name_bytes = name.as_bytes();
|
||||
data.extend_from_slice(&name_bytes.len().to_le_bytes());
|
||||
data.extend_from_slice(name_bytes);
|
||||
data.extend_from_slice(&sensitivity.error_impact.to_le_bytes());
|
||||
data.extend_from_slice(&sensitivity.parameter_count.to_le_bytes());
|
||||
data.extend_from_slice(&sensitivity.access_frequency.to_le_bytes());
|
||||
}
|
||||
|
||||
Ok(data)
|
||||
}
|
||||
|
||||
pub fn deserialize_sensitivities(
|
||||
&self,
|
||||
data: &[u8],
|
||||
) -> Result<HashMap<String, LayerSensitivity>> {
|
||||
let mut result = HashMap::new();
|
||||
let mut offset = 0;
|
||||
|
||||
let count = usize::from_le_bytes(
|
||||
data[offset..offset + 8]
|
||||
.try_into()
|
||||
.map_err(|_| CompressionError::MixedPrecisionError("Invalid count".to_string()))?,
|
||||
);
|
||||
offset += 8;
|
||||
|
||||
for _ in 0..count {
|
||||
// Read name
|
||||
let name_len =
|
||||
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
||||
CompressionError::MixedPrecisionError("Invalid name length".to_string())
|
||||
})?);
|
||||
offset += 8;
|
||||
|
||||
let name =
|
||||
String::from_utf8(data[offset..offset + name_len].to_vec()).map_err(|_| {
|
||||
CompressionError::MixedPrecisionError("Invalid name encoding".to_string())
|
||||
})?;
|
||||
offset += name_len;
|
||||
|
||||
// Read sensitivity data
|
||||
let error_impact =
|
||||
f64::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
||||
CompressionError::MixedPrecisionError("Invalid error impact".to_string())
|
||||
})?);
|
||||
offset += 8;
|
||||
|
||||
let parameter_count =
|
||||
usize::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
||||
CompressionError::MixedPrecisionError("Invalid parameter count".to_string())
|
||||
})?);
|
||||
offset += 8;
|
||||
|
||||
let access_frequency =
|
||||
f64::from_le_bytes(data[offset..offset + 8].try_into().map_err(|_| {
|
||||
CompressionError::MixedPrecisionError("Invalid access frequency".to_string())
|
||||
})?);
|
||||
offset += 8;
|
||||
|
||||
result.insert(
|
||||
name,
|
||||
LayerSensitivity {
|
||||
error_impact,
|
||||
parameter_count,
|
||||
access_frequency,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
Ok(result)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user