Skip to main content

entrenar/prune/config/
method.rs

1//! Pruning method enumeration.
2
3use serde::{Deserialize, Serialize};
4
5/// Pruning method selection.
6///
7/// Each method has different trade-offs between accuracy and computational cost.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
9#[serde(rename_all = "snake_case")]
10pub enum PruneMethod {
11    /// Magnitude-based pruning (Han et al., 2015)
12    /// Simple, fast, no calibration required.
13    #[default]
14    Magnitude,
15
16    /// Wanda: Weight and Activation pruning (Sun et al., 2023)
17    /// Requires calibration data for activation statistics.
18    Wanda,
19
20    /// SparseGPT: Hessian-based pruning (Frantar & Alistarh, 2023)
21    /// Most accurate but computationally expensive.
22    SparseGpt,
23
24    /// Minitron depth pruning - removes entire layers.
25    MinitronDepth,
26
27    /// Minitron width pruning - removes channels.
28    MinitronWidth,
29}
30
31impl PruneMethod {
32    /// Check if this method requires calibration data.
33    pub fn requires_calibration(&self) -> bool {
34        matches!(
35            self,
36            PruneMethod::Wanda
37                | PruneMethod::SparseGpt
38                | PruneMethod::MinitronDepth
39                | PruneMethod::MinitronWidth
40        )
41    }
42
43    /// Get the display name for this method.
44    pub fn display_name(&self) -> &'static str {
45        match self {
46            PruneMethod::Magnitude => "Magnitude",
47            PruneMethod::Wanda => "Wanda",
48            PruneMethod::SparseGpt => "SparseGPT",
49            PruneMethod::MinitronDepth => "Minitron (Depth)",
50            PruneMethod::MinitronWidth => "Minitron (Width)",
51        }
52    }
53}