283 lines
8.6 KiB
Rust
283 lines
8.6 KiB
Rust
//! Pruning module for model compression.
|
|
|
|
use forge_shared::{ModelInfo, PruneMethod, PruneSchedule, PruningConfig};
|
|
|
|
/// Pruner for model weight pruning.
|
|
#[derive(Debug)]
|
|
pub struct Pruner {
|
|
/// Current pruning step.
|
|
current_step: usize,
|
|
/// RNG state.
|
|
rng_state: u64,
|
|
}
|
|
|
|
impl Default for Pruner {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
|
|
impl Pruner {
|
|
/// Create a new pruner.
|
|
pub fn new() -> Self {
|
|
Self {
|
|
current_step: 0,
|
|
rng_state: 42,
|
|
}
|
|
}
|
|
|
|
/// Prune a model and return the fraction of parameters remaining.
|
|
pub fn prune(&mut self, model: &ModelInfo, config: &PruningConfig) -> f32 {
|
|
// Calculate effective sparsity
|
|
let _target_sparsity = config.target_sparsity;
|
|
|
|
// Apply schedule
|
|
let effective_sparsity = self.apply_schedule(config);
|
|
|
|
// Adjust for method characteristics
|
|
let method_factor = self.method_efficiency(config.method);
|
|
|
|
// Calculate remaining parameters
|
|
let remaining = 1.0 - effective_sparsity * method_factor;
|
|
|
|
// Account for excluded layers
|
|
let exclude_fraction = config.exclude_layers.len() as f32 / model.num_layers.max(1) as f32;
|
|
let adjusted_remaining = remaining * (1.0 - exclude_fraction) + exclude_fraction;
|
|
|
|
adjusted_remaining.clamp(0.1, 1.0)
|
|
}
|
|
|
|
/// Apply pruning schedule to get current sparsity.
|
|
fn apply_schedule(&mut self, config: &PruningConfig) -> f32 {
|
|
self.current_step += 1;
|
|
let progress = (self.current_step as f32 / config.num_steps.max(1) as f32).min(1.0);
|
|
|
|
match config.schedule {
|
|
PruneSchedule::OneShot => config.final_sparsity,
|
|
PruneSchedule::Linear => {
|
|
config.initial_sparsity
|
|
+ progress * (config.final_sparsity - config.initial_sparsity)
|
|
}
|
|
PruneSchedule::Gradual => {
|
|
// Gradual: slower start, faster end
|
|
let t = progress * progress;
|
|
config.initial_sparsity + t * (config.final_sparsity - config.initial_sparsity)
|
|
}
|
|
PruneSchedule::Cubic => {
|
|
let t = progress * progress * progress;
|
|
config.initial_sparsity + t * (config.final_sparsity - config.initial_sparsity)
|
|
}
|
|
PruneSchedule::Exponential => {
|
|
let t = 1.0 - (-3.0 * progress).exp();
|
|
config.initial_sparsity + t * (config.final_sparsity - config.initial_sparsity)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Get efficiency factor for pruning method.
|
|
fn method_efficiency(&self, method: PruneMethod) -> f32 {
|
|
match method {
|
|
PruneMethod::Magnitude => 0.95, // Simple, may miss important weights
|
|
PruneMethod::Movement => 0.98, // Considers weight change during training
|
|
PruneMethod::Structured => 1.0, // Full efficiency for structured
|
|
PruneMethod::Unstructured => 0.9, // Sparse storage overhead
|
|
PruneMethod::LayerWise => 0.95,
|
|
PruneMethod::Global => 0.97,
|
|
PruneMethod::LotteryTicket => 0.99, // Very effective
|
|
PruneMethod::NM => 1.0, // Hardware efficient
|
|
}
|
|
}
|
|
|
|
/// Estimate accuracy degradation from pruning.
|
|
pub fn estimate_degradation(&self, config: &PruningConfig) -> f64 {
|
|
let sparsity = config.target_sparsity as f64;
|
|
|
|
// Degradation increases with sparsity
|
|
let base_degradation = sparsity * sparsity * 0.1;
|
|
|
|
// Method factor
|
|
let method_factor = match config.method {
|
|
PruneMethod::Magnitude => 1.2,
|
|
PruneMethod::Movement => 0.8,
|
|
PruneMethod::LotteryTicket => 0.6,
|
|
PruneMethod::Structured => 1.5, // Structured is more aggressive
|
|
_ => 1.0,
|
|
};
|
|
|
|
// Retraining helps recover accuracy
|
|
let retrain_factor = 1.0 / (1.0 + config.retrain_epochs as f64 * 0.1);
|
|
|
|
base_degradation * method_factor * retrain_factor
|
|
}
|
|
|
|
/// Calculate weight importance scores.
|
|
pub fn compute_importance(&mut self, num_weights: usize) -> Vec<f64> {
|
|
let mut importance = Vec::with_capacity(num_weights);
|
|
|
|
for _ in 0..num_weights {
|
|
// Simulate importance scores (would be actual weight magnitudes in practice)
|
|
importance.push(self.random().abs());
|
|
}
|
|
|
|
importance
|
|
}
|
|
|
|
/// Generate pruning mask based on importance.
|
|
pub fn generate_mask(&mut self, importance: &[f64], sparsity: f32) -> Vec<bool> {
|
|
let n = importance.len();
|
|
let num_prune = (n as f32 * sparsity) as usize;
|
|
|
|
// Find threshold
|
|
let mut sorted = importance.to_vec();
|
|
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
|
|
|
let threshold = if num_prune < n {
|
|
sorted[num_prune]
|
|
} else {
|
|
f64::MAX
|
|
};
|
|
|
|
// Generate mask (true = keep, false = prune)
|
|
importance.iter().map(|&imp| imp >= threshold).collect()
|
|
}
|
|
|
|
/// Check if weight pattern matches N:M sparsity.
|
|
pub fn check_nm_pattern(&self, mask: &[bool], n: usize, m: usize) -> bool {
|
|
if !mask.len().is_multiple_of(m) {
|
|
return false;
|
|
}
|
|
|
|
for chunk in mask.chunks(m) {
|
|
let active_count = chunk.iter().filter(|&&x| x).count();
|
|
if active_count != n {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
true
|
|
}
|
|
|
|
/// Random number.
|
|
fn random(&mut self) -> f64 {
|
|
self.rng_state = self
|
|
.rng_state
|
|
.wrapping_mul(6364136223846793005)
|
|
.wrapping_add(1442695040888963407);
|
|
(self.rng_state >> 11) as f64 / (1u64 << 53) as f64
|
|
}
|
|
|
|
/// Reset pruner state.
|
|
pub fn reset(&mut self) {
|
|
self.current_step = 0;
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_pruner_creation() {
|
|
let pruner = Pruner::new();
|
|
assert_eq!(pruner.current_step, 0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_prune_50_percent() {
|
|
let mut pruner = Pruner::new();
|
|
let model = forge_shared::sample_model_info();
|
|
let config = PruningConfig {
|
|
target_sparsity: 0.5,
|
|
final_sparsity: 0.5,
|
|
schedule: PruneSchedule::OneShot,
|
|
..Default::default()
|
|
};
|
|
|
|
let remaining = pruner.prune(&model, &config);
|
|
assert!(remaining > 0.4 && remaining < 0.6);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gradual_schedule() {
|
|
let mut pruner = Pruner::new();
|
|
let config = PruningConfig {
|
|
initial_sparsity: 0.0,
|
|
final_sparsity: 0.9,
|
|
schedule: PruneSchedule::Gradual,
|
|
num_steps: 10,
|
|
..Default::default()
|
|
};
|
|
|
|
let mut prev_sparsity = 0.0;
|
|
for _ in 0..10 {
|
|
let sparsity = pruner.apply_schedule(&config);
|
|
assert!(sparsity >= prev_sparsity);
|
|
prev_sparsity = sparsity;
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_estimate_degradation() {
|
|
let pruner = Pruner::new();
|
|
|
|
let config_low = PruningConfig {
|
|
target_sparsity: 0.3,
|
|
..Default::default()
|
|
};
|
|
let config_high = PruningConfig {
|
|
target_sparsity: 0.9,
|
|
..Default::default()
|
|
};
|
|
|
|
let deg_low = pruner.estimate_degradation(&config_low);
|
|
let deg_high = pruner.estimate_degradation(&config_high);
|
|
|
|
// Higher sparsity should have higher degradation
|
|
assert!(deg_high > deg_low);
|
|
}
|
|
|
|
#[test]
|
|
fn test_compute_importance() {
|
|
let mut pruner = Pruner::new();
|
|
let importance = pruner.compute_importance(100);
|
|
|
|
assert_eq!(importance.len(), 100);
|
|
for &imp in &importance {
|
|
assert!(imp >= 0.0 && imp <= 1.0);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mask() {
|
|
let mut pruner = Pruner::new();
|
|
let importance = vec![0.1, 0.5, 0.2, 0.8, 0.3, 0.9, 0.4, 0.7, 0.6, 0.05];
|
|
|
|
let mask = pruner.generate_mask(&importance, 0.5);
|
|
|
|
assert_eq!(mask.len(), 10);
|
|
let kept = mask.iter().filter(|&&x| x).count();
|
|
assert_eq!(kept, 5); // 50% kept
|
|
}
|
|
|
|
#[test]
|
|
fn test_nm_pattern() {
|
|
let pruner = Pruner::new();
|
|
|
|
// Valid 2:4 pattern
|
|
let mask_valid = vec![true, true, false, false, true, false, true, false];
|
|
assert!(pruner.check_nm_pattern(&mask_valid, 2, 4));
|
|
|
|
// Invalid pattern
|
|
let mask_invalid = vec![true, true, true, false];
|
|
assert!(!pruner.check_nm_pattern(&mask_invalid, 2, 4));
|
|
}
|
|
|
|
#[test]
|
|
fn test_reset() {
|
|
let mut pruner = Pruner::new();
|
|
pruner.current_step = 5;
|
|
pruner.reset();
|
|
assert_eq!(pruner.current_step, 0);
|
|
}
|
|
}
|