fitcompress pruning method

nroope · nroope · commit 42df07304ee4 · 2025-12-19T14:36:30.000+01:00
diff --git a/src/pquant/configs/config_ap.yaml b/src/pquant/configs/config_ap.yaml
@@ -1,9 +1,8 @@
 pruning_parameters:
-  disable_pruning_for_layers: # Disable pruning for these layers, even if enable_pruning is true
-    -
+  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
   enable_pruning: true
   pruning_method: activation_pruning
-  threshold: 0.15
+  threshold: 0.2
   threshold_decay: 0.
   t_delta: 100
   t_start_collecting_batch: 100
diff --git a/src/pquant/configs/config_autosparse.yaml b/src/pquant/configs/config_autosparse.yaml
@@ -3,8 +3,7 @@ pruning_parameters:
   alpha_reset_epoch: 90
   autotune_epochs: 10
   backward_sparsity: false
-  disable_pruning_for_layers: # Disable pruning for these layers, even if enable_pruning is true
-    -
+  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
   enable_pruning: true
   pruning_method: autosparse
   threshold_decay: 0
diff --git a/src/pquant/configs/config_cs.yaml b/src/pquant/configs/config_cs.yaml
@@ -1,11 +1,10 @@
 pruning_parameters:
-  disable_pruning_for_layers: # Disable pruning for these layers, even if enable_pruning is true
-    -
+  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
   enable_pruning: true
   final_temp: 200
   pruning_method: cs
   threshold_decay: 1.0e-09
-  threshold_init: 0
+  threshold_init: 0.
 quantization_parameters:
   default_weight_keep_negatives: 1.
   default_weight_integer_bits: 0.
diff --git a/src/pquant/configs/config_dst.yaml b/src/pquant/configs/config_dst.yaml
@@ -1,7 +1,6 @@
 pruning_parameters:
   alpha: 5.0e-06
-  disable_pruning_for_layers: # Disable pruning for these layers, even if enable_pruning is true
-    -
+  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
   enable_pruning: true
   max_pruning_pct: 0.99
   pruning_method: dst
diff --git a/src/pquant/configs/config_fitcompress.yaml b/src/pquant/configs/config_fitcompress.yaml
@@ -0,0 +1,57 @@
+pruning_parameters:
+  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
+  enable_pruning: true
+  threshold_decay: 0.
+  pruning_method: fitcompress
+  min_frac_bits: 2.
+quantization_parameters:
+  default_weight_keep_negatives: 1.
+  default_weight_integer_bits: 0.
+  default_weight_fractional_bits: 7.
+  default_data_keep_negatives: 0.
+  default_data_integer_bits: 0.
+  default_data_fractional_bits: 7.
+  quantize_input: true
+  quantize_output: false
+  enable_quantization: true
+  hgq_beta: 1e-5
+  hgq_gamma: 0.0003
+  hgq_heterogeneous: True
+  layer_specific: []
+  use_high_granularity_quantization: false
+  use_real_tanh: false
+  use_relu_multiplier: true
+  use_symmetric_quantization: false
+  overflow: SAT
+  round_mode: RND
+fitcompress_parameters:
+  enable_fitcompress : true
+  optimize_quantization : true
+  quantization_schedule : [7., 6., 5., 4.,3.]
+  pruning_schedule : {start : 0, end : -2, steps : 40}
+  compression_goal : 0.007
+  optimize_pruning : true
+  greedy_astar : true
+  approximate : true
+  f_lambda : 1
+training_parameters:
+  epochs: 200
+  fine_tuning_epochs: 0
+  pretraining_epochs: 100
+  pruning_first: false
+  rewind: never
+  rounds: 1
+  save_weights_epoch: -1
+batch_size: 256
+cosine_tmax: 200
+gamma: 0.1
+l2_decay: 0.0001
+label_smoothing: 0.0
+lr: 0.1
+lr_schedule: cosine
+milestones:
+- -1
+- -1
+momentum: 0.9
+optimizer: adam
+plot_frequency: 100
diff --git a/src/pquant/configs/config_mdmm.yaml b/src/pquant/configs/config_mdmm.yaml
@@ -3,8 +3,7 @@
 pruning_parameters:
   pruning_method: mdmm
   enable_pruning: true
-  disable_pruning_for_layers:
-    -
+  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
   constraint_type: "Equality"
   target_value: 0.0
   metric_type: "UnstructuredSparsity"
diff --git a/src/pquant/configs/config_pdp.yaml b/src/pquant/configs/config_pdp.yaml
@@ -1,6 +1,5 @@
 pruning_parameters:
-  disable_pruning_for_layers: # Disable pruning for these layers, even if enable_pruning is true
-    -
+  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
   enable_pruning: true
   epsilon: 0.015
   pruning_method: pdp
@@ -24,13 +23,14 @@ quantization_parameters:
   layer_specific: []
   use_high_granularity_quantization: false
   use_real_tanh: false
-  use_relu_multiplier: true
-  use_symmetric_quantization: false
+  use_relu_multiplier: false
+  overflow: SAT
+  round_mode: RND
 fitcompress_parameters:
   enable_fitcompress : false
   optimize_quantization : true
   quantization_schedule : [7.,4.,3.,2.,1.]
-  pruning_schedule : {start : 0, end : -3, steps : 40}
+  pruning_schedule : {start : 0, end : -2, steps : 40}
   compression_goal : 0.10
   optimize_pruning : true
   greedy_astar : true
@@ -39,7 +39,7 @@ fitcompress_parameters:
 training_parameters:
   epochs: 100
   fine_tuning_epochs: 20
-  pretraining_epochs: 20
+  pretraining_epochs: 1
   pruning_first: false
   rewind: never
   rounds: 1
diff --git a/src/pquant/configs/config_wanda.yaml b/src/pquant/configs/config_wanda.yaml
@@ -1,7 +1,6 @@
 pruning_parameters:
   calculate_pruning_budget: True
-  disable_pruning_for_layers: # Disable pruning for these layers, even if enable_pruning is true
-    -
+  disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true
   enable_pruning: true
   M: null
   N: null
diff --git a/src/pquant/core/constants.py b/src/pquant/core/constants.py
@@ -5,6 +5,7 @@
     AutoSparsePruningModel,
     CSPruningModel,
     DSTPruningModel,
+    FITCompressPruningModel,
     MDMMPruningModel,
     PDPPruningModel,
     WandaPruningModel,
@@ -22,6 +23,7 @@
 PRUNING_MODEL_REGISTRY = {
     "cs": CSPruningModel,
     "dst": DSTPruningModel,
+    "fitcompress": FITCompressPruningModel,
     "pdp": PDPPruningModel,
     "wanda": WandaPruningModel,
     "autosparse": AutoSparsePruningModel,
@@ -44,6 +46,7 @@
 
 try:
     import mlflow
+
     LOG_FUNCTIONS_REGISTRY = {
         "torch": mlflow.pytorch.log_model,
         "tensorflow": mlflow.tensorflow.log_model,
diff --git a/src/pquant/core/finetuning.py b/src/pquant/core/finetuning.py
@@ -10,7 +10,6 @@
 import yaml
 from pydantic import BaseModel, Field, field_validator
 
-
 from pquant.core import constants
 from pquant.data_models.finetuning_model import BaseFinetuningModel
 from pquant.data_models.fitcompress_model import BaseFitCompressModel
@@ -20,6 +19,7 @@
     BasePruningModel,
     CSPruningModel,
     DSTPruningModel,
+    FITCompressPruningModel,
     MDMMPruningModel,
     PDPPruningModel,
     WandaPruningModel,
@@ -70,6 +70,7 @@ class TuningConfig(BaseModel):
         Union[
             CSPruningModel,
             DSTPruningModel,
+            FITCompressPruningModel,
             PDPPruningModel,
             WandaPruningModel,
             AutoSparsePruningModel,
@@ -368,6 +369,13 @@ def dst_config():
     return TuningConfig.load_from_file(path)
 
 
+def fitcompress_config():
+    yaml_name = "config_fitcompress.yaml"
+    parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+    path = os.path.join(parent, "configs", yaml_name)
+    return TuningConfig.load_from_file(path)
+
+
 def mdmm_config():
     yaml_name = "config_mdmm.yaml"
     parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
diff --git a/src/pquant/core/torch/fit_compress.py b/src/pquant/core/torch/fit_compress.py
@@ -485,10 +485,11 @@ def add_quantization(self, model, params, quant_config, reset=False):
             max_abs = torch.max(torch.abs(param_layer.detach().cpu()))
             eps = 1e-12
             int_bits = max(0, math.ceil(math.log2(max_abs + eps)))
-            fractional_bits = quant_config[idx] - int_bits
+            fractional_bits = max(1, quant_config[idx] - int_bits)
 
             all_int_bits.append(int_bits)
             all_frac_bits.append(fractional_bits)
+            logging.info(f"Weights int bits={int_bits}, fractional bits={fractional_bits}")
 
         for idx, param_layer in enumerate(params):
             # If reset is inactive, we quantize weights given the unquantized, but possibly pruned weights
@@ -674,17 +675,20 @@ def set_activation_bits(self, layer):
             k, i, f = layer.get_input_quantization_bits()
             bits = k + i + f
             int_bits = math.ceil(math.log2(max_abs))
-            frac_bits = bits - int_bits - k
+            frac_bits = max(self.config.pruning_parameters.min_frac_bits, bits - int_bits - k)
             layer.saved_inputs = []
             layer.input_quantizer.set_quantization_bits(int_bits, frac_bits)
+            logging.info(f"Set input quantization bits from {i}, {f} to {int_bits}, {frac_bits}")
+
         if layer.quantize_output:
             max_abs = torch.max(torch.tensor([torch.max(torch.abs(e)) for e in layer.saved_outputs]))
             k, i, f = layer.get_output_quantization_bits()
             bits = k + i + f
             int_bits = math.ceil(math.log2(max_abs))
-            frac_bits = bits - int_bits - k
+            frac_bits = max(self.config.pruning_parameters.min_frac_bits, bits - int_bits - k)
             layer.saved_outputs = []
             layer.output_quantizer.set_quantization_bits(int_bits, frac_bits)
+            logging.info(f"Set output quantization bits from {i}, {f} to {int_bits}, {frac_bits}")
 
     def post_fitcompress_calibration(self, best_node_quant_config, config, calibration_epochs=50):
         """
@@ -1145,7 +1149,14 @@ def calculate_current_compression_rate(self, params_layerwise, quant_config):
         uncompressed = 0.0
         for params_layer, quant_conf_layer in zip(params_layerwise, quant_config):
             # Count which parameters are non-zero, non_zero is simply the number of non-zero parameters in the current layer
-            non_zero = torch.sum(torch.where(torch.abs(params_layer) < 2**-quant_conf_layer, 0, 1)).detach().cpu().numpy()
+            int_bits = max(0, math.ceil(math.log2(torch.max(torch.abs(params_layer)))))
+            frac_bits_round_threshold = max(
+                1, quant_conf_layer - int_bits - 2
+            )  # - sign bit and rounding causes the second one
+            non_zero = (
+                torch.sum(torch.where(torch.abs(params_layer) < 2**-frac_bits_round_threshold, 0, 1)).detach().cpu().numpy()
+            )
+            logging.info(f"With {quant_conf_layer}, fractional bits is {frac_bits_round_threshold}")
             active_bytes += (
                 non_zero * quant_conf_layer / 8
             )  # Gives us the number of total bytes needed to store the parameters in the current layer
diff --git a/src/pquant/core/torch/layers.py b/src/pquant/core/torch/layers.py
@@ -1408,7 +1408,7 @@ def call_post_round_functions(model, rewind, rounds, r):
         rewind_weights_functions(model)
     elif rewind == "post-ticket-search" and r == rounds - 1:
         rewind_weights_functions(model)
-    else:
+    elif r != rounds - 1:
         post_round_functions(model)
 
 
diff --git a/src/pquant/core/torch/train.py b/src/pquant/core/torch/train.py
@@ -10,7 +10,7 @@
 )
 
 
-def train_model(model, config, train_func, valid_func, **kwargs):
+def train_model(model, config, train_func, valid_func, input_shape=None, **kwargs):
     """
     Generic training loop, user provides training and validation functions
     """
@@ -25,7 +25,7 @@ def train_model(model, config, train_func, valid_func, **kwargs):
             valid_func(model, epoch=epoch, **kwargs)
             post_epoch_functions(model, e, training_config.pretraining_epochs)
             epoch += 1
-    post_pretrain_functions(model, config, kwargs['trainloader'], kwargs['loss_func'])
+    post_pretrain_functions(model, config, kwargs['trainloader'], kwargs['loss_function'], input_shape=input_shape)
     for r in range(training_config.rounds):
         for e in range(training_config.epochs):
             model.train()
diff --git a/src/pquant/core/utils.py b/src/pquant/core/utils.py
@@ -6,6 +6,7 @@
 from pquant.pruning_methods.autosparse import AutoSparse
 from pquant.pruning_methods.cs import ContinuousSparsification
 from pquant.pruning_methods.dst import DST
+from pquant.pruning_methods.fitcompress import FITCompress
 from pquant.pruning_methods.mdmm import MDMM
 from pquant.pruning_methods.pdp import PDP
 from pquant.pruning_methods.wanda import Wanda
@@ -27,6 +28,8 @@ def get_pruning_layer(config, layer_type):
         return Wanda(config, layer_type)
     elif pruning_method == "mdmm":
         return MDMM(config, layer_type)
+    elif pruning_method == "fitcompress":
+        return FITCompress(config)
 
 
 def get_default_config(pruning_method: str):
@@ -35,6 +38,7 @@ def get_default_config(pruning_method: str):
         "ap",
         "cs",
         "dst",
+        "fitcompress",
         "pdp",
         "wanda",
         "mdmm",
diff --git a/src/pquant/data_models/finetuning_model.py b/src/pquant/data_models/finetuning_model.py
@@ -1,11 +1,13 @@
-from typing import Dict, List, Optional, Union, Any
+from typing import Any, Dict, List, Optional, Union
+
 from pydantic import BaseModel, Field
-from typing_extensions import Literal
+
 
 class HyperparameterSearch(BaseModel):
     numerical: Dict[str, List[Union[int, float]]] = Field(default_factory=dict)
     categorical: Optional[Dict[str, List[str]]] = Field(default_factory=dict)
 
+
 class Sampler(BaseModel):
     type: str = Field(default="TPESampler")
     params: Dict[str, Any] = Field(default_factory=dict)
diff --git a/src/pquant/data_models/pruning_model.py b/src/pquant/data_models/pruning_model.py
@@ -1,5 +1,6 @@
 from enum import Enum
 from typing import List, Literal, Optional
+
 from pydantic import BaseModel, Field
 
 
@@ -12,7 +13,7 @@ class BasePruningModel(BaseModel):
 class CSPruningModel(BasePruningModel):
     pruning_method: Literal["cs"] = "cs"
     final_temp: int = Field(default=200)
-    threshold_init: int = Field(default=0)
+    threshold_init: float = Field(default=0)
 
 
 class DSTPruningModel(BasePruningModel):
@@ -23,6 +24,11 @@ class DSTPruningModel(BasePruningModel):
     threshold_type: str = Field(default="channelwise")
 
 
+class FITCompressPruningModel(BasePruningModel):
+    pruning_method: Literal["fitcompress"] = "fitcompress"
+    min_frac_bit: float = Field(default=2.0)
+
+
 class PDPPruningModel(BasePruningModel):
     pruning_method: Literal["pdp"] = "pdp"
     epsilon: float = Field(default=0.015)
@@ -33,8 +39,8 @@ class PDPPruningModel(BasePruningModel):
 
 class WandaPruningModel(BasePruningModel):
     pruning_method: Literal["wanda"] = "wanda"
-    M: Optional[int] = Field(default=None),
-    N: Optional[int] = Field(default=None),
+    M: Optional[int] = (Field(default=None),)
+    N: Optional[int] = (Field(default=None),)
     sparsity: float = Field(default=0.9)
     t_delta: int = Field(default=100)
     t_start_collecting_batch: int = Field(default=100)
diff --git a/src/pquant/pruning_methods/autosparse.py b/src/pquant/pruning_methods/autosparse.py
diff --git a/src/pquant/pruning_methods/fitcompress.py b/src/pquant/pruning_methods/fitcompress.py
