requirements

Author: Giuseppe5

requirements/requirements-llm.txt

@@ -10,3 +10,4 @@ pydantic
 torch>=2.4
 tqdm
 transformers[sentencepiece]<5.0
+vllm

src/brevitas/export/inference/handler.py

@@ -301,19 +301,23 @@ def __init__(self):
         self.register_buffer('exponent_bit_width', torch.ones(()))
         self.register_buffer('exponent_bias', torch.ones(()))
         self.register_buffer('fp_internal_scale_min', torch.ones(()))
-        self.register_buffer('saturating', torch.ones(()).to(torch.bool))
+        self.register_buffer('saturating_t', torch.ones(()).to(torch.bool))
         self.inf_values = None
         self.nan_values = None
         self.eps = 1e-8  #torch.finfo(self.scale.dtype).tiny
 
+    @property
+    def saturating(self):
+        return bool(self.saturating_t.item())
+
     def prepare_for_export(self, module):
         FloatToIntMixin.prepare_for_export(self, module)
         if module.is_quant_enabled:
 
             self.exponent_bit_width = module.exponent_bit_width()
             self.mantissa_bit_width = module.mantissa_bit_width()
             self.exponent_bias = module.exponent_bias()
-            self.saturating = torch.tensor(module.is_saturating())
+            self.saturating_t = torch.tensor(module.is_saturating())
             self.inf_values = module.inf_values()
             self.nan_values = module.nan_values()
             if module.tensor_quant is not None:

src/brevitas/export/inference/vLLM/manager.py

@@ -13,7 +13,6 @@
 from torch.nn import Module
 import torch.nn as nn
 from vllm.model_executor.layers.linear import LinearBase
-from vllm.model_executor.layers.linear import LinearMethodBase
 from vllm.model_executor.layers.linear import MergedColumnParallelLinear
 from vllm.model_executor.layers.linear import QKVParallelLinear
 from vllm.model_executor.layers.linear import RowParallelLinear
@@ -23,24 +22,7 @@
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 import brevitas.config as config
-from brevitas.export.inference.handler import DynamicFloatInferenceHandler
-from brevitas.export.inference.handler import DynamicIntInferenceHandler
-from brevitas.export.inference.handler import FloatInferencetHandler
-from brevitas.export.inference.handler import FloatWeightInferencetHandler
-from brevitas.export.inference.handler import GroupwiseFloatInferenceHandler
-from brevitas.export.inference.handler import GroupwiseFloatWeightInferenceHandler
-from brevitas.export.inference.handler import GroupwiseIntInferenceHandler
-from brevitas.export.inference.handler import GroupwiseIntWeightInferenceHandler
-from brevitas.export.inference.handler import IntInferencetHandler
-from brevitas.export.inference.handler import IntWeightInferencetHandler
 from brevitas.export.inference.vLLM.handler import QuantLinear
-from brevitas.export.manager import _set_proxy_export_handler
-from brevitas.export.manager import _set_proxy_export_mode
-from brevitas.export.manager import _set_recurrent_layer_export_handler
-from brevitas.export.manager import _set_recurrent_layer_export_mode
-from brevitas.export.manager import BaseManager
-from brevitas.graph.calibrate import QuantizationStatusManager
-from brevitas.nn.equalized_layer import EqualizedModule
 from brevitas.nn.equalized_layer import RotatedModule
 from brevitas.nn.mixin import QuantLayerMixin
 from brevitas.proxy.quant_proxy import QuantProxyFromInjector
@@ -178,8 +160,9 @@ def export(self, model, filepath):
                         proxy_dict['class_type'] = export_handler.__class__.__name__
                 if isinstance(module, self.wrap_layers):
                     layer_dict['rotation_config'] = dict()
-                    layer_dict['rotation_config']['rot_mat_shape'] = module.had_mat.shape[0] if module.had_mat is not None else None
+                    layer_dict['rotation_config']['rot_mat_shape'] = module.had_mat.shape[
+                        0] if module.had_mat is not None else None
                     layer_dict['rotation_config']['k'] = module.k
-                    
+
         with open(json_filename, 'w') as f:
             json.dump(json_to_save, f, cls=EncodeTensor)