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}