check marlin/bitblas dependencies in base model

mobiusml · Jul 16, 2024 · 96a5cc1 · 96a5cc1
1 parent e57104f
commit 96a5cc1
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/hqq/models/base.py b/hqq/models/base.py
@@ -16,9 +16,21 @@
 from ..core.quantize import HQQLinear
 from ..core.peft import PeftUtils, _HQQ_LORA_CLASSES
 from ..backends.torchao import HQQLinearTorchWeightOnlynt4
-from ..backends.marlin import MarlinLinear
 
-_HQQ_BACKEND_CLASSES = [HQQLinearTorchWeightOnlynt4, MarlinLinear]
+_HQQ_BACKEND_CLASSES = [HQQLinearTorchWeightOnlynt4]
+
+try:
+    from ..backends.bitblas import HQQLinearBitBlas
+    _HQQ_BACKEND_CLASSES.append(HQQLinearBitBlas)
+except Exception:
+    pass 
+
+try:
+    from ..backends.marlin import MarlinLinear
+    _HQQ_BACKEND_CLASSES.append(MarlinLinear)
+except Exception:
+    pass 
+
 
 # Defined what is qualified as "linear layer"
 _QUANT_LAYERS = [nn.Linear, HQQLinear] + _HQQ_LORA_CLASSES + _HQQ_BACKEND_CLASSES