Add the Quantizations Methods.

Author: codewithdark-git

quantllm/quant/awq.py

@@ -75,9 +75,9 @@ def quantize(
         self._process_activation_stats()
 
         # Quantize the model layer by layer
-        for name, module in self.model.named_modules():
+        for name, module in self.model.named_modules():            
             if isinstance(module, nn.Linear):
-                self.logger.info(f"Processing layer: {name}")
+                self.logger.log_info(f"Processing layer: {name}")
 
                 # Get activation scale for this layer
                 act_scale = self.act_scales.get(name)

quantllm/quant/gguf.py

@@ -39,7 +39,6 @@ def __init__(
         self.use_packed = use_packed
         self.legacy_format = legacy_format
         self.batch_size = batch_size
-        
     def quantize(
         self,
         calibration_data: Optional[torch.Tensor] = None
@@ -58,7 +57,7 @@ def quantize(
         # Convert linear layers to quantized versions
         for name, module in self.model.named_modules():
             if isinstance(module, nn.Linear):
-                self.logger.info(f"Processing layer: {name}")
+                self.logger.log_info(f"Processing layer: {name}")
 
                 # Create quantized layer
                 layer_stats = stats.get(name, None)

quantllm/quant/gptq.py

@@ -61,9 +61,9 @@ def quantize(self, calibration_data: Optional[torch.Tensor] = None) -> PreTraine
         self.model.eval()
 
         # Process layers
-        for name, module in self.model.named_modules():
+        for name, module in self.model.named_modules():            
             if isinstance(module, nn.Linear):
-                self.logger.info(f"Processing layer: {name}")
+                self.logger.log_info(f"Processing layer: {name}")
 
                 # Compute Hessian approximation
                 self.H[name] = self._compute_hessian(module, calibration_data)

quantllm/quant/quantization_engine.py

@@ -202,9 +202,8 @@ def quantize_model(
             self._quantize_layers(model, stats)
 
             return model
-            
         except Exception as e:
-            self.logger.error(f"Error during quantization: {str(e)}")
+            self.logger.log_error(f"Error during quantization: {str(e)}")
             raise
 
     def _prepare_model(self, model: PreTrainedModel) -> PreTrainedModel:
@@ -375,11 +374,11 @@ def export_model(
                 model.optimize()
                 onnx.save(model, path)
 
-            else:
+            else:                
                 raise ValueError(f"Unsupported export format: {format}")
 
         except Exception as e:
-            self.logger.error(f"Error exporting model: {str(e)}")
+            self.logger.log_error(f"Error exporting model: {str(e)}")
             raise
 
     def benchmark(
@@ -441,8 +440,8 @@ def benchmark(
                 "p99_latency": torch.quantile(latencies, 0.99).item()
             }
 
-        except Exception as e:
-            self.logger.error(f"Error during benchmarking: {str(e)}")
+        except Exception as e:            
+            self.logger.log_error(f"Error during benchmarking: {str(e)}")
             raise
 
 class BaseQuantizer:
@@ -560,15 +559,15 @@ def _prepare_model(self, original_model: PreTrainedModel):
         from transformers import AutoModelForCausalLM
 
         try:
-            # Create new model instance
-            self.logger.info("Creating new model instance...")
+            # Create new model instance            
+            self.logger.log_info("Creating new model instance...")
             new_model = AutoModelForCausalLM.from_config(
                 self.model_config,
                 trust_remote_code=True
             )
 
-            # Copy state dict with proper device handling
-            self.logger.info("Copying model parameters...")
+            # Copy state dict with proper device handling            
+            self.logger.log_info("Copying model parameters...")
             with torch.no_grad():
                 state_dict = {}
                 for name, param in original_model.state_dict().items():
@@ -586,10 +585,10 @@ def _prepare_model(self, original_model: PreTrainedModel):
                 new_model = new_model.to(self.device_manager.primary_device)
 
             self._model = new_model
-            self.logger.info("Model preparation completed successfully")
+            self.logger.log_info("Model preparation completed successfully")
 
-        except Exception as e:
-            self.logger.error(f"Failed to prepare model: {str(e)}")
+        except Exception as e:            
+            self.logger.log_error(f"Failed to prepare model: {str(e)}")
             raise
 
     def prepare_calibration_data(self, calibration_data: torch.Tensor) -> torch.Tensor: