add QUANT_TYPE in qlinear (#736)

* add QUANT_TYPE in qlinear Signed-off-by: jiqing-feng <[email protected]> * add comments Signed-off-by: jiqing-feng <[email protected]> --------- Signed-off-by: jiqing-feng <[email protected]>
ModelCloud · Dec 4, 2024 · 0326109 · 0326109
1 parent 1995602
commit 0326109
Show file tree

Hide file tree

Showing 8 changed files with 16 additions and 0 deletions.
diff --git a/gptqmodel/nn_modules/qlinear/qlinear_bitblas.py b/gptqmodel/nn_modules/qlinear/qlinear_bitblas.py
@@ -92,6 +92,8 @@ class BitBLASQuantLinear(BaseQuantLinear):
         torch.half: "float16",
         torch.int8: "int8",
     }
+    # for transformers/optimum tests compat
+    QUANT_TYPE = "bitblas"
 
     def __init__(
         self,

diff --git a/gptqmodel/nn_modules/qlinear/qlinear_cuda.py b/gptqmodel/nn_modules/qlinear/qlinear_cuda.py
@@ -23,6 +23,8 @@
 class CudaQuantLinear(TorchQuantLinear):
     SUPPORTS_BITS = [2, 3, 4, 8]
     SUPPORTS_DEVICES = [DEVICE.CUDA]
+    # for transformers/optimum tests compat
+    QUANT_TYPE = "cuda"
 
     def __init__(
             self,

diff --git a/gptqmodel/nn_modules/qlinear/qlinear_exllama.py b/gptqmodel/nn_modules/qlinear/qlinear_exllama.py
@@ -45,6 +45,8 @@ class ExllamaQuantLinear(BaseQuantLinear):
     SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [32]
     SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [32]
     SUPPORTS_DEVICES = [DEVICE.CUDA]
+    # for transformers/optimum tests compat
+    QUANT_TYPE = "exllama"
 
     """Linear layer implementation with per-group 4-bit quantization of the weights"""
 

diff --git a/gptqmodel/nn_modules/qlinear/qlinear_exllamav2.py b/gptqmodel/nn_modules/qlinear/qlinear_exllamav2.py
@@ -107,6 +107,8 @@ class ExllamaV2QuantLinear(BaseQuantLinear):
     SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [32]
     SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [32]
     SUPPORTS_DEVICES = [DEVICE.CUDA]
+    # for transformers/optimum tests compat
+    QUANT_TYPE = "exllamav2"
 
     """Linear layer implementation with per-group 4-bit quantization of the weights"""
 

diff --git a/gptqmodel/nn_modules/qlinear/qlinear_ipex.py b/gptqmodel/nn_modules/qlinear/qlinear_ipex.py
@@ -51,6 +51,8 @@ def convert_dtype_torch2str(dtype):
 class IPEXQuantLinear(BaseQuantLinear):
     SUPPORTS_BITS = [4]
     SUPPORTS_DEVICES = [DEVICE.CPU, DEVICE.XPU]
+    # for transformers/optimum tests compat
+    QUANT_TYPE = "ipex"
 
     def __init__(
         self,

diff --git a/gptqmodel/nn_modules/qlinear/qlinear_marlin.py b/gptqmodel/nn_modules/qlinear/qlinear_marlin.py
@@ -145,6 +145,8 @@ class MarlinQuantLinear(BaseQuantLinear):
     SUPPORTS_SYM = [True]
     SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [64]
     SUPPORTS_DEVICES = [DEVICE.CUDA]
+    # for transformers/optimum tests compat
+    QUANT_TYPE = "marlin"
 
     def __init__(self, bits: int, group_size: int, desc_act: bool, sym: bool, infeatures: int, outfeatures: int,
                  bias: bool, **kwargs):

diff --git a/gptqmodel/nn_modules/qlinear/qlinear_torch.py b/gptqmodel/nn_modules/qlinear/qlinear_torch.py
@@ -15,6 +15,8 @@
 class TorchQuantLinear(BaseQuantLinear):
     SUPPORTS_BITS = [2, 3, 4, 8]
     SUPPORTS_DEVICES = [DEVICE.CPU, DEVICE.XPU, DEVICE.CUDA]
+    # for transformers/optimum tests compat
+    QUANT_TYPE = "torch"
 
     def __init__(
         self,

diff --git a/gptqmodel/nn_modules/qlinear/qlinear_tritonv2.py b/gptqmodel/nn_modules/qlinear/qlinear_tritonv2.py
@@ -31,6 +31,8 @@ class TritonV2QuantLinear(BaseQuantLinear, TritonModuleMixin):
     SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [32]
     SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [32]
     SUPPORTS_DEVICES = [DEVICE.CUDA]
+    # for transformers/optimum tests compat
+    QUANT_TYPE = "tritonv2"
 
     """
     Triton v2 quantized linear layer.