1052 lines
36 KiB
Rust
1052 lines
36 KiB
Rust
//! Magnitude-based weight pruning for neural networks
|
|
//!
|
|
//! This module implements magnitude-based pruning techniques that remove parameters
|
|
//! with the smallest absolute values, a fundamental approach in neural network compression.
|
|
//! Magnitude pruning is based on the hypothesis that small weights contribute less to
|
|
//! the network's output and can be safely removed.
|
|
//!
|
|
//! ## Supported Pruning Types
|
|
//!
|
|
//! - **Global Magnitude Pruning**: Remove smallest weights across the entire network
|
|
//! - **Layer-wise Magnitude Pruning**: Remove smallest weights within each layer independently
|
|
//! - **Gradual Pruning**: Apply pruning schedules during training (linear, polynomial, exponential)
|
|
//! - **Fine-grained Pruning**: Element-wise weight removal (unstructured)
|
|
//! - **Structured Magnitude Pruning**: Channel-wise, filter-wise pruning patterns
|
|
//!
|
|
//! ## Key Features
|
|
//!
|
|
//! - Configurable sparsity ratios (0-100%)
|
|
//! - Multiple pruning granularities (element, channel, filter)
|
|
//! - Pruning schedules for gradual sparsification during training
|
|
//! - Mask generation and application for efficient training
|
|
//! - Integration with existing sparse tensor infrastructure
|
|
//! - Support for both weight and activation pruning
|
|
//!
|
|
//! ## Design Principles
|
|
//!
|
|
//! - **Memory Efficiency**: Leverage sparse tensor formats for storage
|
|
//! - **Training Integration**: Seamless integration with optimizers and training loops
|
|
//! - **Hardware Acceleration**: Support for structured patterns that map to modern accelerators
|
|
//! - **Flexibility**: Support various pruning strategies and schedules
|
|
//!
|
|
//! ## Usage Example
|
|
//!
|
|
//! ```rust
|
|
//! use rtx_transformers::layers::{MagnitudePruning, PruningConfig, PruningStrategy};
|
|
//! use rtx_tensor::{Tensor, Device};
|
|
//!
|
|
//! let device = Device::cpu();
|
|
//! let weights = Tensor::randn(&[512, 1024], &device).unwrap();
|
|
//!
|
|
//! // Global magnitude pruning with 50% sparsity
|
|
//! let config = PruningConfig::global(0.5);
|
|
//! let pruned_weights = MagnitudePruning::prune(&weights, &config).unwrap();
|
|
//!
|
|
//! // Layer-wise pruning with gradual schedule
|
|
//! let schedule_config = PruningConfig::layerwise(0.8)
|
|
//! .with_schedule(PruningSchedule::polynomial(100, 3.0));
|
|
//! let mask = MagnitudePruning::create_mask(&weights, &schedule_config).unwrap();
|
|
//! ```
|
|
|
|
use crate::{Result, TransformerError};
|
|
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
|
|
/// Configuration for magnitude-based pruning
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub struct PruningConfig {
|
|
/// Pruning strategy to apply
|
|
pub strategy: PruningStrategy,
|
|
/// Target sparsity ratio (0.0 to 1.0)
|
|
pub sparsity_ratio: f64,
|
|
/// Pruning granularity
|
|
pub granularity: PruningGranularity,
|
|
/// Optional pruning schedule for gradual pruning
|
|
pub schedule: Option<PruningSchedule>,
|
|
/// Random seed for deterministic pruning
|
|
pub seed: Option<u64>,
|
|
}
|
|
|
|
/// Strategy for selecting weights to prune
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum PruningStrategy {
|
|
/// Global magnitude pruning across all parameters
|
|
Global,
|
|
/// Layer-wise magnitude pruning within each layer
|
|
LayerWise,
|
|
/// Structured pruning (channels, filters)
|
|
Structured {
|
|
/// Type of structured pattern
|
|
pattern_type: StructuredPattern,
|
|
},
|
|
}
|
|
|
|
/// Granularity of pruning operations
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum PruningGranularity {
|
|
/// Element-wise pruning (unstructured)
|
|
Element,
|
|
/// Channel-wise pruning
|
|
Channel,
|
|
/// Filter-wise pruning
|
|
Filter,
|
|
/// Block-wise pruning
|
|
Block { height: usize, width: usize },
|
|
}
|
|
|
|
/// Structured pruning patterns
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum StructuredPattern {
|
|
/// Channel pruning
|
|
Channel,
|
|
/// Filter pruning
|
|
Filter,
|
|
/// N:M sparsity pattern
|
|
NM { n: usize, m: usize },
|
|
}
|
|
|
|
/// Pruning schedule for gradual sparsification during training
|
|
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
|
pub enum PruningSchedule {
|
|
/// Linear schedule from initial to final sparsity
|
|
Linear {
|
|
/// Initial sparsity ratio
|
|
initial_sparsity: f64,
|
|
/// Final sparsity ratio
|
|
final_sparsity: f64,
|
|
/// Duration in training steps
|
|
duration_steps: usize,
|
|
},
|
|
/// Polynomial schedule
|
|
Polynomial {
|
|
/// Initial sparsity ratio
|
|
initial_sparsity: f64,
|
|
/// Final sparsity ratio
|
|
final_sparsity: f64,
|
|
/// Duration in training steps
|
|
duration_steps: usize,
|
|
/// Polynomial exponent
|
|
exponent: f64,
|
|
},
|
|
/// Exponential schedule
|
|
Exponential {
|
|
/// Initial sparsity ratio
|
|
initial_sparsity: f64,
|
|
/// Final sparsity ratio
|
|
final_sparsity: f64,
|
|
/// Duration in training steps
|
|
duration_steps: usize,
|
|
/// Decay rate
|
|
decay_rate: f64,
|
|
},
|
|
}
|
|
|
|
/// Binary mask for indicating which weights to keep/prune
|
|
#[derive(Debug, Clone)]
|
|
pub struct PruningMask {
|
|
/// Binary mask values (true = keep, false = prune)
|
|
mask: Vec<bool>,
|
|
/// Shape of the mask
|
|
shape: Vec<usize>,
|
|
/// Number of parameters kept
|
|
kept_parameters: usize,
|
|
/// Total number of parameters
|
|
total_parameters: usize,
|
|
}
|
|
|
|
/// Statistics about pruning results
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct PruningStats {
|
|
/// Original parameter count
|
|
pub original_params: usize,
|
|
/// Parameters after pruning
|
|
pub pruned_params: usize,
|
|
/// Achieved sparsity ratio
|
|
pub sparsity_ratio: f64,
|
|
/// Compression ratio
|
|
pub compression_ratio: f64,
|
|
/// Memory savings (in bytes, estimated)
|
|
pub memory_savings: usize,
|
|
}
|
|
|
|
/// Optimizer integration for pruned training
|
|
#[derive(Debug, Clone)]
|
|
pub struct PrunedOptimizer {
|
|
/// Active pruning masks for each parameter group
|
|
pub masks: HashMap<String, PruningMask>,
|
|
/// Update count for scheduling
|
|
pub update_count: usize,
|
|
/// Pruning schedule configuration
|
|
pub schedule_config: Option<PruningSchedule>,
|
|
}
|
|
|
|
/// Training state for pruned networks
|
|
#[derive(Debug, Clone)]
|
|
pub struct PrunedTrainingState {
|
|
/// Current training step
|
|
pub current_step: usize,
|
|
/// Active sparsity ratio
|
|
pub current_sparsity: f64,
|
|
/// Accumulated pruning statistics
|
|
pub stats: PruningStats,
|
|
/// Whether gradual pruning is active
|
|
pub gradual_pruning_active: bool,
|
|
}
|
|
|
|
/// Gradual pruning controller for training loops
|
|
#[derive(Debug)]
|
|
pub struct GradualPruningController {
|
|
/// Pruning configuration
|
|
config: PruningConfig,
|
|
/// Current training state
|
|
state: PrunedTrainingState,
|
|
/// Parameter tensors being pruned
|
|
parameters: HashMap<String, DenseTensor>,
|
|
/// Current masks for each parameter
|
|
masks: HashMap<String, PruningMask>,
|
|
}
|
|
|
|
/// Main struct for magnitude-based pruning operations
|
|
#[derive(Debug)]
|
|
pub struct MagnitudePruning;
|
|
|
|
/// Placeholder for tensor-like data structure
|
|
/// This will be replaced with actual Tensor integration once tensor API is stable
|
|
#[derive(Debug, Clone)]
|
|
pub struct DenseTensor {
|
|
pub data: Vec<f32>,
|
|
pub shape: Vec<usize>,
|
|
}
|
|
|
|
impl DenseTensor {
|
|
pub fn new(data: Vec<f32>, shape: Vec<usize>) -> Self {
|
|
assert_eq!(data.len(), shape.iter().product::<usize>());
|
|
Self { data, shape }
|
|
}
|
|
|
|
pub fn shape(&self) -> &[usize] {
|
|
&self.shape
|
|
}
|
|
|
|
pub fn data(&self) -> &[f32] {
|
|
&self.data
|
|
}
|
|
|
|
pub fn zeros(shape: Vec<usize>) -> Self {
|
|
let total_elements = shape.iter().product::<usize>();
|
|
Self {
|
|
data: vec![0.0; total_elements],
|
|
shape,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl PruningConfig {
|
|
/// Create global magnitude pruning configuration
|
|
pub fn global(sparsity_ratio: f64) -> Self {
|
|
Self {
|
|
strategy: PruningStrategy::Global,
|
|
sparsity_ratio,
|
|
granularity: PruningGranularity::Element,
|
|
schedule: None,
|
|
seed: None,
|
|
}
|
|
}
|
|
|
|
/// Create layer-wise magnitude pruning configuration
|
|
pub fn layerwise(sparsity_ratio: f64) -> Self {
|
|
Self {
|
|
strategy: PruningStrategy::LayerWise,
|
|
sparsity_ratio,
|
|
granularity: PruningGranularity::Element,
|
|
schedule: None,
|
|
seed: None,
|
|
}
|
|
}
|
|
|
|
/// Create structured magnitude pruning configuration
|
|
pub fn structured(sparsity_ratio: f64, pattern: StructuredPattern) -> Self {
|
|
let granularity = match pattern {
|
|
StructuredPattern::Channel => PruningGranularity::Channel,
|
|
StructuredPattern::Filter => PruningGranularity::Filter,
|
|
StructuredPattern::NM { .. } => PruningGranularity::Element,
|
|
};
|
|
|
|
Self {
|
|
strategy: PruningStrategy::Structured {
|
|
pattern_type: pattern,
|
|
},
|
|
sparsity_ratio,
|
|
granularity,
|
|
schedule: None,
|
|
seed: None,
|
|
}
|
|
}
|
|
|
|
/// Add pruning schedule
|
|
pub fn with_schedule(mut self, schedule: PruningSchedule) -> Self {
|
|
self.schedule = Some(schedule);
|
|
self
|
|
}
|
|
|
|
/// Set random seed for deterministic pruning
|
|
pub fn with_seed(mut self, seed: u64) -> Self {
|
|
self.seed = Some(seed);
|
|
self
|
|
}
|
|
|
|
/// Validate configuration parameters
|
|
pub fn validate(&self) -> Result<()> {
|
|
// Validate sparsity ratio
|
|
if !(0.0..=1.0).contains(&self.sparsity_ratio) {
|
|
return Err(TransformerError::Config(
|
|
"Sparsity ratio must be between 0.0 and 1.0".into()
|
|
));
|
|
}
|
|
|
|
// Validate schedule if present
|
|
if let Some(ref schedule) = self.schedule {
|
|
schedule.validate()?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl PruningSchedule {
|
|
/// Create linear pruning schedule
|
|
pub fn linear(duration_steps: usize, final_sparsity: f64) -> Self {
|
|
Self::Linear {
|
|
initial_sparsity: 0.0,
|
|
final_sparsity,
|
|
duration_steps,
|
|
}
|
|
}
|
|
|
|
/// Create polynomial pruning schedule
|
|
pub fn polynomial(duration_steps: usize, exponent: f64) -> Self {
|
|
Self::Polynomial {
|
|
initial_sparsity: 0.0,
|
|
final_sparsity: 0.9, // Default to 90% sparsity
|
|
duration_steps,
|
|
exponent,
|
|
}
|
|
}
|
|
|
|
/// Create exponential pruning schedule
|
|
pub fn exponential(duration_steps: usize, decay_rate: f64) -> Self {
|
|
Self::Exponential {
|
|
initial_sparsity: 0.0,
|
|
final_sparsity: 0.9, // Default to 90% sparsity
|
|
duration_steps,
|
|
decay_rate,
|
|
}
|
|
}
|
|
|
|
/// Compute sparsity ratio at given training step
|
|
pub fn compute_sparsity(&self, current_step: usize) -> f64 {
|
|
match self {
|
|
Self::Linear { initial_sparsity, final_sparsity, duration_steps } => {
|
|
if current_step >= *duration_steps {
|
|
return *final_sparsity;
|
|
}
|
|
let progress = current_step as f64 / *duration_steps as f64;
|
|
initial_sparsity + progress * (final_sparsity - initial_sparsity)
|
|
}
|
|
Self::Polynomial { initial_sparsity, final_sparsity, duration_steps, exponent } => {
|
|
if current_step >= *duration_steps {
|
|
return *final_sparsity;
|
|
}
|
|
let progress = current_step as f64 / *duration_steps as f64;
|
|
let poly_progress = progress.powf(*exponent);
|
|
initial_sparsity + poly_progress * (final_sparsity - initial_sparsity)
|
|
}
|
|
Self::Exponential { initial_sparsity, final_sparsity, duration_steps, decay_rate } => {
|
|
if current_step >= *duration_steps {
|
|
return *final_sparsity;
|
|
}
|
|
let progress = current_step as f64 / *duration_steps as f64;
|
|
let exp_progress = 1.0 - (-decay_rate * progress).exp();
|
|
initial_sparsity + exp_progress * (final_sparsity - initial_sparsity)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Validate schedule parameters
|
|
pub fn validate(&self) -> Result<()> {
|
|
match self {
|
|
Self::Linear { initial_sparsity, final_sparsity, duration_steps } => {
|
|
Self::validate_common_params(*initial_sparsity, *final_sparsity, *duration_steps)?;
|
|
}
|
|
Self::Polynomial { initial_sparsity, final_sparsity, duration_steps, exponent } => {
|
|
Self::validate_common_params(*initial_sparsity, *final_sparsity, *duration_steps)?;
|
|
if *exponent <= 0.0 {
|
|
return Err(TransformerError::Config("Polynomial exponent must be positive".into()));
|
|
}
|
|
}
|
|
Self::Exponential { initial_sparsity, final_sparsity, duration_steps, decay_rate } => {
|
|
Self::validate_common_params(*initial_sparsity, *final_sparsity, *duration_steps)?;
|
|
if *decay_rate <= 0.0 {
|
|
return Err(TransformerError::Config("Exponential decay rate must be positive".into()));
|
|
}
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn validate_common_params(initial: f64, final_val: f64, duration: usize) -> Result<()> {
|
|
if !(0.0..=1.0).contains(&initial) {
|
|
return Err(TransformerError::Config("Initial sparsity must be between 0.0 and 1.0".into()));
|
|
}
|
|
if !(0.0..=1.0).contains(&final_val) {
|
|
return Err(TransformerError::Config("Final sparsity must be between 0.0 and 1.0".into()));
|
|
}
|
|
if duration == 0 {
|
|
return Err(TransformerError::Config("Duration steps must be greater than 0".into()));
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl PruningMask {
|
|
/// Create new pruning mask
|
|
pub fn new(mask: Vec<bool>, shape: Vec<usize>) -> Self {
|
|
let total_parameters = shape.iter().product::<usize>();
|
|
let kept_parameters = mask.iter().filter(|&&x| x).count();
|
|
|
|
Self {
|
|
mask,
|
|
shape,
|
|
kept_parameters,
|
|
total_parameters,
|
|
}
|
|
}
|
|
|
|
/// Get mask values
|
|
pub fn mask(&self) -> &[bool] {
|
|
&self.mask
|
|
}
|
|
|
|
/// Get mask shape
|
|
pub fn shape(&self) -> &[usize] {
|
|
&self.shape
|
|
}
|
|
|
|
/// Get sparsity ratio
|
|
pub fn sparsity_ratio(&self) -> f64 {
|
|
1.0 - (self.kept_parameters as f64 / self.total_parameters as f64)
|
|
}
|
|
|
|
/// Apply mask to tensor (element-wise multiplication)
|
|
pub fn apply(&self, tensor: &DenseTensor) -> Result<DenseTensor> {
|
|
if tensor.shape != self.shape {
|
|
return Err(TransformerError::Config("Tensor and mask shapes must match".into()));
|
|
}
|
|
|
|
let masked_data: Vec<f32> = tensor.data.iter()
|
|
.zip(self.mask.iter())
|
|
.map(|(&value, &keep)| if keep { value } else { 0.0 })
|
|
.collect();
|
|
|
|
Ok(DenseTensor::new(masked_data, tensor.shape.clone()))
|
|
}
|
|
|
|
/// Invert the mask (convert keep->prune and prune->keep)
|
|
pub fn invert(&self) -> Self {
|
|
let inverted_mask: Vec<bool> = self.mask.iter().map(|&x| !x).collect();
|
|
Self::new(inverted_mask, self.shape.clone())
|
|
}
|
|
|
|
/// Apply mask to gradients (used in optimizer integration)
|
|
pub fn apply_to_gradients(&self, gradients: &DenseTensor) -> Result<DenseTensor> {
|
|
if gradients.shape != self.shape {
|
|
return Err(TransformerError::Config("Gradient and mask shapes must match".into()));
|
|
}
|
|
|
|
// Zero out gradients for pruned weights
|
|
let masked_gradients: Vec<f32> = gradients.data.iter()
|
|
.zip(self.mask.iter())
|
|
.map(|(&grad, &keep)| if keep { grad } else { 0.0 })
|
|
.collect();
|
|
|
|
Ok(DenseTensor::new(masked_gradients, gradients.shape.clone()))
|
|
}
|
|
}
|
|
|
|
impl PruningStats {
|
|
/// Create new pruning statistics
|
|
pub fn new(original_params: usize, pruned_params: usize) -> Self {
|
|
let sparsity_ratio = 1.0 - (pruned_params as f64 / original_params as f64);
|
|
let compression_ratio = if pruned_params == 0 {
|
|
f64::INFINITY
|
|
} else {
|
|
original_params as f64 / pruned_params as f64
|
|
};
|
|
|
|
// Estimate memory savings (assuming 4 bytes per f32)
|
|
let memory_savings = (original_params - pruned_params) * 4;
|
|
|
|
Self {
|
|
original_params,
|
|
pruned_params,
|
|
sparsity_ratio,
|
|
compression_ratio,
|
|
memory_savings,
|
|
}
|
|
}
|
|
|
|
/// Check if tensor meets target sparsity threshold
|
|
pub fn meets_sparsity_target(&self, target_sparsity: f64) -> bool {
|
|
self.sparsity_ratio >= target_sparsity
|
|
}
|
|
|
|
/// Merge statistics from multiple pruning operations
|
|
pub fn merge(&self, other: &PruningStats) -> Self {
|
|
let total_original = self.original_params + other.original_params;
|
|
let total_pruned = self.pruned_params + other.pruned_params;
|
|
Self::new(total_original, total_pruned)
|
|
}
|
|
}
|
|
|
|
impl MagnitudePruning {
|
|
/// Apply magnitude-based pruning to tensor
|
|
///
|
|
/// # Arguments
|
|
/// * `tensor` - Input tensor to prune
|
|
/// * `config` - Pruning configuration
|
|
///
|
|
/// # Returns
|
|
/// Pruned tensor with smallest magnitude weights set to zero
|
|
pub fn prune(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
|
|
// Validate configuration
|
|
config.validate()?;
|
|
|
|
match config.strategy {
|
|
PruningStrategy::Global => {
|
|
Self::prune_global(tensor, config)
|
|
}
|
|
PruningStrategy::LayerWise => {
|
|
Self::prune_layerwise(tensor, config)
|
|
}
|
|
PruningStrategy::Structured { ref pattern_type } => {
|
|
Self::prune_structured(tensor, config, pattern_type)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Create pruning mask without modifying the original tensor
|
|
///
|
|
/// # Arguments
|
|
/// * `tensor` - Input tensor to analyze
|
|
/// * `config` - Pruning configuration
|
|
///
|
|
/// # Returns
|
|
/// Binary mask indicating which weights to keep (true) or prune (false)
|
|
pub fn create_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
|
|
// Validate configuration
|
|
config.validate()?;
|
|
|
|
match config.strategy {
|
|
PruningStrategy::Global => {
|
|
Self::create_global_mask(tensor, config)
|
|
}
|
|
PruningStrategy::LayerWise => {
|
|
Self::create_layerwise_mask(tensor, config)
|
|
}
|
|
PruningStrategy::Structured { ref pattern_type } => {
|
|
Self::create_structured_mask(tensor, config, pattern_type)
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Compute importance scores based on magnitude
|
|
///
|
|
/// # Arguments
|
|
/// * `tensor` - Input tensor
|
|
///
|
|
/// # Returns
|
|
/// Vector of importance scores (absolute values)
|
|
pub fn compute_importance_scores(tensor: &DenseTensor) -> Result<Vec<f32>> {
|
|
let scores: Vec<f32> = tensor.data.iter().map(|&x| x.abs()).collect();
|
|
Ok(scores)
|
|
}
|
|
|
|
/// Analyze pruning results
|
|
///
|
|
/// # Arguments
|
|
/// * `original` - Original tensor before pruning
|
|
/// * `pruned` - Tensor after pruning
|
|
///
|
|
/// # Returns
|
|
/// Detailed statistics about the pruning results
|
|
pub fn analyze_pruning(original: &DenseTensor, pruned: &DenseTensor) -> Result<PruningStats> {
|
|
if original.shape != pruned.shape {
|
|
return Err(TransformerError::Config("Tensor shapes must match for analysis".into()));
|
|
}
|
|
|
|
let original_nonzeros = original.data.iter().filter(|&&x| x != 0.0).count();
|
|
let pruned_nonzeros = pruned.data.iter().filter(|&&x| x != 0.0).count();
|
|
|
|
Ok(PruningStats::new(original_nonzeros, pruned_nonzeros))
|
|
}
|
|
|
|
// Private implementation methods - will be implemented in subsequent TDD iterations
|
|
|
|
fn prune_global(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
|
|
let mask = Self::create_global_mask(tensor, config)?;
|
|
mask.apply(tensor)
|
|
}
|
|
|
|
fn prune_layerwise(tensor: &DenseTensor, config: &PruningConfig) -> Result<DenseTensor> {
|
|
let mask = Self::create_layerwise_mask(tensor, config)?;
|
|
mask.apply(tensor)
|
|
}
|
|
|
|
fn prune_structured(
|
|
tensor: &DenseTensor,
|
|
config: &PruningConfig,
|
|
pattern: &StructuredPattern,
|
|
) -> Result<DenseTensor> {
|
|
let mask = Self::create_structured_mask(tensor, config, pattern)?;
|
|
mask.apply(tensor)
|
|
}
|
|
|
|
fn create_global_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
|
|
let total_elements = tensor.data.len();
|
|
let elements_to_keep = ((1.0 - config.sparsity_ratio) * total_elements as f64).round() as usize;
|
|
|
|
// Compute importance scores (absolute magnitudes)
|
|
let importance_scores = Self::compute_importance_scores(tensor)?;
|
|
|
|
// Create (index, score) pairs for sorting
|
|
let mut indexed_scores: Vec<(usize, f32)> = importance_scores.iter()
|
|
.enumerate()
|
|
.map(|(i, &score)| (i, score))
|
|
.collect();
|
|
|
|
// Sort by importance score in descending order (keep highest magnitude values)
|
|
indexed_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
|
|
|
|
// Create mask: true for kept elements, false for pruned elements
|
|
let mut mask_data = vec![false; total_elements];
|
|
for i in 0..elements_to_keep {
|
|
if i < indexed_scores.len() {
|
|
let (index, _) = indexed_scores[i];
|
|
mask_data[index] = true;
|
|
}
|
|
}
|
|
|
|
Ok(PruningMask::new(mask_data, tensor.shape.clone()))
|
|
}
|
|
|
|
fn create_layerwise_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
|
|
// For layer-wise pruning, we currently support 2D tensors (treating each row as a layer)
|
|
if tensor.shape.len() != 2 {
|
|
return Err(TransformerError::Config(
|
|
"Layer-wise pruning currently only supports 2D tensors".into()
|
|
));
|
|
}
|
|
|
|
let rows = tensor.shape[0];
|
|
let cols = tensor.shape[1];
|
|
let elements_per_row = cols;
|
|
let elements_to_keep_per_row = ((1.0 - config.sparsity_ratio) * elements_per_row as f64).round() as usize;
|
|
|
|
let mut mask_data = vec![false; tensor.data.len()];
|
|
|
|
// Process each row independently
|
|
for row in 0..rows {
|
|
let row_start = row * cols;
|
|
let row_end = row_start + cols;
|
|
|
|
// Get magnitudes for this row
|
|
let row_scores: Vec<(usize, f32)> = (row_start..row_end)
|
|
.map(|i| (i, tensor.data[i].abs()))
|
|
.collect();
|
|
|
|
// Sort by magnitude in descending order
|
|
let mut sorted_scores = row_scores;
|
|
sorted_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
|
|
|
|
// Keep the top elements_to_keep_per_row elements in this row
|
|
for i in 0..elements_to_keep_per_row.min(sorted_scores.len()) {
|
|
let (original_index, _) = sorted_scores[i];
|
|
mask_data[original_index] = true;
|
|
}
|
|
}
|
|
|
|
Ok(PruningMask::new(mask_data, tensor.shape.clone()))
|
|
}
|
|
|
|
fn create_structured_mask(
|
|
tensor: &DenseTensor,
|
|
config: &PruningConfig,
|
|
pattern: &StructuredPattern,
|
|
) -> Result<PruningMask> {
|
|
match pattern {
|
|
StructuredPattern::Channel => {
|
|
Self::create_channel_pruning_mask(tensor, config)
|
|
}
|
|
StructuredPattern::Filter => {
|
|
Self::create_filter_pruning_mask(tensor, config)
|
|
}
|
|
StructuredPattern::NM { n, m } => {
|
|
Self::create_nm_pruning_mask(tensor, config, *n, *m)
|
|
}
|
|
}
|
|
}
|
|
|
|
fn create_channel_pruning_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
|
|
// Channel pruning for 2D tensors (each row is a channel)
|
|
if tensor.shape.len() != 2 {
|
|
return Err(TransformerError::Config(
|
|
"Channel pruning currently only supports 2D tensors".into()
|
|
));
|
|
}
|
|
|
|
let rows = tensor.shape[0];
|
|
let cols = tensor.shape[1];
|
|
let channels_to_keep = ((1.0 - config.sparsity_ratio) * rows as f64).round() as usize;
|
|
|
|
// Calculate channel importance (sum of absolute values per row)
|
|
let mut channel_scores: Vec<(usize, f32)> = Vec::new();
|
|
|
|
for row in 0..rows {
|
|
let row_start = row * cols;
|
|
let row_end = row_start + cols;
|
|
let row_sum: f32 = tensor.data[row_start..row_end]
|
|
.iter()
|
|
.map(|&x| x.abs())
|
|
.sum();
|
|
channel_scores.push((row, row_sum));
|
|
}
|
|
|
|
// Sort by importance in descending order
|
|
channel_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
|
|
|
|
// Create mask: keep top channels, zero out others
|
|
let mut mask_data = vec![false; tensor.data.len()];
|
|
for i in 0..channels_to_keep {
|
|
if i < channel_scores.len() {
|
|
let (channel_idx, _) = channel_scores[i];
|
|
let row_start = channel_idx * cols;
|
|
let row_end = row_start + cols;
|
|
for j in row_start..row_end {
|
|
mask_data[j] = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(PruningMask::new(mask_data, tensor.shape.clone()))
|
|
}
|
|
|
|
fn create_filter_pruning_mask(tensor: &DenseTensor, config: &PruningConfig) -> Result<PruningMask> {
|
|
// Filter pruning for 2D tensors (each column is a filter)
|
|
if tensor.shape.len() != 2 {
|
|
return Err(TransformerError::Config(
|
|
"Filter pruning currently only supports 2D tensors".into()
|
|
));
|
|
}
|
|
|
|
let rows = tensor.shape[0];
|
|
let cols = tensor.shape[1];
|
|
let filters_to_keep = ((1.0 - config.sparsity_ratio) * cols as f64).round() as usize;
|
|
|
|
// Calculate filter importance (sum of absolute values per column)
|
|
let mut filter_scores: Vec<(usize, f32)> = Vec::new();
|
|
|
|
for col in 0..cols {
|
|
let col_sum: f32 = (0..rows)
|
|
.map(|row| tensor.data[row * cols + col].abs())
|
|
.sum();
|
|
filter_scores.push((col, col_sum));
|
|
}
|
|
|
|
// Sort by importance in descending order
|
|
filter_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
|
|
|
|
// Create mask: keep top filters, zero out others
|
|
let mut mask_data = vec![false; tensor.data.len()];
|
|
for i in 0..filters_to_keep {
|
|
if i < filter_scores.len() {
|
|
let (filter_idx, _) = filter_scores[i];
|
|
for row in 0..rows {
|
|
mask_data[row * cols + filter_idx] = true;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(PruningMask::new(mask_data, tensor.shape.clone()))
|
|
}
|
|
|
|
fn create_nm_pruning_mask(tensor: &DenseTensor, config: &PruningConfig, n: usize, m: usize) -> Result<PruningMask> {
|
|
// N:M sparsity: keep N elements out of every M consecutive elements
|
|
if n == 0 || m == 0 {
|
|
return Err(TransformerError::Config("N and M must be non-zero for N:M sparsity".into()));
|
|
}
|
|
if n > m {
|
|
return Err(TransformerError::Config("N cannot be greater than M in N:M sparsity".into()));
|
|
}
|
|
|
|
let total_elements = tensor.data.len();
|
|
let mut mask_data = vec![false; total_elements];
|
|
|
|
// Process in groups of M elements
|
|
for group_start in (0..total_elements).step_by(m) {
|
|
let group_end = (group_start + m).min(total_elements);
|
|
let group_size = group_end - group_start;
|
|
let elements_to_keep = n.min(group_size);
|
|
|
|
// Get magnitudes for this group
|
|
let mut group_scores: Vec<(usize, f32)> = (group_start..group_end)
|
|
.map(|i| (i, tensor.data[i].abs()))
|
|
.collect();
|
|
|
|
// Sort by magnitude in descending order
|
|
group_scores.sort_by(|a, b| b.1.total_cmp(&a.1));
|
|
|
|
// Keep top N elements in this group
|
|
for i in 0..elements_to_keep {
|
|
let (original_index, _) = group_scores[i];
|
|
mask_data[original_index] = true;
|
|
}
|
|
}
|
|
|
|
Ok(PruningMask::new(mask_data, tensor.shape.clone()))
|
|
}
|
|
}
|
|
|
|
impl PrunedOptimizer {
|
|
/// Create new pruned optimizer state
|
|
pub fn new() -> Self {
|
|
Self {
|
|
masks: HashMap::new(),
|
|
update_count: 0,
|
|
schedule_config: None,
|
|
}
|
|
}
|
|
|
|
/// Create pruned optimizer with schedule
|
|
pub fn with_schedule(schedule: PruningSchedule) -> Self {
|
|
Self {
|
|
masks: HashMap::new(),
|
|
update_count: 0,
|
|
schedule_config: Some(schedule),
|
|
}
|
|
}
|
|
|
|
/// Add parameter mask for tracking
|
|
pub fn add_parameter_mask(&mut self, name: &str, mask: PruningMask) {
|
|
self.masks.insert(name.to_string(), mask);
|
|
}
|
|
|
|
/// Get current sparsity ratio based on schedule
|
|
pub fn get_current_sparsity(&self) -> f64 {
|
|
if let Some(ref schedule) = self.schedule_config {
|
|
schedule.compute_sparsity(self.update_count)
|
|
} else {
|
|
// If no schedule, return average sparsity of all masks
|
|
if self.masks.is_empty() {
|
|
0.0
|
|
} else {
|
|
let total_sparsity: f64 = self.masks.values()
|
|
.map(|mask| mask.sparsity_ratio())
|
|
.sum();
|
|
total_sparsity / self.masks.len() as f64
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Apply pruning masks to gradients before optimizer step
|
|
pub fn apply_masks_to_gradients(&self, gradients: &HashMap<String, DenseTensor>) -> Result<HashMap<String, DenseTensor>> {
|
|
let mut masked_gradients = HashMap::new();
|
|
|
|
for (param_name, grad_tensor) in gradients {
|
|
if let Some(mask) = self.masks.get(param_name) {
|
|
let masked_grad = mask.apply_to_gradients(grad_tensor)?;
|
|
masked_gradients.insert(param_name.clone(), masked_grad);
|
|
} else {
|
|
// No mask for this parameter, keep gradients as-is
|
|
masked_gradients.insert(param_name.clone(), grad_tensor.clone());
|
|
}
|
|
}
|
|
|
|
Ok(masked_gradients)
|
|
}
|
|
|
|
/// Update masks based on current parameters and schedule
|
|
pub fn update_masks(&mut self, parameters: &HashMap<String, DenseTensor>, config: &PruningConfig) -> Result<()> {
|
|
let target_sparsity = if let Some(ref schedule) = self.schedule_config {
|
|
schedule.compute_sparsity(self.update_count)
|
|
} else {
|
|
config.sparsity_ratio
|
|
};
|
|
|
|
// Create updated config with current target sparsity
|
|
let mut updated_config = config.clone();
|
|
updated_config.sparsity_ratio = target_sparsity;
|
|
|
|
// Update masks for each parameter
|
|
for (param_name, tensor) in parameters {
|
|
let new_mask = MagnitudePruning::create_mask(tensor, &updated_config)?;
|
|
self.masks.insert(param_name.clone(), new_mask);
|
|
}
|
|
|
|
self.update_count += 1;
|
|
Ok(())
|
|
}
|
|
|
|
/// Get pruning statistics across all parameters
|
|
pub fn get_overall_stats(&self, parameters: &HashMap<String, DenseTensor>) -> Result<PruningStats> {
|
|
let mut total_original = 0;
|
|
let mut total_pruned = 0;
|
|
|
|
for (param_name, tensor) in parameters {
|
|
let original_count = tensor.data.len();
|
|
total_original += original_count;
|
|
|
|
if let Some(mask) = self.masks.get(param_name) {
|
|
total_pruned += mask.kept_parameters;
|
|
} else {
|
|
// No mask means all parameters are kept
|
|
total_pruned += original_count;
|
|
}
|
|
}
|
|
|
|
Ok(PruningStats::new(total_original, total_pruned))
|
|
}
|
|
}
|
|
|
|
impl PrunedTrainingState {
|
|
/// Create new training state
|
|
pub fn new() -> Self {
|
|
Self {
|
|
current_step: 0,
|
|
current_sparsity: 0.0,
|
|
stats: PruningStats::new(0, 0),
|
|
gradual_pruning_active: false,
|
|
}
|
|
}
|
|
|
|
/// Update training state
|
|
pub fn update_step(&mut self, new_stats: PruningStats, sparsity: f64) {
|
|
self.current_step += 1;
|
|
self.current_sparsity = sparsity;
|
|
self.stats = new_stats;
|
|
}
|
|
|
|
/// Check if gradual pruning should be applied
|
|
pub fn should_update_pruning(&self, update_frequency: usize) -> bool {
|
|
self.gradual_pruning_active && (self.current_step % update_frequency == 0)
|
|
}
|
|
}
|
|
|
|
impl GradualPruningController {
|
|
/// Create new gradual pruning controller
|
|
pub fn new(config: PruningConfig) -> Self {
|
|
Self {
|
|
config,
|
|
state: PrunedTrainingState::new(),
|
|
parameters: HashMap::new(),
|
|
masks: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
/// Add parameter to be pruned
|
|
pub fn add_parameter(&mut self, name: &str, tensor: DenseTensor) {
|
|
self.parameters.insert(name.to_string(), tensor);
|
|
}
|
|
|
|
/// Update parameters and recompute masks if needed
|
|
pub fn update_parameters(&mut self, parameters: &HashMap<String, DenseTensor>) -> Result<()> {
|
|
// Update stored parameters
|
|
for (name, tensor) in parameters {
|
|
self.parameters.insert(name.clone(), tensor.clone());
|
|
}
|
|
|
|
// Compute current sparsity based on schedule
|
|
let current_sparsity = if let Some(ref schedule) = self.config.schedule {
|
|
schedule.compute_sparsity(self.state.current_step)
|
|
} else {
|
|
self.config.sparsity_ratio
|
|
};
|
|
|
|
// Update config with current sparsity
|
|
let mut current_config = self.config.clone();
|
|
current_config.sparsity_ratio = current_sparsity;
|
|
|
|
// Recompute masks for all parameters
|
|
for (name, tensor) in &self.parameters {
|
|
let mask = MagnitudePruning::create_mask(tensor, ¤t_config)?;
|
|
self.masks.insert(name.clone(), mask);
|
|
}
|
|
|
|
// Update training state
|
|
let overall_stats = self.compute_overall_stats();
|
|
self.state.update_step(overall_stats, current_sparsity);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get current pruning masks
|
|
pub fn get_masks(&self) -> &HashMap<String, PruningMask> {
|
|
&self.masks
|
|
}
|
|
|
|
/// Get current training state
|
|
pub fn get_state(&self) -> &PrunedTrainingState {
|
|
&self.state
|
|
}
|
|
|
|
/// Apply current masks to parameters
|
|
pub fn apply_pruning(&self, parameters: &HashMap<String, DenseTensor>) -> Result<HashMap<String, DenseTensor>> {
|
|
let mut pruned_parameters = HashMap::new();
|
|
|
|
for (name, tensor) in parameters {
|
|
if let Some(mask) = self.masks.get(name) {
|
|
let pruned_tensor = mask.apply(tensor)?;
|
|
pruned_parameters.insert(name.clone(), pruned_tensor);
|
|
} else {
|
|
// No mask, keep original
|
|
pruned_parameters.insert(name.clone(), tensor.clone());
|
|
}
|
|
}
|
|
|
|
Ok(pruned_parameters)
|
|
}
|
|
|
|
/// Enable gradual pruning mode
|
|
pub fn enable_gradual_pruning(&mut self) {
|
|
self.state.gradual_pruning_active = true;
|
|
}
|
|
|
|
/// Disable gradual pruning mode
|
|
pub fn disable_gradual_pruning(&mut self) {
|
|
self.state.gradual_pruning_active = false;
|
|
}
|
|
|
|
/// Compute overall statistics
|
|
fn compute_overall_stats(&self) -> PruningStats {
|
|
let mut total_original = 0;
|
|
let mut total_kept = 0;
|
|
|
|
for (name, tensor) in &self.parameters {
|
|
total_original += tensor.data.len();
|
|
|
|
if let Some(mask) = self.masks.get(name) {
|
|
total_kept += mask.kept_parameters;
|
|
} else {
|
|
total_kept += tensor.data.len();
|
|
}
|
|
}
|
|
|
|
PruningStats::new(total_original, total_kept)
|
|
}
|
|
}
|
|
|
|
#[cfg(all(test, feature = "disabled_tests"))]
|
|
mod magnitude_pruning_tests {
|
|
use super::*;
|
|
|
|
// Tests will be implemented following TDD principles
|
|
} |