Pruning Pipeline

The pruning pipeline manages the end-to-end workflow from model loading to compressed export. This chapter covers the pipeline stages and state machine.

Pipeline Stages

The pruning process follows a defined sequence of stages:

#![allow(unused)]
fn main() {
use entrenar::prune::PruningStage;

let stages = [
    PruningStage::Idle,               // Initial state
    PruningStage::Calibrating,        // Collecting activation statistics
    PruningStage::ComputingImportance,// Scoring weight importance
    PruningStage::Pruning,            // Applying sparsity masks
    PruningStage::FineTuning,         // Recovering accuracy
    PruningStage::Evaluating,         // Validating quality
    PruningStage::Exporting,          // Saving compressed model
    PruningStage::Complete,           // Pipeline finished
];
}

Stage Details

Idle

Initial state before pruning begins.

#![allow(unused)]
fn main() {
let stage = PruningStage::Idle;
assert!(!stage.is_active());
assert!(!stage.is_terminal());
}

Calibrating

Collects activation statistics for Wanda/SparseGPT methods.

#![allow(unused)]
fn main() {
let stage = PruningStage::Calibrating;
assert!(stage.is_active());
println!("{}", stage.display_name());  // "Calibrating"
}

Activities:

• Forward pass through calibration data
• Compute per-layer activation norms
• Build Hessian approximations (SparseGPT)

Skipped when: Using magnitude pruning (no calibration needed)

ComputingImportance

Scores each weight's importance using the configured method.

#![allow(unused)]
fn main() {
let stage = PruningStage::ComputingImportance;
assert!(stage.is_active());
}

Activities:

• Apply importance formula (|w|, |w|*activation, Hessian-based)
• Compute statistics (min, max, mean, std)
• Validate for numerical stability (NaN, Inf)

Pruning

Generates and applies sparsity masks to zero out weights.

#![allow(unused)]
fn main() {
let stage = PruningStage::Pruning;
assert!(stage.is_active());
}

Activities:

• Sort weights by importance
• Generate masks based on target sparsity and pattern
• Apply masks to model weights
• Verify achieved sparsity matches target

FineTuning

Recovers accuracy lost during pruning through continued training.

#![allow(unused)]
fn main() {
let stage = PruningStage::FineTuning;
assert!(stage.is_active());
}

Activities:

• Continue training with frozen sparsity pattern
• Adjust learning rate (typically lower)
• Monitor loss convergence
• Apply gradient updates only to non-zero weights

Duration: Typically 10-20% of original training steps

Evaluating

Validates pruned model quality against benchmarks.

#![allow(unused)]
fn main() {
let stage = PruningStage::Evaluating;
assert!(stage.is_active());
}

Activities:

• Run evaluation benchmarks
• Compare to baseline (unpruned) model
• Check quality gates (max accuracy drop)
• Generate evaluation report

Exporting

Saves the compressed model in deployment format.

#![allow(unused)]
fn main() {
let stage = PruningStage::Exporting;
assert!(stage.is_active());
}

Activities:

• Convert to sparse storage format
• Apply additional compression (quantization)
• Save model weights and configuration
• Verify exported model loads correctly

Complete

Terminal state indicating successful pipeline completion.

#![allow(unused)]
fn main() {
let stage = PruningStage::Complete;
assert!(!stage.is_active());
assert!(stage.is_terminal());
}

Stage Properties

Each stage exposes useful properties:

#![allow(unused)]
fn main() {
let stage = PruningStage::FineTuning;

// Display name for UI
println!("Stage: {}", stage.display_name());  // "Fine-tuning"

// Check if actively processing
if stage.is_active() {
    println!("Pipeline is working...");
}

// Check if pipeline is done
if stage.is_terminal() {
    println!("Pipeline complete!");
}
}

Pipeline Flow Visualization

┌─────────┐
│  Idle   │
└────┬────┘
     │
     ▼
┌─────────────┐
│ Calibrating │ (skip if Magnitude)
└──────┬──────┘
       │
       ▼
┌───────────────────┐
│ComputingImportance│
└─────────┬─────────┘
          │
          ▼
    ┌─────────┐
    │ Pruning │
    └────┬────┘
         │
         ▼
   ┌───────────┐
   │FineTuning │
   └─────┬─────┘
         │
         ▼
   ┌───────────┐
   │Evaluating │
   └─────┬─────┘
         │
         ▼
   ┌───────────┐
   │ Exporting │
   └─────┬─────┘
         │
         ▼
   ┌──────────┐
   │ Complete │
   └──────────┘

Monitoring Pipeline Progress

Display pipeline status with stage indicators:

#![allow(unused)]
fn main() {
fn print_pipeline_status(current: PruningStage) {
    let stages = [
        PruningStage::Idle,
        PruningStage::Calibrating,
        PruningStage::ComputingImportance,
        PruningStage::Pruning,
        PruningStage::FineTuning,
        PruningStage::Evaluating,
        PruningStage::Exporting,
        PruningStage::Complete,
    ];

    for stage in &stages {
        let indicator = if *stage == current {
            if stage.is_active() { "🟢" } else { "✅" }
        } else {
            "⚪"
        };
        println!("{} {}", indicator, stage.display_name());
    }
}
}

Output:

✅ Idle
✅ Calibrating
🟢 Computing Importance  ← Currently here
⚪ Pruning
⚪ Fine-tuning
⚪ Evaluating
⚪ Exporting
⚪ Complete

Error Handling

Pipeline stages can fail. Handle errors gracefully:

#![allow(unused)]
fn main() {
fn run_stage(stage: PruningStage) -> Result<(), PruneError> {
    match stage {
        PruningStage::Calibrating => {
            // Could fail if out of memory
        }
        PruningStage::ComputingImportance => {
            // Could fail with NaN/Inf weights
        }
        PruningStage::Pruning => {
            // Could fail with invalid pattern
        }
        // ...
    }
    Ok(())
}
}

Best Practices

1. Log stage transitions - Track timing and progress
2. Checkpoint between stages - Enable restart on failure
3. Validate at each stage - Catch issues early
4. Monitor memory - Calibration and importance computation can spike
5. Set quality gates - Define acceptable accuracy drop thresholds

Integration Example

Complete pipeline integration:

#![allow(unused)]
fn main() {
use entrenar::prune::{
    PruningConfig, PruneMethod, PruningSchedule,
    SparsityPatternConfig, CalibrationConfig, PruningStage
};

fn run_pruning_pipeline() {
    // Configure pruning
    let config = PruningConfig::new()
        .with_method(PruneMethod::Wanda)
        .with_target_sparsity(0.5)
        .with_pattern(SparsityPatternConfig::nm_2_4())
        .with_schedule(PruningSchedule::OneShot { step: 0 });

    // Configure calibration
    let calibration = CalibrationConfig::new()
        .with_num_samples(128)
        .with_sequence_length(2048)
        .with_dataset("c4");

    // Validate before starting
    config.validate().expect("Invalid config");

    // Run pipeline stages
    let stages = [
        PruningStage::Calibrating,
        PruningStage::ComputingImportance,
        PruningStage::Pruning,
        PruningStage::FineTuning,
        PruningStage::Evaluating,
        PruningStage::Exporting,
        PruningStage::Complete,
    ];

    for stage in &stages {
        println!("Starting: {}", stage.display_name());
        // Execute stage...
        println!("Completed: {}", stage.display_name());
    }
}
}